ENH: datetime-arange: The arange function largely works now

4 files changed, 158 insertions, 58 deletions

diff --git a/numpy/core/src/multiarray/_datetime.h b/numpy/core/src/multiarray/_datetime.h
index 34b3b5eba..c2645b34d 100644
--- a/numpy/core/src/multiarray/_datetime.h
+++ b/numpy/core/src/multiarray/_datetime.h
@@ -376,7 +376,7 @@ is_any_numpy_datetime_or_timedelta(PyObject *obj);
 /*
  * Implements a datetime-specific arange
  */
-NPY_NO_EXPORT PyObject *
+NPY_NO_EXPORT PyArrayObject *
 datetime_arange(PyObject *start, PyObject *stop, PyObject *step,
                 PyArray_Descr *dtype);
 
diff --git a/numpy/core/src/multiarray/ctors.c b/numpy/core/src/multiarray/ctors.c
index b400f1fec..35ed706c4 100644
--- a/numpy/core/src/multiarray/ctors.c
+++ b/numpy/core/src/multiarray/ctors.c
@@ -3085,7 +3085,7 @@ PyArray_ArangeObj(PyObject *start, PyObject *stop, PyObject *step, PyArray_Descr
             (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);
+        return (PyObject *)datetime_arange(start, stop, step, dtype);
     }
 
     if (!dtype) {
diff --git a/numpy/core/src/multiarray/datetime.c b/numpy/core/src/multiarray/datetime.c
index 3a208bb3b..6c41fe708 100644
--- a/numpy/core/src/multiarray/datetime.c
+++ b/numpy/core/src/multiarray/datetime.c
@@ -1178,30 +1178,70 @@ bad_input:
     return -1;
 }
 
