diff options
| -rw-r--r-- | numpy/core/src/multiarray/_datetime.h | 14 | ||||
| -rw-r--r-- | numpy/core/src/multiarray/ctors.c | 9 | ||||
| -rw-r--r-- | numpy/core/src/multiarray/datetime.c | 73 |
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; +} |