diff options
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/core/src/multiarray/datetime.c | 40 | ||||
| -rw-r--r-- | numpy/core/tests/test_datetime.py | 2 |
2 files changed, 23 insertions, 19 deletions
diff --git a/numpy/core/src/multiarray/datetime.c b/numpy/core/src/multiarray/datetime.c index d21bb9776..de81bcea1 100644 --- a/numpy/core/src/multiarray/datetime.c +++ b/numpy/core/src/multiarray/datetime.c @@ -3221,18 +3221,6 @@ NPY_NO_EXPORT PyArrayObject * datetime_arange(PyObject *start, PyObject *stop, PyObject *step, PyArray_Descr *dtype) { - PyArray_DatetimeMetaData meta; - /* - * Both datetime and timedelta are stored as int64, so they can - * share value variables. - */ - npy_int64 values[3]; - PyObject *objs[3]; - int type_nums[3]; - - npy_intp i, length; - PyArrayObject *ret; - npy_int64 *ret_data; /* * First normalize the input parameters so there is no Py_None, @@ -3265,6 +3253,8 @@ datetime_arange(PyObject *start, PyObject *stop, PyObject *step, /* Check if the units of the given dtype are generic, in which * case we use the code path that detects the units */ + int type_nums[3]; + PyArray_DatetimeMetaData meta; if (dtype != NULL) { PyArray_DatetimeMetaData *meta_tmp; @@ -3313,6 +3303,7 @@ datetime_arange(PyObject *start, PyObject *stop, PyObject *step, } /* Set up to convert the objects to a common datetime unit metadata */ + PyObject *objs[3]; objs[0] = start; objs[1] = stop; objs[2] = step; @@ -3333,11 +3324,22 @@ datetime_arange(PyObject *start, PyObject *stop, PyObject *step, type_nums[2] = NPY_TIMEDELTA; } - /* Convert all the arguments */ + /* Convert all the arguments + * + * Both datetime and timedelta are stored as int64, so they can + * share value variables. + */ + npy_int64 values[3]; if (convert_pyobjects_to_datetimes(3, objs, type_nums, NPY_SAME_KIND_CASTING, values, &meta) < 0) { return NULL; } + /* If no start was provided, default to 0 */ + if (start == NULL) { + /* enforced above */ + assert(type_nums[0] == NPY_TIMEDELTA); + values[0] = 0; + } /* If no step was provided, default to 1 */ if (step == NULL) { @@ -3362,6 +3364,7 @@ datetime_arange(PyObject *start, PyObject *stop, PyObject *step, } /* Calculate the array length */ + npy_intp length; if (values[2] > 0 && values[1] > values[0]) { length = (values[1] - values[0] + (values[2] - 1)) / values[2]; } @@ -3389,19 +3392,20 @@ datetime_arange(PyObject *start, PyObject *stop, PyObject *step, } /* Create the result array */ - ret = (PyArrayObject *)PyArray_NewFromDescr( - &PyArray_Type, dtype, 1, &length, NULL, - NULL, 0, NULL); + PyArrayObject *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); + npy_int64 *ret_data = (npy_int64 *)PyArray_DATA(ret); /* Create the timedeltas or datetimes */ - for (i = 0; i < length; ++i) { + for (npy_intp i = 0; i < length; ++i) { *ret_data = values[0]; values[0] += values[2]; ret_data++; diff --git a/numpy/core/tests/test_datetime.py b/numpy/core/tests/test_datetime.py index f99c0f72b..a10e6dbe0 100644 --- a/numpy/core/tests/test_datetime.py +++ b/numpy/core/tests/test_datetime.py @@ -1831,7 +1831,7 @@ class TestDateTime(object): def test_timedelta_arange_no_dtype(self): d = np.array(5, dtype="m8[D]") assert_equal(np.arange(d, d + 1), d) - assert_raises(ValueError, np.arange, d) + assert_equal(np.arange(d), np.arange(0, d)) def test_datetime_maximum_reduce(self): a = np.array(['2010-01-02', '1999-03-14', '1833-03'], dtype='M8[D]') |