ENH: datetime-casting: Refine the casting rules with a date vs time unit barrier

4 files changed, 294 insertions, 80 deletions

diff --git a/numpy/core/src/multiarray/_datetime.h b/numpy/core/src/multiarray/_datetime.h
index 8a9ffff7b..fd1c3e22f 100644
--- a/numpy/core/src/multiarray/_datetime.h
+++ b/numpy/core/src/multiarray/_datetime.h
@@ -99,11 +99,27 @@ convert_datetime_divisor_to_multiple(PyArray_DatetimeMetaData *meta,
  */
 NPY_NO_EXPORT npy_bool
 datetime_metadata_divides(
-                        PyArray_Descr *dividend,
-                        PyArray_Descr *divisor,
+                        PyArray_DatetimeMetaData *dividend,
+                        PyArray_DatetimeMetaData *divisor,
                         int strict_with_nonlinear_units);
 
 /*
+ * This provides the casting rules for the DATETIME data type metadata.
+ */
+NPY_NO_EXPORT npy_bool
+can_cast_datetime64_metadata(PyArray_DatetimeMetaData *src_meta,
+                             PyArray_DatetimeMetaData *dst_meta,
+                             NPY_CASTING casting);
+
+/*
+ * This provides the casting rules for the TIMEDELTA data type metadata.
+ */
+NPY_NO_EXPORT npy_bool
+can_cast_timedelta64_metadata(PyArray_DatetimeMetaData *src_meta,
+                             PyArray_DatetimeMetaData *dst_meta,
+                             NPY_CASTING casting);
+
+/*
  * Computes the GCD of the two date-time metadata values. Raises
  * an exception if there is no reasonable GCD, such as with
  * years and days.
diff --git a/numpy/core/src/multiarray/convert_datatype.c b/numpy/core/src/multiarray/convert_datatype.c
index 9e0b93a56..a1864c2ac 100644
--- a/numpy/core/src/multiarray/convert_datatype.c
+++ b/numpy/core/src/multiarray/convert_datatype.c
@@ -184,21 +184,21 @@ PyArray_CanCastSafely(int fromtype, int totype)
     }
     /* Special-cases for some types */
     switch (fromtype) {
-        case PyArray_DATETIME:
-        case PyArray_TIMEDELTA:
-        case PyArray_OBJECT:
-        case PyArray_VOID:
+        case NPY_DATETIME:
+        case NPY_TIMEDELTA:
+        case NPY_OBJECT:
+        case NPY_VOID:
             return 0;
-        case PyArray_BOOL:
+        case NPY_BOOL:
             return 1;
     }
     switch (totype) {
-        case PyArray_BOOL:
-        case PyArray_DATETIME:
-        case PyArray_TIMEDELTA:
+        case NPY_BOOL:
+        case NPY_DATETIME:
+        case NPY_TIMEDELTA:
             return 0;
-        case PyArray_OBJECT:
-        case PyArray_VOID:
+        case NPY_OBJECT:
+        case NPY_VOID:
             return 1;
     }
 
@@ -210,7 +210,7 @@ PyArray_CanCastSafely(int fromtype, int totype)
     if (from->f->cancastto) {
         int *curtype = from->f->cancastto;
 
-        while (*curtype != PyArray_NOTYPE) {
+        while (*curtype != NPY_NOTYPE) {
             if (*curtype++ == totype) {
                 return 1;
             }
@@ -228,23 +228,23 @@ PyArray_CanCastSafely(int fromtype, int totype)
 NPY_NO_EXPORT npy_bool
 PyArray_CanCastTo(PyArray_Descr *from, PyArray_Descr *to)
 {
-    int fromtype=from->type_num;
-    int totype=to->type_num;
+    int from_type_num = from->type_num;
+    int to_type_num = to->type_num;
     npy_bool ret;
 
-    ret = (npy_bool) PyArray_CanCastSafely(fromtype, totype);
+    ret = (npy_bool) PyArray_CanCastSafely(from_type_num, to_type_num);
     if (ret) {
         /* Check String and Unicode more closely */
-        if (fromtype == NPY_STRING) {
-            if (totype == NPY_STRING) {
+        if (from_type_num == NPY_STRING) {
+            if (to_type_num == NPY_STRING) {
                 ret = (from->elsize <= to->elsize);
             }
-            else if (totype == NPY_UNICODE) {
+            else if (to_type_num == NPY_UNICODE) {
                 ret = (from->elsize << 2 <= to->elsize);
             }
         }
-        else if (fromtype == NPY_UNICODE) {
-            if (totype == NPY_UNICODE) {
+        else if (from_type_num == NPY_UNICODE) {
+            if (to_type_num == NPY_UNICODE) {
                 ret = (from->elsize <= to->elsize);
             }
         }
@@ -252,14 +252,41 @@ PyArray_CanCastTo(PyArray_Descr *from, PyArray_Descr *to)
          * For datetime/timedelta, only treat casts moving towards
          * more precision as safe.
          */
-        else if (fromtype == NPY_DATETIME && totype == NPY_DATETIME) {
-            return datetime_metadata_divides(from, to, 0);
+        else if (from_type_num == NPY_DATETIME && to_type_num == NPY_DATETIME) {
+            PyArray_DatetimeMetaData *meta1, *meta2;
+            meta1 = get_datetime_metadata_from_dtype(from);
+            if (meta1 == NULL) {
+                PyErr_Clear();
+                return 0;
+            }
+            meta2 = get_datetime_metadata_from_dtype(to);
+            if (meta2 == NULL) {
+                PyErr_Clear();
+                return 0;
+            }
+
+            return can_cast_datetime64_metadata(meta1, meta2,
+                                                NPY_SAFE_CASTING);
         }
-        else if (fromtype == NPY_TIMEDELTA && totype == NPY_TIMEDELTA) {
-            return datetime_metadata_divides(from, to, 1);
+        else if (from_type_num == NPY_TIMEDELTA &&
+                                    to_type_num == NPY_TIMEDELTA) {
+            PyArray_DatetimeMetaData *meta1, *meta2;
+            meta1 = get_datetime_metadata_from_dtype(from);
+            if (meta1 == NULL) {
+                PyErr_Clear();
+                return 0;
+            }
+            meta2 = get_datetime_metadata_from_dtype(to);
+            if (meta2 == NULL) {
+                PyErr_Clear();
+                return 0;
+            }
+
+            return can_cast_timedelta64_metadata(meta1, meta2,
+                                                 NPY_SAFE_CASTING);
         }
         /*
-         * TODO: If totype is STRING or unicode
+         * TODO: If to_type_num is STRING or unicode
          * see if the length is long enough to hold the
          * stringified value of the object.
          */
@@ -374,22 +401,50 @@ PyArray_CanCastTypeTo(PyArray_Descr *from, PyArray_Descr *to,
         }
 
         switch (from->type_num) {
-            case NPY_DATETIME:
-            case NPY_TIMEDELTA:
-                switch (casting) {
-                    case NPY_NO_CASTING:
-                        return PyArray_ISNBO(from->byteorder) ==
-                                            PyArray_ISNBO(to->byteorder) &&
-                                has_equivalent_datetime_metadata(from, to);
-                    case NPY_EQUIV_CASTING:
-                        return has_equivalent_datetime_metadata(from, to);
-                    case NPY_SAFE_CASTING:
-                        return datetime_metadata_divides(from, to,
-                                            from->type_num == NPY_TIMEDELTA);
-                    default:
-                        return 1;
+            case NPY_DATETIME: {
+                PyArray_DatetimeMetaData *meta1, *meta2;
+                meta1 = get_datetime_metadata_from_dtype(from);
+                if (meta1 == NULL) {
+                    PyErr_Clear();
+                    return 0;
                 }
-                break;
+                meta2 = get_datetime_metadata_from_dtype(to);
+                if (meta2 == NULL) {
+                    PyErr_Clear();
+                    return 0;
+                }
+
+                if (casting == NPY_NO_CASTING) {
+                    return PyArray_ISNBO(from->byteorder) ==
+                                        PyArray_ISNBO(to->byteorder) &&
+                            can_cast_datetime64_metadata(meta1, meta2, casting);
+                }
+                else {
+                    return can_cast_datetime64_metadata(meta1, meta2, casting);
+                }
+            }
+            case NPY_TIMEDELTA: {
+                PyArray_DatetimeMetaData *meta1, *meta2;
+                meta1 = get_datetime_metadata_from_dtype(from);
+                if (meta1 == NULL) {
+                    PyErr_Clear();
+                    return 0;
+                }
+                meta2 = get_datetime_metadata_from_dtype(to);
+                if (meta2 == NULL) {
+                    PyErr_Clear();
+                    return 0;
+                }
+
+                if (casting == NPY_NO_CASTING) {
+                    return PyArray_ISNBO(from->byteorder) ==
+                                        PyArray_ISNBO(to->byteorder) &&
+                        can_cast_timedelta64_metadata(meta1, meta2, casting);
+                }
+                else {
+                    return can_cast_timedelta64_metadata(meta1, meta2, casting);
+                }
+            }
             default:
                 switch (casting) {
                     case NPY_NO_CASTING:
diff --git a/numpy/core/src/multiarray/datetime.c b/numpy/core/src/multiarray/datetime.c
index 16ac5c4fb..27f4c11ed 100644
--- a/numpy/core/src/multiarray/datetime.c
+++ b/numpy/core/src/multiarray/datetime.c
@@ -1537,58 +1537,38 @@ get_datetime_conversion_factor(PyArray_DatetimeMetaData *src_meta,
  */
 NPY_NO_EXPORT npy_bool
 datetime_metadata_divides(
-                        PyArray_Descr *dividend,
-                        PyArray_Descr *divisor,
+                        PyArray_DatetimeMetaData *dividend,
+                        PyArray_DatetimeMetaData *divisor,
                         int strict_with_nonlinear_units)
 {
-    PyArray_DatetimeMetaData *meta1, *meta2;
     npy_uint64 num1, num2;
 
-    /* Must be datetime types */
-    if ((dividend->type_num != NPY_DATETIME &&
-                        dividend->type_num != NPY_TIMEDELTA) ||
-                    (divisor->type_num != NPY_DATETIME &&
-                        divisor->type_num != NPY_TIMEDELTA)) {
-        return 0;
-    }
-
-    meta1 = get_datetime_metadata_from_dtype(dividend);
-    if (meta1 == NULL) {
-        PyErr_Clear();
-        return 0;
-    }
-    meta2 = get_datetime_metadata_from_dtype(divisor);
-    if (meta2 == NULL) {
-        PyErr_Clear();
-        return 0;
-    }
-
     /* Generic units divide into anything */
-    if (meta2->base == NPY_FR_GENERIC) {
+    if (divisor->base == NPY_FR_GENERIC) {
         return 1;
     }
     /* Non-generic units never divide into generic units */
-    else if (meta1->base == NPY_FR_GENERIC) {
+    else if (dividend->base == NPY_FR_GENERIC) {
         return 0;
     }
 
     /* Events must match */
-    if (meta1->events != meta2->events) {
+    if (dividend->events != divisor->events) {
         return 0;
     }
 
-    num1 = (npy_uint64)meta1->num;
-    num2 = (npy_uint64)meta2->num;
+    num1 = (npy_uint64)dividend->num;
+    num2 = (npy_uint64)divisor->num;
 
     /* If the bases are different, factor in a conversion */
-    if (meta1->base != meta2->base) {
+    if (dividend->base != divisor->base) {
         /*
          * Years and Months are incompatible with
          * all other units (except years and months are compatible
          * with each other).
          */
-        if (meta1->base == NPY_FR_Y) {
-            if (meta2->base == NPY_FR_M) {
+        if (dividend->base == NPY_FR_Y) {
+            if (divisor->base == NPY_FR_M) {
                 num1 *= 12;
             }
             else if (strict_with_nonlinear_units) {
@@ -1599,8 +1579,8 @@ datetime_metadata_divides(
                 return 1;
             }
         }
-        else if (meta2->base == NPY_FR_Y) {
-            if (meta1->base == NPY_FR_M) {
+        else if (divisor->base == NPY_FR_Y) {
+            if (dividend->base == NPY_FR_M) {
                 num2 *= 12;
             }
             else if (strict_with_nonlinear_units) {
@@ -1611,7 +1591,7 @@ datetime_metadata_divides(
                 return 1;
             }
         }
-        else if (meta1->base == NPY_FR_M || meta2->base == NPY_FR_M) {
+        else if (dividend->base == NPY_FR_M || divisor->base == NPY_FR_M) {
             if (strict_with_nonlinear_units) {
                 return 0;
             }
@@ -1622,14 +1602,14 @@ datetime_metadata_divides(
         }
 
         /* Take the greater base (unit sizes are decreasing in enum) */
-        if (meta1->base > meta2->base) {
-            num2 *= get_datetime_units_factor(meta2->base, meta1->base);
+        if (dividend->base > divisor->base) {
+            num2 *= get_datetime_units_factor(divisor->base, dividend->base);
             if (num2 == 0) {
                 return 0;
             }
         }
         else {
-            num1 *= get_datetime_units_factor(meta1->base, meta2->base);
+            num1 *= get_datetime_units_factor(dividend->base, divisor->base);
             if (num1 == 0) {
                 return 0;
             }
@@ -1645,6 +1625,162 @@ datetime_metadata_divides(
 }
 
 /*
+ * This provides the casting rules for the DATETIME data type units.
+ *
+ * Notably, there is a barrier between 'date units' and 'time units'
+ * for all but 'unsafe' casting.
+ */
+NPY_NO_EXPORT npy_bool
+can_cast_datetime64_units(NPY_DATETIMEUNIT src_unit,
+                          NPY_DATETIMEUNIT dst_unit,
+                          NPY_CASTING casting)
+{
+    switch (casting) {
+        /* Allow anything with unsafe casting */
+        case NPY_UNSAFE_CASTING:
+            return 1;
+
+        /*
+         * Only enforce the 'date units' vs 'time units' barrier with
+         * 'same_kind' casting.
+         */
+        case NPY_SAME_KIND_CASTING:
+            if (src_unit == NPY_FR_GENERIC || dst_unit == NPY_FR_GENERIC) {
+                return src_unit == dst_unit;
+            }
+            else {
+                return (src_unit <= NPY_FR_D && dst_unit <= NPY_FR_D) ||
+                       (src_unit > NPY_FR_D && dst_unit > NPY_FR_D);
+            }
+
+        /*
+         * Enforce the 'date units' vs 'time units' barrier and that
+         * casting is only allowed towards more precise units with
+         * 'safe' casting.
+         */
+        case NPY_SAFE_CASTING:
+            if (src_unit == NPY_FR_GENERIC || dst_unit == NPY_FR_GENERIC) {
+                return src_unit == dst_unit;
+            }
+            else {
+                return (src_unit <= dst_unit) &&
+                       ((src_unit <= NPY_FR_D && dst_unit <= NPY_FR_D) ||
+                        (src_unit > NPY_FR_D && dst_unit > NPY_FR_D));
+            }
+
+        /* Enforce equality with 'no' or 'equiv' casting */
+        default:
+            return src_unit == dst_unit;
+    }
+}
+
+/*
+ * This provides the casting rules for the TIMEDELTA data type units.
+ *
+ * Notably, there is a barrier between the nonlinear years and
+ * months units, and all the other units.
+ */
+NPY_NO_EXPORT npy_bool
+can_cast_timedelta64_units(NPY_DATETIMEUNIT src_unit,
+                          NPY_DATETIMEUNIT dst_unit,
+                          NPY_CASTING casting)
+{
+    switch (casting) {
+        /* Allow anything with unsafe casting */
+        case NPY_UNSAFE_CASTING:
+            return 1;
+
+        /*
+         * Only enforce the 'date units' vs 'time units' barrier with
+         * 'same_kind' casting.
+         */
+        case NPY_SAME_KIND_CASTING:
+            if (src_unit == NPY_FR_GENERIC || dst_unit == NPY_FR_GENERIC) {
+                return src_unit == dst_unit;
+            }
+            else {
+                return (src_unit <= NPY_FR_M && dst_unit <= NPY_FR_M) ||
+                       (src_unit > NPY_FR_M && dst_unit > NPY_FR_M);
+            }
+
+        /*
+         * Enforce the 'date units' vs 'time units' barrier and that
+         * casting is only allowed towards more precise units with
+         * 'safe' casting.
+         */
+        case NPY_SAFE_CASTING:
+            if (src_unit == NPY_FR_GENERIC || dst_unit == NPY_FR_GENERIC) {
+                return src_unit == dst_unit;
+            }
+            else {
+                return (src_unit <= dst_unit) &&
+                       ((src_unit <= NPY_FR_M && dst_unit <= NPY_FR_M) ||
+                        (src_unit > NPY_FR_M && dst_unit > NPY_FR_M));
+            }
+
+        /* Enforce equality with 'no' or 'equiv' casting */
+        default:
+            return src_unit == dst_unit;
+    }
+}
+
+/*
+ * This provides the casting rules for the DATETIME data type metadata.
+ */
+NPY_NO_EXPORT npy_bool
+can_cast_datetime64_metadata(PyArray_DatetimeMetaData *src_meta,
+                             PyArray_DatetimeMetaData *dst_meta,
+                             NPY_CASTING casting)
+{
+    switch (casting) {
+        case NPY_UNSAFE_CASTING:
+            return 1;
+
+        case NPY_SAME_KIND_CASTING:
+            return can_cast_datetime64_units(src_meta->base, dst_meta->base,
+                                             casting);
+
+        case NPY_SAFE_CASTING:
+            return can_cast_datetime64_units(src_meta->base, dst_meta->base,
+                                                             casting) &&
+                   datetime_metadata_divides(src_meta, dst_meta, 0);
+
+        default:
+            return src_meta->base == dst_meta->base &&
+                   src_meta->num == dst_meta->num &&
+                   src_meta->events == dst_meta->events;
+    }
+}
+
+/*
+ * This provides the casting rules for the TIMEDELTA data type metadata.
+ */
+NPY_NO_EXPORT npy_bool
+can_cast_timedelta64_metadata(PyArray_DatetimeMetaData *src_meta,
+                             PyArray_DatetimeMetaData *dst_meta,
+                             NPY_CASTING casting)
+{
+    switch (casting) {
+        case NPY_UNSAFE_CASTING:
+            return 1;
+
+        case NPY_SAME_KIND_CASTING:
+            return can_cast_timedelta64_units(src_meta->base, dst_meta->base,
+                                             casting);
+
+        case NPY_SAFE_CASTING:
+            return can_cast_timedelta64_units(src_meta->base, dst_meta->base,
+                                                             casting) &&
+                   datetime_metadata_divides(src_meta, dst_meta, 1);
+
+        default:
+            return src_meta->base == dst_meta->base &&
+                   src_meta->num == dst_meta->num &&
+                   src_meta->events == dst_meta->events;
+    }
+}
+
+/*
  * Computes the GCD of the two date-time metadata values. Raises
  * an exception if there is no reasonable GCD, such as with
  * years and days.
diff --git a/numpy/core/tests/test_datetime.py b/numpy/core/tests/test_datetime.py
index c533f3b2c..863cefea3 100644
--- a/numpy/core/tests/test_datetime.py
+++ b/numpy/core/tests/test_datetime.py
@@ -77,11 +77,18 @@ class TestDateTime(TestCase):
         # Cannot cast timedelta safely from months/years to days
         assert_(not np.can_cast('m8[M]', 'm8[D]', casting='safe'))
         assert_(not np.can_cast('m8[Y]', 'm8[D]', casting='safe'))
-        # Can cast same_kind from months/years to days
+        # Can cast datetime same_kind from months/years to days
         assert_(np.can_cast('M8[M]', 'M8[D]', casting='same_kind'))
         assert_(np.can_cast('M8[Y]', 'M8[D]', casting='same_kind'))
-        assert_(np.can_cast('m8[M]', 'm8[D]', casting='same_kind'))
-        assert_(np.can_cast('m8[Y]', 'm8[D]', casting='same_kind'))
+        # Can't cast timedelta same_kind from months/years to days
+        assert_(not np.can_cast('m8[M]', 'm8[D]', casting='same_kind'))
+        assert_(not np.can_cast('m8[Y]', 'm8[D]', casting='same_kind'))
+        # Can't cast datetime same_kind across the date/time boundary
+        assert_(not np.can_cast('M8[D]', 'M8[h]', casting='same_kind'))
+        assert_(not np.can_cast('M8[h]', 'M8[D]', casting='same_kind'))
+        # Can cast timedelta same_kind across the date/time boundary
+        assert_(np.can_cast('m8[D]', 'm8[h]', casting='same_kind'))
+        assert_(np.can_cast('m8[h]', 'm8[D]', casting='same_kind'))
 
         # Cannot cast safely if the integer multiplier doesn't divide
         assert_(not np.can_cast('M8[7h]', 'M8[3h]', casting='safe'))