ENH: datetime-arange: Add boilerplate for the specialized datetime_arange

3 files changed, 96 insertions, 0 deletions

diff --git a/numpy/core/src/multiarray/_datetime.h b/numpy/core/src/multiarray/_datetime.h
index 05dab36a4..20e40eec5 100644
--- a/numpy/core/src/multiarray/_datetime.h
+++ b/numpy/core/src/multiarray/_datetime.h
@@ -352,4 +352,18 @@ cast_timedelta_to_timedelta(PyArray_DatetimeMetaData *src_meta,
                           npy_timedelta src_dt,
                           npy_timedelta *dst_dt);
 
+/*
+ * Returns true if the object is something that is best considered
+ * a Datetime or Timedelta, false otherwise.
+ */
+NPY_NO_EXPORT npy_bool
+is_any_numpy_datetime_or_timedelta(PyObject *obj);
+
+/*
+ * Implements a datetime-specific arange
+ */
+NPY_NO_EXPORT PyObject *
+datetime_arange(PyObject *start, PyObject *stop, PyObject *step,
+                PyArray_Descr *dtype);
+
 #endif
diff --git a/numpy/core/src/multiarray/ctors.c b/numpy/core/src/multiarray/ctors.c
index a292a1a82..f1bc7bcef 100644
--- a/numpy/core/src/multiarray/ctors.c
+++ b/numpy/core/src/multiarray/ctors.c
@@ -3083,6 +3083,15 @@ PyArray_ArangeObj(PyObject *start, PyObject *stop, PyObject *step, PyArray_Descr
     PyArray_Descr *native = NULL;
     int swap;
 
+    /* Datetime arange is handled specially */
+    if ((dtype != NULL && (dtype->type_num == NPY_DATETIME ||
+                           dtype->type_num == NPY_TIMEDELTA)) ||
+            (dtype == NULL && is_any_numpy_datetime_or_timedelta(start) ||
+                              is_any_numpy_datetime_or_timedelta(stop) ||
+                              is_any_numpy_datetime_or_timedelta(step))) {
+        return datetime_arange(start, stop, step, dtype);
+    }
+
     if (!dtype) {
         PyArray_Descr *deftype;
         PyArray_Descr *newtype;
diff --git a/numpy/core/src/multiarray/datetime.c b/numpy/core/src/multiarray/datetime.c
index 17733fc1c..3d0a9658a 100644
--- a/numpy/core/src/multiarray/datetime.c
+++ b/numpy/core/src/multiarray/datetime.c
@@ -4295,3 +4295,76 @@ cast_timedelta_to_timedelta(PyArray_DatetimeMetaData *src_meta,
     return 0;
 }
 
+/*
+ * Returns true if the object is something that is best considered
+ * a Datetime, false otherwise.
+ */
+static npy_bool
+is_any_numpy_datetime(PyObject *obj)
+{
+    return (PyArray_IsScalar(obj, Datetime) ||
+            (PyArray_Check(obj) && (
+                PyArray_DESCR(obj)->type_num == NPY_DATETIME)) ||
+            PyDate_Check(obj) ||
+            PyDatetime_Check(obj));
+}
+
+/*
+ * Returns true if the object is something that is best considered
+ * a Timedelta, false otherwise.
+ */
+static npy_bool
+is_any_numpy_timedelta(PyObject *obj)
+{
+    return (PyArray_IsScalar(obj, Timedelta) ||
+            (PyArray_Check(obj) && (
+                PyArray_DESCR(obj)->type_num == NPY_TIMEDELTA)) ||
+            PyDelta_Check(obj));
+}
+
+/*
+ * Returns true if the object is something that is best considered
+ * a Datetime or Timedelta, false otherwise.
+ */
+NPY_NO_EXPORT npy_bool
+is_any_numpy_datetime_or_timedelta(PyObject *obj)
+{
+    return obj != NULL &&
+           (is_any_numpy_datetime(obj) ||
+            is_any_numpy_timedelta(obj));
+}
+
+NPY_NO_EXPORT PyObject *
+datetime_arange(PyObject *start, PyObject *stop, PyObject *step,
+                PyArray_Descr *dtype)
+{
+    PyArray_DatetimeMetaData meta_start, meta_stop, meta_step;
+    /*
+     * First normalize the input parameters so there is no Py_None,
+     * and start is moved to stop if stop is unspecified.
+     */
+    if (step == Py_None) {
+        step = NULL;
+    }
+    if (stop == NULL || stop == Py_None) {
+        stop = start;
+        start = NULL;
+        /* If start was NULL or None, raise an exception */
+        if (stop == NULL || stop == Py_None) {
+            PyErr_SetString(PyExc_ValueError,
+                    "arange needs at least a stopping value");
+            return NULL;
+        }
+    }
+    if (start == Py_None) {
+        start = NULL;
+    }
+
+    /*
+     * Now figure out whether we're doing a datetime or a timedelta
+     * arange, and extract the values.
+     */
+    
+    PyErr_SetString(PyExc_RuntimeError, "datetime_arange is incomplete");
+    return NULL;
+}