diff options
| -rw-r--r-- | numpy/core/src/multiarray/datetime.c | 21 | ||||
| -rw-r--r-- | numpy/core/src/multiarray/descriptor.c | 70 | ||||
| -rw-r--r-- | numpy/core/tests/test_datetime.py | 9 |
3 files changed, 87 insertions, 13 deletions
diff --git a/numpy/core/src/multiarray/datetime.c b/numpy/core/src/multiarray/datetime.c index 1eb100756..2a29015cd 100644 --- a/numpy/core/src/multiarray/datetime.c +++ b/numpy/core/src/multiarray/datetime.c @@ -1386,7 +1386,8 @@ compute_datetime_metadata_greatest_common_divisor( * Years, Months, and Business days are incompatible with * all other units. */ - if (meta1->base == NPY_FR_Y || meta1->base == NPY_FR_M || + if (meta1->base == NPY_FR_Y || + meta1->base == NPY_FR_M || meta1->base == NPY_FR_B || meta2->base == NPY_FR_Y || meta2->base == NPY_FR_M || @@ -1552,6 +1553,7 @@ convert_datetime_metadata_tuple_to_metacobj(PyObject *tuple) int den = 1; if (!PyTuple_Check(tuple)) { + PyObject_Print(tuple, stderr, 0); PyErr_SetString(PyExc_TypeError, "Require tuple for tuple->metacobj conversion"); return NULL; @@ -1579,12 +1581,29 @@ convert_datetime_metadata_tuple_to_metacobj(PyObject *tuple) /* Convert the values to longs */ dt_data->num = PyInt_AsLong(PyTuple_GET_ITEM(tuple, 1)); + if (dt_data->num == -1 && PyErr_Occurred()) { + PyArray_free(dt_data); + return NULL; + } + if (tuple_size == 3) { dt_data->events = PyInt_AsLong(PyTuple_GET_ITEM(tuple, 2)); + if (dt_data->events == -1 && PyErr_Occurred()) { + PyArray_free(dt_data); + return NULL; + } } else { den = PyInt_AsLong(PyTuple_GET_ITEM(tuple, 2)); + if (den == -1 && PyErr_Occurred()) { + PyArray_free(dt_data); + return NULL; + } dt_data->events = PyInt_AsLong(PyTuple_GET_ITEM(tuple, 3)); + if (dt_data->events == -1 && PyErr_Occurred()) { + PyArray_free(dt_data); + return NULL; + } } if (dt_data->num <= 0 || dt_data->events <= 0 || den <= 0) { diff --git a/numpy/core/src/multiarray/descriptor.c b/numpy/core/src/multiarray/descriptor.c index 58f573fef..d42624f68 100644 --- a/numpy/core/src/multiarray/descriptor.c +++ b/numpy/core/src/multiarray/descriptor.c @@ -1740,7 +1740,55 @@ arraydescr_new(PyTypeObject *NPY_UNUSED(subtype), PyObject *args, PyObject *kwds return (PyObject *)conv; } -/* return a tuple of (callable object, args, state). */ +/* + * Return a tuple of + * (cleaned metadata dictionary, tuple with (str, num, events)) + */ +static PyObject * +_get_pickleabletype_from_datetime_metadata(PyArray_Descr *dtype) +{ + PyObject *newdict; + PyObject *ret, *dt_tuple; + PyArray_DatetimeMetaData *meta; + + /* Create the 2-item tuple to return */ + ret = PyTuple_New(2); + if (ret == NULL) { + return NULL; + } + + /* Make a cleaned copy of the metadata dictionary */ + newdict = PyDict_Copy(dtype->metadata); + if (newdict == NULL) { + Py_DECREF(ret); + return NULL; + } + PyDict_DelItemString(newdict, NPY_METADATA_DTSTR); + PyTuple_SET_ITEM(ret, 0, newdict); + + /* Convert the datetime metadata into a tuple */ + meta = get_datetime_metadata_from_dtype(dtype); + if (meta == NULL) { + Py_DECREF(ret); + return NULL; + } + dt_tuple = convert_datetime_metadata_to_tuple(meta); + if (dt_tuple == NULL) { + Py_DECREF(ret); + return NULL; + } + PyTuple_SET_ITEM(ret, 1, dt_tuple); + + return ret; +} + +/* + * return a tuple of (callable object, args, state). + * + * TODO: This method needs to change so that unpickling doesn't + * use __setstate__. This is required for the dtype + * to be an immutable object. + */ static PyObject * arraydescr_reduce(PyArray_Descr *self, PyObject *NPY_UNUSED(args)) { @@ -1801,19 +1849,19 @@ arraydescr_reduce(PyArray_Descr *self, PyObject *NPY_UNUSED(args)) state = PyTuple_New(9); PyTuple_SET_ITEM(state, 0, PyInt_FromLong(version)); if (PyDataType_ISDATETIME(self)) { - PyArray_DatetimeMetaData *meta; - PyObject *dt_tuple; - meta = get_datetime_metadata_from_dtype(self); - if (meta == NULL) { - Py_DECREF(ret); - return NULL; - } - dt_tuple = convert_datetime_metadata_to_tuple(meta); - if (dt_tuple == NULL) { + PyObject *newobj; + /* Handle CObject in NPY_METADATA_DTSTR key separately */ + /* + * newobj is a tuple of cleaned metadata dictionary + * and tuple of date_time info (str, num, den, events) + */ + newobj = _get_pickleabletype_from_datetime_metadata(self); + if (newobj == NULL) { + Py_DECREF(state); Py_DECREF(ret); return NULL; } - PyTuple_SET_ITEM(state, 8, dt_tuple); + PyTuple_SET_ITEM(state, 8, newobj); } else { Py_INCREF(self->metadata); diff --git a/numpy/core/tests/test_datetime.py b/numpy/core/tests/test_datetime.py index 6c37b01f8..f9209bf71 100644 --- a/numpy/core/tests/test_datetime.py +++ b/numpy/core/tests/test_datetime.py @@ -1,4 +1,4 @@ -from os import path +import os, pickle import numpy as np from numpy.testing import * from numpy.compat import asbytes @@ -25,6 +25,13 @@ class TestDateTime(TestCase): assert_raises(TypeError, np.dtype, 'M16') assert_raises(TypeError, np.dtype, 'm16') + def test_pickle(self): + # Check that pickle roundtripping works + dt = np.dtype('M8[7D]//3') + assert_equal(dt, pickle.loads(pickle.dumps(dt))) + dt = np.dtype('M8[B]') + assert_equal(dt, pickle.loads(pickle.dumps(dt))) + def test_dtype_promotion(self): # datetime <op> datetime requires matching units assert_equal(np.promote_types(np.dtype('M8[Y]'), np.dtype('M8[Y]')), |