Merge pull request #9174 from pitrou/pickle_generic_datetime

BUG: allow pickling generic datetime

2 files changed, 59 insertions, 2 deletions

diff --git a/numpy/core/src/multiarray/datetime.c b/numpy/core/src/multiarray/datetime.c
index 3cf9a2bd5..a6a1ba1ac 100644
--- a/numpy/core/src/multiarray/datetime.c
+++ b/numpy/core/src/multiarray/datetime.c
@@ -1718,8 +1718,6 @@ datetime_type_promotion(PyArray_Descr *type1, PyArray_Descr *type2)
  * a date time unit enum value. The 'metastr' parameter
  * is used for error messages, and may be NULL.
  *
- * Generic units have no representation as a string in this form.
- *
  * Returns 0 on success, -1 on failure.
  */
 NPY_NO_EXPORT NPY_DATETIMEUNIT
@@ -1761,6 +1759,9 @@ parse_datetime_unit_from_string(char *str, Py_ssize_t len, char *metastr)
                 return NPY_FR_as;
         }
     }
+    else if (len == 7 && !strncmp(str, "generic", 7)) {
+        return NPY_FR_GENERIC;
+    }
 
     /* If nothing matched, it's an error */
     if (metastr == NULL) {
diff --git a/numpy/core/tests/test_dtype.py b/numpy/core/tests/test_dtype.py
index 452cbd4bd..7829cb6b9 100644
--- a/numpy/core/tests/test_dtype.py
+++ b/numpy/core/tests/test_dtype.py
@@ -1,5 +1,6 @@
 from __future__ import division, absolute_import, print_function
 
+import pickle
 import sys
 
 import numpy as np
@@ -624,6 +625,61 @@ class TestDtypeAttributes(TestCase):
         assert_equal(np.dtype(user_def_subcls).name, 'user_def_subcls')
 
 
+class TestPickling(TestCase):
+
+    def check_pickling(self, dtype):
+        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+            pickled = pickle.loads(pickle.dumps(dtype, proto))
+            assert_equal(pickled, dtype)
+            assert_equal(pickled.descr, dtype.descr)
+            if dtype.metadata:
+                assert_equal(pickled.metadata, dtype.metadata)
+            else:
+                self.assertFalse(pickled.metadata)  # may be None
+            # Check the reconstructed dtype is functional
+            x = np.zeros(3, dtype=dtype)
+            y = np.zeros(3, dtype=pickled)
+            assert_equal(x, y)
+            assert_equal(x[0], y[0])
+
+    def test_builtin(self):
+        for t in [np.int, np.float, np.complex, np.int32, np.str, np.object,
+                  np.unicode, np.bool]:
+            self.check_pickling(np.dtype(t))
+
+    def test_structured(self):
+        dt = np.dtype(([('a', '>f4', (2, 1)), ('b', '<f8', (1, 3))], (2, 2)))
+        self.check_pickling(dt)
+        dt = np.dtype('i4, i1', align=True)
+        self.check_pickling(dt)
+        dt = np.dtype('i4, i1', align=False)
+        self.check_pickling(dt)
+        dt = np.dtype({
+            'names': ['A', 'B'],
+            'formats': ['f4', 'f4'],
+            'offsets': [0, 8],
+            'itemsize': 16})
+        self.check_pickling(dt)
+        dt = np.dtype({'names': ['r', 'b'],
+                       'formats': ['u1', 'u1'],
+                       'titles': ['Red pixel', 'Blue pixel']})
+        self.check_pickling(dt)
+
+    def test_datetime(self):
+        for base in ['m8', 'M8']:
+            for unit in ['', 'Y', 'M', 'W', 'D', 'h', 'm', 's', 'ms',
+                         'us', 'ns', 'ps', 'fs', 'as']:
+                dt = np.dtype('%s[%s]' % (base, unit) if unit else base)
+                self.check_pickling(dt)
+                if unit:
+                    dt = np.dtype('%s[7%s]' % (base, unit))
+                    self.check_pickling(dt)
+
+    def test_metadata(self):
+        dt = np.dtype(int, metadata={'datum': 1})
+        self.check_pickling(dt)
+
+
 def test_rational_dtype():
     # test for bug gh-5719
     a = np.array([1111], dtype=rational).astype