-NPY_NO_EXPORT PyObject *
-parse_datetime_metacobj_from_metastr(char *metastr, Py_ssize_t len)
+/*
+ * Creates a datetime or timedelta dtype using the provided metadata.
+ */
+NPY_NO_EXPORT PyArray_Descr *
+create_datetime_dtype(int type_num, PyArray_DatetimeMetaData *meta)
 {
+    PyArray_Descr *dtype = NULL;
     PyArray_DatetimeMetaData *dt_data;
+    PyObject *metacobj = NULL;
+
+    /* Create a default datetime or timedelta */
+    if (type_num == NPY_DATETIME || type_num == NPY_TIMEDELTA) {
+        dtype = PyArray_DescrNewFromType(type_num);
+    }
+    else {
+        PyErr_SetString(PyExc_RuntimeError,
+                "Asked to create a datetime type with a non-datetime "
+                "type number");
+        return NULL;
+    }
+
+    if (dtype == NULL) {
+        return NULL;
+    }
+
+    /*
+     * Remove any reference to old metadata dictionary
+     * And create a new one for this new dtype
+     */
+    Py_XDECREF(dtype->metadata);
+    dtype->metadata = PyDict_New();
+    if (dtype->metadata == NULL) {
+        Py_DECREF(dtype);
+        return NULL;
+    }
 
+    /* Create a metadata capsule to copy the provided metadata */
     dt_data = PyArray_malloc(sizeof(PyArray_DatetimeMetaData));
     if (dt_data == NULL) {
-        return PyErr_NoMemory();
+        Py_DECREF(dtype);
+        PyErr_NoMemory();
+        return NULL;
     }
 
-    /* If there's no metastr, use the default */
-    if (len == 0) {
-        dt_data->base = NPY_DATETIME_DEFAULTUNIT;
-        dt_data->num = 1;
-        dt_data->events = 1;
+    /* Copy the metadata */
+    *dt_data = *meta;
+
+    /* Allocate a capsule for it (this claims ownership of dt_data) */
+    metacobj = NpyCapsule_FromVoidPtr((void *)dt_data, simple_capsule_dtor);
+    if (metacobj == NULL) {
+        Py_DECREF(dtype);
+        return NULL;
     }
-    else {
-        if (parse_datetime_metadata_from_metastr(metastr, len, dt_data) < 0) {
-            PyArray_free(dt_data);
-            return NULL;
-        }
+
+    /* Set the metadata object in the dictionary. */
+    if (PyDict_SetItemString(dtype->metadata, NPY_METADATA_DTSTR,
+                                                    metacobj) < 0) {
+        Py_DECREF(dtype);
+        Py_DECREF(metacobj);
+        return NULL;
     }
+    Py_DECREF(metacobj);
 
-    return NpyCapsule_FromVoidPtr((void *)dt_data, simple_capsule_dtor);
+    return dtype;
 }
 
 /*
@@ -1211,11 +1251,10 @@ parse_datetime_metacobj_from_metastr(char *metastr, Py_ssize_t len)
 NPY_NO_EXPORT PyArray_Descr *
 parse_dtype_from_datetime_typestr(char *typestr, Py_ssize_t len)
 {
-    PyArray_Descr *dtype = NULL;
+    PyArray_DatetimeMetaData meta;
     char *metastr = NULL;
     int is_timedelta = 0;
     Py_ssize_t metalen = 0;
-    PyObject *metacobj = NULL;
 
     if (len < 2) {
         PyErr_Format(PyExc_TypeError,
@@ -1255,45 +1294,13 @@ parse_dtype_from_datetime_typestr(char *typestr, Py_ssize_t len)
         return NULL;
     }
 
-    /* Create a default datetime or timedelta */
-    if (is_timedelta) {
-        dtype = PyArray_DescrNewFromType(NPY_TIMEDELTA);
-    }
-    else {
-        dtype = PyArray_DescrNewFromType(NPY_DATETIME);
-    }
-    if (dtype == NULL) {
+    /* Parse the metadata string into a metadata struct */
+    if (parse_datetime_metadata_from_metastr(metastr, metalen, &meta) < 0) {
         return NULL;
     }
 
-    /*
-     * Remove any reference to old metadata dictionary
-     * And create a new one for this new dtype
-     */
-    Py_XDECREF(dtype->metadata);
-    dtype->metadata = PyDict_New();
-    if (dtype->metadata == NULL) {
-        Py_DECREF(dtype);
-        return NULL;
-    }
-
-    /* Parse the metadata string into a metadata capsule */
-    metacobj = parse_datetime_metacobj_from_metastr(metastr, metalen);
-    if (metacobj == NULL) {
-        Py_DECREF(dtype);
-        return NULL;
-    }
-
-    /* Set the metadata object in the dictionary. */
-    if (PyDict_SetItemString(dtype->metadata, NPY_METADATA_DTSTR,
-                                                    metacobj) < 0) {
-        Py_DECREF(dtype);
-        Py_DECREF(metacobj);
-        return NULL;
-    }
-    Py_DECREF(metacobj);
-
-    return dtype;
+    return create_datetime_dtype(is_timedelta ? NPY_TIMEDELTA : NPY_DATETIME,
+                                    &meta);
 }
 
 static NPY_DATETIMEUNIT _multiples_table[16][4] = {
@@ -4445,7 +4452,7 @@ is_any_numpy_datetime_or_timedelta(PyObject *obj)
             is_any_numpy_timedelta(obj));
 }
 
-NPY_NO_EXPORT PyObject *
+NPY_NO_EXPORT PyArrayObject *
 datetime_arange(PyObject *start, PyObject *stop, PyObject *step,
                 PyArray_Descr *dtype)
 {
@@ -4457,6 +4464,10 @@ datetime_arange(PyObject *start, PyObject *stop, PyObject *step,
      */
     npy_int64 start_value = 0, stop_value = 0, step_value;
 
+    npy_intp i, length;
+    PyArrayObject *ret;
+    npy_int64 *ret_data;
+
     /*
      * First normalize the input parameters so there is no Py_None,
      * and start is moved to stop if stop is unspecified.
@@ -4562,6 +4573,17 @@ datetime_arange(PyObject *start, PyObject *stop, PyObject *step,
                 }
             }
         }
+
+        /* Get the 'step' value */
+        if (step == NULL) {
+            step_value = 1;
+        }
+        else {
+            if (convert_pyobject_to_timedelta(&meta, step,
+                                            &step_value) < 0) {
+                return NULL;
+            }
+        }
     }
     /* if dtype is NULL */
     else {
@@ -4735,7 +4757,71 @@ datetime_arange(PyObject *start, PyObject *stop, PyObject *step,
             }
         }
     }
-    
-    PyErr_SetString(PyExc_RuntimeError, "datetime_arange is incomplete");
-    return NULL;
+
+    /* Now start, stop, and step have their values and matching metadata */
+    if (start_value == NPY_DATETIME_NAT ||
+                    stop_value == NPY_DATETIME_NAT ||
+                    step_value == NPY_DATETIME_NAT) {
+        PyErr_SetString(PyExc_ValueError,
+                    "arange: cannot use NaT (not-a-time) datetime values");
+        return NULL;
+    }
+
+    /* Calculate the array length */
+    if (step_value > 0) {
+        if (stop_value > start_value) {
+            length = (stop_value - start_value) / step_value;
+        }
+        else {
+            length = 0;
+        }
+    }
+    else if (step_value < 0) {
+        if (stop_value < start_value) {
+            length = (start_value - stop_value) / (-step_value);
+        }
+        else {
+            length = 0;
+        }
+    }
+    else {
+        PyErr_SetString(PyExc_ValueError,
+                    "arange: step may not be zero");
+        return NULL;
+    }
+
+    /* Create the dtype of the result */
+    if (dtype != NULL) {
+        Py_INCREF(dtype);
+    }
+    else {
+        dtype = create_datetime_dtype(
+                        is_timedelta ? NPY_TIMEDELTA : NPY_DATETIME,
+                        &meta);
+        if (dtype == NULL) {
+            return NULL;
+        }
+    }
+
+    /* Create the result array */
+    ret = (PyArrayObject *)PyArray_NewFromDescr(
+                            &PyArray_Type, dtype, 1, &length, NULL,
+                            NULL, 0, NULL);
+    if (ret == NULL) {
+        return NULL;
+    }
+
+    if (length > 0) {
+        /* Extract the data pointer */
+        ret_data = (npy_int64 *)PyArray_DATA(ret);
+
+        /* Create the timedeltas or datetimes */
+        for (i = 0; i < length; ++i) {
+            *ret_data = start_value;
+            start_value += step_value;
+            ret_data++;
+        }
+    }
+
+    return ret;
 }
