Merge pull request matplotlib#8870 from dstansby/plot-timedelta

phobson · web-flow · commit 95e56559fe7c · 2017-08-07T08:36:02.000-07:00
Add num2timedelta method with test
diff --git a/lib/matplotlib/dates.py b/lib/matplotlib/dates.py
@@ -136,7 +136,7 @@
 import matplotlib.ticker as ticker
 
 
-__all__ = ('date2num', 'num2date', 'drange', 'epoch2num',
+__all__ = ('date2num', 'num2date', 'num2timedelta', 'drange', 'epoch2num',
            'num2epoch', 'mx2num', 'DateFormatter',
            'IndexDateFormatter', 'AutoDateFormatter', 'DateLocator',
            'RRuleLocator', 'AutoDateLocator', 'YearLocator',
@@ -406,6 +406,38 @@ def num2date(x, tz=None):
         return _from_ordinalf_np_vectorized(x, tz).tolist()
 
 
+def _ordinalf_to_timedelta(x):
+    return datetime.timedelta(days=x)
+
+
+_ordinalf_to_timedelta_np_vectorized = np.vectorize(_ordinalf_to_timedelta)
+
+
+def num2timedelta(x):
+    """
+    Converts number of days to a :class:`timdelta` object.
+    If *x* is a sequence, a sequence of :class:`timedelta` objects will
+    be returned.
+
+    Parameters
+    ----------
+    x : float, sequence of floats
+        Number of days (fraction part represents hours, minutes, seconds)
+
+    Returns
+    -------
+    :class:`timedelta` or list[:class:`timedelta`]
+
+    """
+    if not cbook.iterable(x):
+        return _ordinalf_to_timedelta(x)
+    else:
+        x = np.asarray(x)
+        if not x.size:
+            return x
+        return _ordinalf_to_timedelta_np_vectorized(x).tolist()
+
+
 def drange(dstart, dend, delta):
     """
     Return a date range as float Gregorian ordinals.  *dstart* and
diff --git a/lib/matplotlib/tests/test_dates.py b/lib/matplotlib/tests/test_dates.py
@@ -457,3 +457,12 @@ def test_DayLocator():
 def test_tz_utc():
     dt = datetime.datetime(1970, 1, 1, tzinfo=mdates.UTC)
     dt.tzname()
+
+
+@pytest.mark.parametrize("x, tdelta",
+                         [(1, datetime.timedelta(days=1)),
+                          ([1, 1.5], [datetime.timedelta(days=1),
+                                      datetime.timedelta(days=1.5)])])
+def test_num2timedelta(x, tdelta):
+    dt = mdates.num2timedelta(x)
+    assert dt == tdelta
