FIX: translate timedelta64 in _dt64_to_ordinalf

jklymak · jklymak · commit c569ca7d43c9 · 2018-11-22T13:28:42.000-08:00
diff --git a/lib/matplotlib/dates.py b/lib/matplotlib/dates.py
@@ -255,6 +255,9 @@ def _dt64_to_ordinalf(d):
     because we do times compared to ``0001-01-01T00:00:00`` (plus one day).
     """
 
+    if isinstance(d, np.timedelta64):
+        return d / np.timedelta64(1, 'D')
+
     # the "extra" ensures that we at least allow the dynamic range out to
     # seconds.  That should get out to +/-2e11 years.
     # NOTE: First cast truncates; second cast back is for NumPy 1.10.
@@ -275,6 +278,13 @@ def _dt64_to_ordinalf(d):
     return dt
 
 
+def _dt64_to_ordinalf_iterable(d):
+    dnew = np.zeros(len(d))
+    for nn, dd in enumerate(d):
+        dnew[nn] = _dt64_to_ordinalf(dd)
+    return(dnew)
+
+
 def _from_ordinalf(x, tz=None):
     """
     Convert Gregorian float of the date, preserving hours, minutes,
@@ -413,19 +423,32 @@ def date2num(d):
     if hasattr(d, "values"):
         # this unpacks pandas series or dataframes...
         d = d.values
-    if not np.iterable(d):
-        if (isinstance(d, np.datetime64) or (isinstance(d, np.ndarray) and
-                np.issubdtype(d.dtype, np.datetime64))):
-            return _dt64_to_ordinalf(d)
-        return _to_ordinalf(d)
 
-    else:
-        d = np.asarray(d)
-        if np.issubdtype(d.dtype, np.datetime64):
+    if hasattr(d, 'size') and not d.size:
+        # this is an empty
+        return d
+
+    if not np.iterable(d) and not isinstance(d, np.ndarray):
+        # single value logic...
+        if (isinstance(d, np.datetime64) or isinstance(d, np.timedelta64)):
             return _dt64_to_ordinalf(d)
-        if not d.size:
-            return d
+        else:
+            return _to_ordinalf(d)
+
+    elif (isinstance(d, np.ndarray) and
+            (np.issubdtype(d.dtype, np.datetime64) or
+            np.issubdtype(d.dtype, np.timedelta64))):
+        # array with all one type of datetime64 object.
+        return _dt64_to_ordinalf(d)
+
+    elif len(d):
+        # this is a list or tuple...
+        if (isinstance(d[0], np.datetime64) or
+                isinstance(d[0], np.timedelta64)):
+            return _dt64_to_ordinalf_iterable(d)
         return _to_ordinalf_np_vectorized(d)
+    else:
+        return []
 
 
 def julian2num(j):
diff --git a/lib/matplotlib/tests/test_dates.py b/lib/matplotlib/tests/test_dates.py
@@ -680,3 +680,21 @@ def test_datetime64_in_list():
     dt = [np.datetime64('2000-01-01'), np.datetime64('2001-01-01')]
     dn = mdates.date2num(dt)
     assert np.array_equal(dn, [730120.,  730486.])
+
+
+def test_timedelta():
+    """
+    test that timedelta objects are properly translated into days
+    """
+    dt = [datetime.datetime(2000, 1, 1, 0, 0, 0),
+          datetime.timedelta(days=1, hours=2)]
+    assert mdates.date2num(dt[1]) == 1 + 2 / 24
+    # check that mixed lists work....
+    assert mdates.date2num(dt)[0] == 730120.0
+    assert mdates.date2num(dt)[1] == 1 + 2 / 24
+    dt = (np.datetime64('2000-01-01'),
+          np.timedelta64(26, 'h'))
+    assert mdates.date2num(dt[1]) == 1 + 2 / 24
+    # check that mixed lists work....
+    assert mdates.date2num(dt)[0] == 730120.0
+    assert mdates.date2num(dt)[1] == 1 + 2 / 24