diff --git a/numpy/core/tests/test_datetime.py b/numpy/core/tests/test_datetime.py
index 2d13d399e..9470c16d7 100644
--- a/numpy/core/tests/test_datetime.py
+++ b/numpy/core/tests/test_datetime.py
@@ -1100,6 +1100,20 @@ class TestDateTime(TestCase):
         assert_equal(np.datetime_as_string(a, local=True, tzoffset=-5*60),
                      '2010-03-15T01:30-0500')
 
+    def test_datetime_arange(self):
+        # Default mode
+        a = np.arange('2010-01-05', '2010-01-10', dtype='M8[D]')
+        assert_equal(a.dtype, np.dtype('M8[D]'))
+        assert_equal(a,
+            np.array(['2010-01-05', '2010-01-06', '2010-01-07',
+                      '2010-01-08', '2010-01-09'], dtype='M8[D]'))
+
+        a = np.arange('1950-02-10', '1950-02-06', -1, dtype='M8[D]')
+        assert_equal(a.dtype, np.dtype('M8[D]'))
+        assert_equal(a,
+            np.array(['1950-02-10', '1950-02-09', '1950-02-08',
+                      '1950-02-07'], dtype='M8[D]'))
+
 class TestDateTimeData(TestCase):
 
     def test_basic(self):