FuncFormatter can take func(x) and func(x,pos), and FuncFormatter tests

story645 · story645 · commit eed5b94d71c1 · 2020-05-01T10:44:29.000-04:00
diff --git a/lib/matplotlib/tests/test_ticker.py b/lib/matplotlib/tests/test_ticker.py
@@ -1041,6 +1041,27 @@ def test_basic(self):
         assert '00002' == tmp_form(2)
 
 
+class TestFuncFormatter:
+    def test_input_x_pos(self):
+        func = lambda x, pos : f"{x}+{pos}"
+        funcform = mticker.FuncFormatter(func)
+        assert "1+2" == funcform(1,2)
+
+    def text_input_x(self):
+        func = lambda x: f"hello {x}"
+        funcform = mticker.FuncFormatter(func)
+        assert "hello 0" == funcform(0)
+
+    def test_error(self):
+        with pytest.raises(TypeError, match=r"FuncFormatter*"):
+            func = lambda x, y, z: " ".join([x, y, z])
+            funcform = mticker.FuncFormatter(func)
+            funcform(1,2,3)
+
+    def test_built_in(self):
+        funcform = mticker.FuncFormatter("{} hi!".format)
+        assert "42 hi!" == funcform(42)
+
 class TestStrMethodFormatter:
     test_data = [
         ('{x:05d}', (2,), '00002'),
diff --git a/lib/matplotlib/ticker.py b/lib/matplotlib/ticker.py
@@ -168,6 +168,7 @@
 import logging
 import locale
 import math
+import inspect
 from numbers import Integral
 
 import numpy as np
@@ -377,19 +378,34 @@ class FuncFormatter(Formatter):
     """
     Use a user-defined function for formatting.
 
-    The function should take in two inputs (a tick value ``x`` and a
-    position ``pos``), and return a string containing the corresponding
-    tick label.
+    The function can take in at most two inputs (a required tick value ``x`` and 
+    an optional position ``pos``), and must return a string containing the 
+    corresponding tick label.
     """
     def __init__(self, func):
-        self.func = func
+        try:
+            self.nargs = len(inspect.signature(func).parameters)
+            if self.nargs == 2:
+                self.func = func 
+            elif self.nargs == 1:
+                def func_pos(x, pos):
+                    return func(x)
+                self.func = func_pos
+            else:
+                raise TypeError("""FuncFormatter functions take at most 
+                                2 parameters: x (required), pos (optional)""")
+        except ValueError as e:
+            #built ins like str.format don't have signatures
+            self.func = func
+
 
     def __call__(self, x, pos=None):
         """
         Return the value of the user defined function.
 
         *x* and *pos* are passed through as-is.
         """
+    
         return self.func(x, pos)
 
 
