Merge pull request #2707 from pelson/callable_date_formatter

tacaswell · tacaswell · commit f3d407d5631d · 2014-01-14T07:43:33.000-08:00
Callable date formatter
diff --git a/lib/matplotlib/dates.py b/lib/matplotlib/dates.py
@@ -517,9 +517,25 @@ class AutoDateFormatter(ticker.Formatter):
     dictionary by doing::
 
 
-      formatter = AutoDateFormatter()
-      formatter.scaled[1/(24.*60.)] = '%M:%S' # only show min and sec
-
+    >>> formatter = AutoDateFormatter()
+    >>> formatter.scaled[1/(24.*60.)] = '%M:%S' # only show min and sec
+
+    Custom `FunctionFormatter`s can also be used. The following example shows
+    how to use a custom format function to strip trailing zeros from decimal
+    seconds and adds the date to the first ticklabel::
+
+    >>> def my_format_function(x, pos=None):
+    ...     x = matplotlib.dates.num2date(x)
+    ...     if pos == 0:
+    ...         fmt = '%D %H:%M:%S.%f'
+    ...     else:
+    ...         fmt = '%H:%M:%S.%f'
+    ...     label = x.strftime(fmt)
+    ...     label = label.rstrip("0")
+    ...     label = label.rstrip(".")
+    ...     return label
+    >>> from matplotlib.ticker import FuncFormatter
+    >>> formatter.scaled[1/(24.*60.)] = FuncFormatter(my_format_function)
     """
 
     # This can be improved by providing some user-level direction on
@@ -535,8 +551,9 @@ class AutoDateFormatter(ticker.Formatter):
 
     def __init__(self, locator, tz=None, defaultfmt='%Y-%m-%d'):
         """
-        Autofmt the date labels.  The default format is the one to use
-        if none of the times in scaled match
+        Autoformat the date labels.  The default format is the one to use
+        if none of the values in ``self.scaled`` are greater than the unit
+        returned by ``locator._get_unit()``.
         """
         self._locator = locator
         self._tz = tz
@@ -548,17 +565,25 @@ def __init__(self, locator, tz=None, defaultfmt='%Y-%m-%d'):
                        1. / 24.: '%H:%M:%S',
                        1. / (24. * 60.): '%H:%M:%S.%f'}
 
-    def __call__(self, x, pos=0):
-        scale = float(self._locator._get_unit())
+    def __call__(self, x, pos=None):
+        locator_unit_scale = float(self._locator._get_unit())
         fmt = self.defaultfmt
 
-        for k in sorted(self.scaled):
-            if k >= scale:
-                fmt = self.scaled[k]
+        # Pick the first scale which is greater than the locator unit.
+        for possible_scale in sorted(self.scaled):
+            if possible_scale >= locator_unit_scale:
+                fmt = self.scaled[possible_scale]
                 break
 
-        self._formatter = DateFormatter(fmt, self._tz)
-        return self._formatter(x, pos)
+        if isinstance(fmt, six.string_types):
+            self._formatter = DateFormatter(fmt, self._tz)
+            result = self._formatter(x, pos)
+        elif six.callable(fmt):
+            result = fmt(x, pos)
+        else:
+            raise TypeError('Unexpected type passed to {!r}.'.formatter(self))
+
+        return result
 
 
 class rrulewrapper:
diff --git a/lib/matplotlib/tests/baseline_images/test_dates/date_axvspan.png b/lib/matplotlib/tests/baseline_images/test_dates/date_axvspan.png
diff --git a/lib/matplotlib/tests/test_dates.py b/lib/matplotlib/tests/test_dates.py
@@ -8,13 +8,13 @@
 import warnings
 import tempfile
 
-from nose.tools import assert_raises, assert_equal
 import dateutil
+import mock
+from nose.tools import assert_raises, assert_equal
 
 from matplotlib.testing.decorators import image_comparison, cleanup
 import matplotlib.pyplot as plt
 import matplotlib.dates as mdates
-from matplotlib.dates import DayLocator
 
 
 @image_comparison(baseline_images=['date_empty'], extensions=['png'])
@@ -155,6 +155,18 @@ def test_DateFormatter():
     fig.autofmt_xdate()
 
 
+def test_date_formatter_callable():
+    scale = -11
+    locator = mock.Mock(_get_unit=mock.Mock(return_value=scale))
+    callable_formatting_function = lambda dates, _: \
+                        [dt.strftime('%d-%m//%Y') for dt in dates]
+    
+    formatter = mdates.AutoDateFormatter(locator)
+    formatter.scaled[-10] = callable_formatting_function
+    assert_equal(formatter([datetime.datetime(2014, 12, 25)]),
+                 ['25-12//2014'])
+
+
 def test_drange():
     """
     This test should check if drange works as expected, and if all the
