summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--numpy/core/src/multiarray/datetime.c454
-rw-r--r--numpy/core/tests/test_datetime.py39
2 files changed, 235 insertions, 258 deletions
diff --git a/numpy/core/src/multiarray/datetime.c b/numpy/core/src/multiarray/datetime.c
index 606d69c50..04f68077c 100644
--- a/numpy/core/src/multiarray/datetime.c
+++ b/numpy/core/src/multiarray/datetime.c
@@ -4452,17 +4452,154 @@ is_any_numpy_datetime_or_timedelta(PyObject *obj)
is_any_numpy_timedelta(obj));
}
+/*
+ * Converts an array of PyObject * into datetimes and/or timedeltas,
+ * based on the values in type_nums.
+ *
+ * If inout_meta->base is -1, uses GCDs to calculate the metadata, filling
+ * in 'inout_meta' with the resulting metadata. Otherwise uses the provided
+ * 'inout_meta' for all the conversions.
+ *
+ * When obj[i] is NULL, out_value[i] will be set to NPY_DATETIME_NAT.
+ *
+ * Returns 0 on success, -1 on failure.
+ */
+NPY_NO_EXPORT int
+convert_pyobjects_to_datetimes(int count,
+ PyObject **objs, int *type_nums,
+ npy_int64 *out_values,
+ PyArray_DatetimeMetaData *inout_meta)
+{
+ int i, is_out_strict;
+ PyArray_DatetimeMetaData *meta;
+
+ /* No values trivially succeeds */
+ if (count == 0) {
+ return 0;
+ }
+
+ /* Use the inputs to resolve the unit metadata if requested */
+ if (inout_meta->base == -1) {
+ /* Allocate an array of metadata corresponding to the objects */
+ meta = PyArray_malloc(count * sizeof(PyArray_DatetimeMetaData));
+ if (meta == NULL) {
+ PyErr_NoMemory();
+ return -1;
+ }
+
+ /* Convert all the objects into timedeltas or datetimes */
+ for (i = 0; i < count; ++i) {
+ meta[i].base = -1;
+ meta[i].num = 1;
+ meta[i].events = 1;
+
+ /* NULL -> NaT */
+ if (objs[i] == NULL) {
+ out_values[i] = NPY_DATETIME_NAT;
+ meta[i].base = NPY_FR_GENERIC;
+ }
+ else if (type_nums[i] == NPY_DATETIME) {
+ if (convert_pyobject_to_datetime(&meta[i], objs[i],
+ &out_values[i]) < 0) {
+ PyArray_free(meta);
+ return -1;
+ }
+ }
+ else if (type_nums[i] == NPY_TIMEDELTA) {
+ if (convert_pyobject_to_timedelta(&meta[i], objs[i],
+ &out_values[i]) < 0) {
+ PyArray_free(meta);
+ return -1;
+ }
+ }
+ else {
+ PyErr_SetString(PyExc_ValueError,
+ "convert_pyobjects_to_datetimes requires that "
+ "all the type_nums provided be datetime or timedelta");
+ PyArray_free(meta);
+ return -1;
+ }
+ }
+
+ /* Merge all the metadatas, starting with the first one */
+ *inout_meta = meta[0];
+ is_out_strict = (type_nums[0] == NPY_TIMEDELTA);
+
+ for (i = 1; i < count; ++i) {
+ if (compute_datetime_metadata_greatest_common_divisor(
+ &meta[i], inout_meta, inout_meta,
+ type_nums[i] == NPY_TIMEDELTA,
+ is_out_strict) < 0) {
+ PyArray_free(meta);
+ return -1;
+ }
+ is_out_strict = is_out_strict || (type_nums[i] == NPY_TIMEDELTA);
+ }
+
+ /* Convert all the values into the resolved unit metadata */
+ for (i = 0; i < count; ++i) {
+ if (type_nums[i] == NPY_DATETIME) {
+ if (cast_datetime_to_datetime(&meta[i], inout_meta,
+ out_values[i], &out_values[i]) < 0) {
+ PyArray_free(meta);
+ return -1;
+ }
+ }
+ else if (type_nums[i] == NPY_TIMEDELTA) {
+ if (cast_timedelta_to_timedelta(&meta[i], inout_meta,
+ out_values[i], &out_values[i]) < 0) {
+ PyArray_free(meta);
+ return -1;
+ }
+ }
+ }
+
+ PyArray_free(meta);
+ }
+ /* Otherwise convert to the provided unit metadata */
+ else {
+ /* Convert all the objects into timedeltas or datetimes */
+ for (i = 0; i < count; ++i) {
+ /* NULL -> NaT */
+ if (objs[i] == NULL) {
+ out_values[i] = NPY_DATETIME_NAT;
+ }
+ else if (type_nums[i] == NPY_DATETIME) {
+ if (convert_pyobject_to_datetime(inout_meta, objs[i],
+ &out_values[i]) < 0) {
+ return -1;
+ }
+ }
+ else if (type_nums[i] == NPY_TIMEDELTA) {
+ if (convert_pyobject_to_timedelta(inout_meta, objs[i],
+ &out_values[i]) < 0) {
+ return -1;
+ }
+ }
+ else {
+ PyErr_SetString(PyExc_ValueError,
+ "convert_pyobjects_to_datetimes requires that "
+ "all the type_nums provided be datetime or timedelta");
+ return -1;
+ }
+ }
+ }
+
+ return 0;
+}
+
NPY_NO_EXPORT PyArrayObject *
datetime_arange(PyObject *start, PyObject *stop, PyObject *step,
PyArray_Descr *dtype)
{
PyArray_DatetimeMetaData meta;
- int type_num = 0;
/*
* Both datetime and timedelta are stored as int64, so they can
* share value variables.
*/
- npy_int64 start_value = 0, stop_value = 0, step_value;
+ npy_int64 values[3];
+ PyObject *objs[3];
+ int type_nums[3];
npy_intp i, length;
PyArrayObject *ret;
@@ -4502,8 +4639,8 @@ datetime_arange(PyObject *start, PyObject *stop, PyObject *step,
if (dtype != NULL) {
PyArray_DatetimeMetaData *meta_tmp;
- type_num = dtype->type_num;
- if (type_num != NPY_DATETIME && type_num != NPY_TIMEDELTA) {
+ type_nums[0] = dtype->type_num;
+ if (type_nums[0] != NPY_DATETIME && type_nums[0] != NPY_TIMEDELTA) {
PyErr_SetString(PyExc_ValueError,
"datetime_arange was given a non-datetime dtype");
return NULL;
@@ -4520,6 +4657,7 @@ datetime_arange(PyObject *start, PyObject *stop, PyObject *step,
*/
if (meta_tmp->base == NPY_FR_GENERIC) {
dtype = NULL;
+ meta.base = -1;
}
/* Otherwise use the provided metadata */
else {
@@ -4528,280 +4666,80 @@ datetime_arange(PyObject *start, PyObject *stop, PyObject *step,
}
else {
if (is_any_numpy_datetime(start) || is_any_numpy_datetime(stop)) {
- type_num = NPY_DATETIME;
+ type_nums[0] = NPY_DATETIME;
}
else {
- type_num = NPY_TIMEDELTA;
+ type_nums[0] = NPY_TIMEDELTA;
}
- }
- /*
- * Now figure out whether we're doing a datetime or a timedelta
- * arange, and get the metadata for start and stop.
- */
- if (dtype != NULL) {
- if (type_num == NPY_TIMEDELTA) {
- /* Get 'start', defaulting to 0 if not provided */
- if (start == NULL) {
- start_value = 0;
- }
- else {
- if (convert_pyobject_to_timedelta(&meta, start,
- &start_value) < 0) {
- return NULL;
- }
- }
- /* Get 'stop' */
- if (convert_pyobject_to_timedelta(&meta, stop,
- &stop_value) < 0) {
- return NULL;
- }
- }
- else {
- /* Get 'start', raising an exception if it is not provided */
- if (start == NULL) {
- PyErr_SetString(PyExc_ValueError,
- "arange requires both a start and a stop for "
- "NumPy datetime64 ranges");
- return NULL;
- }
- else {
- if (convert_pyobject_to_datetime(&meta, start,
- &start_value) < 0) {
- return NULL;
- }
- }
- /*
- * Get 'stop', treating it specially if it's an integer
- * or a timedelta.
- */
- if (PyInt_Check(stop) ||
- PyLong_Check(stop) ||
- PyArray_IsScalar(stop, Integer) ||
- is_any_numpy_timedelta(stop)) {
- if (convert_pyobject_to_timedelta(&meta, stop,
- &stop_value) < 0) {
- return NULL;
- }
+ meta.base = -1;
+ }
- /* Add the start value to make it a datetime */
- stop_value += start_value;
- }
- else {
- if (convert_pyobject_to_datetime(&meta, stop,
- &stop_value) < 0) {
- return NULL;
- }
- }
- }
+ if (type_nums[0] == NPY_DATETIME && start == NULL) {
+ PyErr_SetString(PyExc_ValueError,
+ "arange requires both a start and a stop for "
+ "NumPy datetime64 ranges");
+ return NULL;
+ }
- /* Get the 'step' value */
- if (step == NULL) {
- step_value = 1;
- }
- else {
- if (convert_pyobject_to_timedelta(&meta, step,
- &step_value) < 0) {
- return NULL;
- }
- }
+ /* Set up to convert the objects to a common datetime unit metadata */
+ objs[0] = start;
+ objs[1] = stop;
+ objs[2] = step;
+ if (type_nums[0] == NPY_TIMEDELTA) {
+ type_nums[1] = NPY_TIMEDELTA;
+ type_nums[2] = NPY_TIMEDELTA;
}
- /* if dtype is NULL */
else {
- /*
- * Start meta as generic units, so taking the gcd
- * does the right thing
- */
- meta.base = NPY_FR_GENERIC;
-
- /* Timedelta arange */
- if (type_num == NPY_TIMEDELTA) {
- PyArray_DatetimeMetaData meta_tmp, meta_mrg;
-
- /* Convert start to a timedelta */
- if (start == NULL) {
- start_value = 0;
- }
- else {
- /* Directly copy to meta, no point in gcd for 'start' */
- meta.base = -1;
- /* Convert 'start' to a timedelta */
- if (convert_pyobject_to_timedelta(&meta, start,
- &start_value) < 0) {
- return NULL;
- }
- }
-
- /* Temporary metadata, to merge with 'meta' */
- meta_tmp.base = -1;
- /* Convert 'stop' to a timedelta */
- if (convert_pyobject_to_timedelta(&meta_tmp, stop,
- &stop_value) < 0) {
- return NULL;
- }
-
- /* Merge the metadata */
- if (compute_datetime_metadata_greatest_common_divisor(
- &meta, &meta_tmp, &meta_mrg, 1, 1) < 0) {
- return NULL;
- }
-
- /* Cast 'start' and 'stop' to the merged metadata */
- if (cast_timedelta_to_timedelta(&meta, &meta_mrg,
- start_value, &start_value) < 0) {
- return NULL;
- }
- if (cast_timedelta_to_timedelta(&meta_tmp, &meta_mrg,
- stop_value, &stop_value) < 0) {
- return NULL;
- }
-
- meta = meta_mrg;
-
- if (step == NULL) {
- step_value = 1;
- }
- else {
- /* Temporary metadata, to merge with 'meta' */
- meta_tmp.base = -1;
- /* Convert 'step' to a timedelta */
- if (convert_pyobject_to_timedelta(&meta_tmp, step,
- &step_value) < 0) {
- return NULL;
- }
-
- /* Merge the metadata */
- if (compute_datetime_metadata_greatest_common_divisor(
- &meta, &meta_tmp, &meta_mrg, 1, 1) < 0) {
- return NULL;
- }
-
- /* Cast 'start', 'stop', and 'step' to the merged metadata */
- if (cast_timedelta_to_timedelta(&meta, &meta_mrg,
- start_value, &start_value) < 0) {
- return NULL;
- }
- if (cast_timedelta_to_timedelta(&meta, &meta_mrg,
- stop_value, &stop_value) < 0) {
- return NULL;
- }
- if (cast_timedelta_to_timedelta(&meta_tmp, &meta_mrg,
- step_value, &step_value) < 0) {
- return NULL;
- }
-
- meta = meta_mrg;
- }
+ if (PyInt_Check(objs[1]) ||
+ PyLong_Check(objs[1]) ||
+ PyArray_IsScalar(objs[1], Integer) ||
+ is_any_numpy_timedelta(objs[1])) {
+ type_nums[1] = NPY_TIMEDELTA;
}
- /* Datetime arange */
else {
- PyArray_DatetimeMetaData meta_tmp, meta_mrg;
-
- /* Get 'start', raising an exception if it is not provided */
- if (start == NULL) {
- PyErr_SetString(PyExc_ValueError,
- "arange requires both a start and a stop for "
- "NumPy datetime64 ranges");
- return NULL;
- }
- else {
- /* Directly copy to meta, no point in gcd for 'start' */
- meta.base = -1;
- /* Convert 'start' to a datetime */
- if (convert_pyobject_to_datetime(&meta, start,
- &start_value) < 0) {
- return NULL;
- }
- }
-
- /* Temporary metadata, to merge with 'meta' */
- meta_tmp.base = -1;
- /* Convert 'stop' to a datetime */
- if (convert_pyobject_to_datetime(&meta_tmp, stop,
- &stop_value) < 0) {
- return NULL;
- }
-
- /* Merge the metadata */
- if (compute_datetime_metadata_greatest_common_divisor(
- &meta, &meta_tmp, &meta_mrg, 0, 0) < 0) {
- return NULL;
- }
-
- /* Cast 'start' and 'stop' to the merged metadata */
- if (cast_datetime_to_datetime(&meta, &meta_mrg,
- start_value, &start_value) < 0) {
- return NULL;
- }
- if (cast_datetime_to_datetime(&meta_tmp, &meta_mrg,
- stop_value, &stop_value) < 0) {
- return NULL;
- }
-
- meta = meta_mrg;
-
- if (step == NULL) {
- step_value = 1;
- }
- else {
- /* Temporary metadata, to merge with 'meta' */
- meta_tmp.base = -1;
- /* Convert 'step' to a timedelta */
- if (convert_pyobject_to_timedelta(&meta_tmp, step,
- &step_value) < 0) {
- return NULL;
- }
+ type_nums[1] = NPY_DATETIME;
+ }
+ type_nums[2] = NPY_TIMEDELTA;
+ }
- /* Merge the metadata */
- if (compute_datetime_metadata_greatest_common_divisor(
- &meta, &meta_tmp, &meta_mrg, 0, 1) < 0) {
- return NULL;
- }
+ /* Convert all the arguments */
+ if (convert_pyobjects_to_datetimes(3, objs, type_nums,
+ values, &meta) < 0) {
+ return NULL;
+ }
- /* Cast 'start', 'stop', and 'step' to the merged metadata */
- if (cast_datetime_to_datetime(&meta, &meta_mrg,
- start_value, &start_value) < 0) {
- return NULL;
- }
- if (cast_datetime_to_datetime(&meta, &meta_mrg,
- stop_value, &stop_value) < 0) {
- return NULL;
- }
- if (cast_timedelta_to_timedelta(&meta_tmp, &meta_mrg,
- step_value, &step_value) < 0) {
- return NULL;
- }
+ /* If no step was provided, default to 1 */
+ if (step == NULL) {
+ values[2] = 1;
+ }
- meta = meta_mrg;
- }
- }
+ /*
+ * In the case of arange(datetime, timedelta), convert
+ * the timedelta into a datetime by adding the start datetime.
+ */
+ if (type_nums[0] == NPY_DATETIME && type_nums[1] == NPY_TIMEDELTA) {
+ values[1] += values[0];
}
/* 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) {
+ if (values[0] == NPY_DATETIME_NAT ||
+ values[1] == NPY_DATETIME_NAT ||
+ values[2] == 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;
- }
+ if (values[2] > 0 && values[1] > values[0]) {
+ length = (values[1] - values[0] + (values[2] - 1)) / values[2];
}
- else if (step_value < 0) {
- if (stop_value < start_value) {
- length = (start_value - stop_value) / (-step_value);
- }
- else {
- length = 0;
- }
+ else if (values[2] < 0 && values[1] < values[0]) {
+ length = (values[1] - values[0] + (values[2] + 1)) / values[2];
+ }
+ else if (values[2] != 0) {
+ length = 0;
}
else {
PyErr_SetString(PyExc_ValueError,
@@ -4814,7 +4752,7 @@ datetime_arange(PyObject *start, PyObject *stop, PyObject *step,
Py_INCREF(dtype);
}
else {
- dtype = create_datetime_dtype(type_num, &meta);
+ dtype = create_datetime_dtype(type_nums[0], &meta);
if (dtype == NULL) {
return NULL;
}
@@ -4834,8 +4772,8 @@ datetime_arange(PyObject *start, PyObject *stop, PyObject *step,
/* Create the timedeltas or datetimes */
for (i = 0; i < length; ++i) {
- *ret_data = start_value;
- start_value += step_value;
+ *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 a8f1be8fc..f75b5c1a9 100644
--- a/numpy/core/tests/test_datetime.py
+++ b/numpy/core/tests/test_datetime.py
@@ -1120,6 +1120,45 @@ class TestDateTime(TestCase):
assert_equal(a,
np.datetime64('1969-05') + np.arange(12, step=2))
+ # datetime, integer|timedelta works as well
+ # produces arange (start, start + stop) in this case
+ a = np.arange('1969', 18, 3, dtype='M8')
+ assert_equal(a.dtype, np.dtype('M8[Y]'))
+ assert_equal(a,
+ np.datetime64('1969') + np.arange(18, step=3))
+ a = np.arange('1969-12-19', 22, np.timedelta64(2), dtype='M8')
+ assert_equal(a.dtype, np.dtype('M8[D]'))
+ assert_equal(a,
+ np.datetime64('1969-12-19') + np.arange(22, step=2))
+
+ # Step of 0 is disallowed
+ assert_raises(ValueError, np.arange, np.datetime64('today'),
+ np.datetime64('today') + 3, 0)
+ # Promotion across nonlinear unit boundaries is disallowed
+ assert_raises(TypeError, np.arange, np.datetime64('2011-03-01','D'),
+ np.timedelta64(5,'M'))
+ assert_raises(TypeError, np.arange,
+ np.datetime64('2012-02-03T14Z','s'),
+ np.timedelta64(5,'Y'))
+
+ def test_timedelta_arange(self):
+ a = np.arange(3, 10, dtype='m8')
+ assert_equal(a.dtype, np.dtype('m8'))
+ assert_equal(a, np.timedelta64(0) + np.arange(3, 10))
+
+ a = np.arange(np.timedelta64(3,'s'), 10, 2, dtype='m8')
+ assert_equal(a.dtype, np.dtype('m8[s]'))
+ assert_equal(a, np.timedelta64(0, 's') + np.arange(3, 10, 2))
+
+ # Step of 0 is disallowed
+ assert_raises(ValueError, np.arange, np.timedelta64(0),
+ np.timedelta64(5), 0)
+ # Promotion across nonlinear unit boundaries is disallowed
+ assert_raises(TypeError, np.arange, np.timedelta64(0,'D'),
+ np.timedelta64(5,'M'))
+ assert_raises(TypeError, np.arange, np.timedelta64(0,'Y'),
+ np.timedelta64(5,'D'))
+
class TestDateTimeData(TestCase):
def test_basic(self):