Add num2timedelta method with test

dstansby · dstansby · commit b8c2abf75452 · 2017-07-25T14:11:05.000+01:00
diff --git a/lib/matplotlib/dates.py b/lib/matplotlib/dates.py
@@ -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`
+
+    """
+    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
