ENH: datetime: Add timedelta -> timedelta scalar casting

3 files changed, 123 insertions, 1 deletions

diff --git a/numpy/core/src/multiarray/_datetime.h b/numpy/core/src/multiarray/_datetime.h
index e4bc2eaf0..4c7e4457b 100644
--- a/numpy/core/src/multiarray/_datetime.h
+++ b/numpy/core/src/multiarray/_datetime.h
@@ -280,5 +280,16 @@ cast_datetime_to_datetime(PyArray_DatetimeMetaData *src_meta,
                           npy_datetime src_dt,
                           npy_datetime *dst_dt);
 
+/*
+ * Casts a single timedelta from having src_meta metadata into
+ * dst_meta metadata.
+ *
+ * Returns 0 on success, -1 on failure.
+ */
+NPY_NO_EXPORT int
+cast_timedelta_to_timedelta(PyArray_DatetimeMetaData *src_meta,
+                          PyArray_DatetimeMetaData *dst_meta,
+                          npy_timedelta src_dt,
+                          npy_timedelta *dst_dt);
 
 #endif
diff --git a/numpy/core/src/multiarray/datetime.c b/numpy/core/src/multiarray/datetime.c
index 99704c186..84297198a 100644
--- a/numpy/core/src/multiarray/datetime.c
+++ b/numpy/core/src/multiarray/datetime.c
@@ -2169,7 +2169,6 @@ convert_pyobject_to_datetime_metadata(PyObject *obj,
     PyObject *ascii = NULL;
     char *str = NULL;
     Py_ssize_t len = 0;
-    NPY_DATETIMEUNIT unit;
 
     if (PyTuple_Check(obj)) {
         return convert_datetime_metadata_tuple_to_datetime_metadata(obj,
@@ -3137,6 +3136,31 @@ convert_pyobject_to_timedelta(PyArray_DatetimeMetaData *meta, PyObject *obj,
         *out = PyLong_AsLongLong(obj);
         return 0;
     }
+    /* Timedelta scalar */
+    else if (PyArray_IsScalar(obj, Timedelta)) {
+        PyTimedeltaScalarObject *dts = (PyTimedeltaScalarObject *)obj;
+
+        return cast_timedelta_to_timedelta(&dts->obmeta, meta,
+                                            dts->obval, out);
+    }
+    /* Timedelta zero-dimensional array */
+    else if (PyArray_Check(obj) &&
+                    PyArray_NDIM(obj) == 0 &&
+                    PyArray_DESCR(obj)->type_num == NPY_TIMEDELTA) {
+        PyArray_DatetimeMetaData *obj_meta;
+        npy_timedelta dt = 0;
+
+        obj_meta = get_datetime_metadata_from_dtype(PyArray_DESCR(obj));
+        if (obj_meta == NULL) {
+            return -1;
+        }
+        PyArray_DESCR(obj)->f->copyswap(&dt,
+                                        PyArray_DATA(obj),
+                                        !PyArray_ISNOTSWAPPED(obj),
+                                        obj);
+
+        return cast_timedelta_to_timedelta(obj_meta, meta, dt, out);
+    }
     /* TODO: Finish this function */
 
     PyErr_SetString(PyExc_ValueError,
@@ -3287,3 +3311,63 @@ cast_datetime_to_datetime(PyArray_DatetimeMetaData *src_meta,
     return 0;
 }
 
+/*
+ * Casts a single timedelta from having src_meta metadata into
+ * dst_meta metadata.
+ *
+ * Returns 0 on success, -1 on failure.
+ */
+NPY_NO_EXPORT int
+cast_timedelta_to_timedelta(PyArray_DatetimeMetaData *src_meta,
+                          PyArray_DatetimeMetaData *dst_meta,
+                          npy_timedelta src_dt,
+                          npy_timedelta *dst_dt)
+{
+    npy_int64 num = 0, denom = 0;
+    int event = 0;
+
+    /* If the metadata is the same, short-circuit the conversion */
+    if (src_meta->base == dst_meta->base &&
+            src_meta->num == dst_meta->num &&
+            src_meta->events == dst_meta->events) {
+        *dst_dt = src_dt;
+        return 0;
+    }
+
+    /* Get the conversion factor */
+    get_datetime_conversion_factor(src_meta, dst_meta, &num, &denom);
+
+    if (num == 0) {
+        PyErr_SetString(PyExc_OverflowError,
+                    "Integer overflow getting a conversion factor between "
+                    "different timedelta types");
+        return -1;
+    }
+
+    /* Remove the event number from the value */
+    if (src_meta->events > 1) {
+        event = (int)(src_dt % src_meta->events);
+        src_dt = src_dt / src_meta->events;
+        if (event < 0) {
+            --src_dt;
+            event += src_meta->events;
+        }
+    }
+
+    /* Apply the scaling */
+    if (src_dt < 0) {
+        *dst_dt = (src_dt * num - (denom - 1)) / denom;
+    }
+    else {
+        *dst_dt = src_dt * num / denom;
+    }
+
+    /* Add the event number back in */
+    if (dst_meta->events > 1) {
+        event = event % dst_meta->events;
+        *dst_dt = (*dst_dt) * dst_meta->events + event;
+    }
+
+    return 0;
+}
+
diff --git a/numpy/core/tests/test_datetime.py b/numpy/core/tests/test_datetime.py
index 7e4e48963..fe130fb05 100644
--- a/numpy/core/tests/test_datetime.py
+++ b/numpy/core/tests/test_datetime.py
@@ -99,6 +99,33 @@ class TestDateTime(TestCase):
         assert_equal(np.datetime64(b, 's'), a)
         assert_equal(np.datetime64(b, 's').dtype, np.dtype('M8[s]'))
 
+    def test_timedelta_scalar_construction(self):
+        # Construct with different units
+        assert_equal(np.timedelta64(7, 'D'),
+                     np.timedelta64(1, 'W'))
+        assert_equal(np.timedelta64(120, 's'),
+                     np.timedelta64(2, 'm'))
+
+        # When constructing from a scalar or zero-dimensional array,
+        # it either keeps the units or you can override them.
+        a = np.timedelta64(2, 'h')
+        b = np.array(2, dtype='m8[h]')
+
+        assert_equal(a.dtype, np.dtype('m8[h]'))
+        assert_equal(b.dtype, np.dtype('m8[h]'))
+
+        assert_equal(np.timedelta64(a), a);
+        assert_equal(np.timedelta64(a).dtype, np.dtype('m8[h]'))
+
+        assert_equal(np.timedelta64(b), a)
+        assert_equal(np.timedelta64(b).dtype, np.dtype('m8[h]'))
+
+        assert_equal(np.timedelta64(a, 's'), a)
+        assert_equal(np.timedelta64(a, 's').dtype, np.dtype('m8[s]'))
+
+        assert_equal(np.timedelta64(b, 's'), a)
+        assert_equal(np.timedelta64(b, 's').dtype, np.dtype('m8[s]'))
+
     def test_datetime_nat_casting(self):
         a = np.array('NaT', dtype='M8[D]')
         b = np.datetime64('NaT', '[D]')