More careful checking of datetime metadata in dtype __setstate__()

2 files changed, 27 insertions, 9 deletions

diff --git a/numpy/core/src/multiarray/descriptor.c b/numpy/core/src/multiarray/descriptor.c
index adf671262..43f8f57b8 100644
--- a/numpy/core/src/multiarray/descriptor.c
+++ b/numpy/core/src/multiarray/descriptor.c
@@ -2530,25 +2530,35 @@ arraydescr_setstate(PyArray_Descr *self, PyObject *args)
         metadata = NULL;
     }
 
-    Py_XDECREF(self->metadata);
     if (PyDataType_ISDATETIME(self) && (metadata != NULL)) {
-        PyArray_DatetimeMetaData *dt_data;
+        PyObject *new_metadata, *errmsg;
+        PyArray_DatetimeMetaData temp_dt_data;
 
-        /* The Python metadata */
-        self->metadata = PyTuple_GET_ITEM(metadata, 0);
-        Py_XINCREF(self->metadata);
+        if ((! PyTuple_Check(metadata)) || (PyTuple_Size(metadata) != 2)) {
+            errmsg = PyUString_FromString("Invalid datetime dtype (metadata, c_metadata): ");
+            PyUString_ConcatAndDel(&errmsg, PyObject_Repr(metadata));
+            PyErr_SetObject(PyExc_ValueError, errmsg);
+            return NULL;
+        }
 
-        /* The datetime metadata */
-        dt_data = &(((PyArray_DatetimeDTypeMetaData *)self->c_metadata)->meta);
         if (convert_datetime_metadata_tuple_to_datetime_metadata(
                                     PyTuple_GET_ITEM(metadata, 1),
-                                    dt_data) < 0) {
+                                    &temp_dt_data) < 0) {
             return NULL;
         }
+
+        new_metadata = PyTuple_GET_ITEM(metadata, 0);
+        Py_XINCREF(new_metadata);
+        Py_XDECREF(self->metadata);
+        self->metadata = new_metadata;
+        memcpy((char *) &((PyArray_DatetimeDTypeMetaData *)self->c_metadata)->meta,
+               (char *) &temp_dt_data,
+               sizeof(PyArray_DatetimeMetaData));
     }
     else {
-        self->metadata = metadata;
         Py_XINCREF(metadata);
+        Py_XDECREF(self->metadata);
+        self->metadata = metadata;
     }
 
     Py_INCREF(Py_None);
diff --git a/numpy/core/tests/test_datetime.py b/numpy/core/tests/test_datetime.py
index cb62182cd..fbde86b57 100644
--- a/numpy/core/tests/test_datetime.py
+++ b/numpy/core/tests/test_datetime.py
@@ -547,6 +547,14 @@ class TestDateTime(TestCase):
               "I1\nI1\nI1\ntp7\ntp8\ntp9\nb."
         assert_equal(pickle.loads(asbytes(pkl)), np.dtype('>M8[us]'))
 
+    def test_setstate(self):
+        "Verify that datetime dtype __setstate__ can handle bad arguments"
+        dt = np.dtype('>M8[us]')
+        assert_raises(ValueError, dt.__setstate__, (4, '>', None, None, None, -1, -1, 0, 1))
+        assert (dt.__reduce__()[2] == np.dtype('>M8[us]').__reduce__()[2])
+        assert_raises(TypeError, dt.__setstate__, (4, '>', None, None, None, -1, -1, 0, ({}, 'xxx')))
+        assert (dt.__reduce__()[2] == np.dtype('>M8[us]').__reduce__()[2])
+
     def test_dtype_promotion(self):
         # datetime <op> datetime computes the metadata gcd
         # timedelta <op> timedelta computes the metadata gcd