BUG: Allow empty strings to be pickled

2 files changed, 22 insertions, 8 deletions

diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c
index efa97dd65..a908e648d 100644
--- a/numpy/core/src/multiarray/methods.c
+++ b/numpy/core/src/multiarray/methods.c
@@ -13,6 +13,7 @@
 #include "npy_import.h"
 #include "ufunc_override.h"
 #include "common.h"
+#include "templ_common.h" /* for npy_mul_with_overflow_intp */
 #include "ctors.h"
 #include "calculation.h"
 #include "convert_datatype.h"
@@ -1671,6 +1672,8 @@ array_setstate(PyArrayObject *self, PyObject *args)
     Py_ssize_t len;
     npy_intp size, dimensions[NPY_MAXDIMS];
     int nd;
+    npy_intp nbytes;
+    int overflowed;
 
     PyArrayObject_fields *fa = (PyArrayObject_fields *)self;
 
@@ -1712,13 +1715,15 @@ array_setstate(PyArrayObject *self, PyObject *args)
         return NULL;
     }
     size = PyArray_MultiplyList(dimensions, nd);
-    if (PyArray_DESCR(self)->elsize == 0) {
-        PyErr_SetString(PyExc_ValueError, "Invalid data-type size.");
-        return NULL;
+    if (size < 0) {
+        /* More items than are addressable */
+        return PyErr_NoMemory();
     }
-    if (size < 0 || size > NPY_MAX_INTP / PyArray_DESCR(self)->elsize) {
-        PyErr_NoMemory();
-        return NULL;
+    overflowed = npy_mul_with_overflow_intp(
+        &nbytes, size, PyArray_DESCR(self)->elsize);
+    if (overflowed) {
+        /* More bytes than are addressable */
+        return PyErr_NoMemory();
     }
 
     if (PyDataType_FLAGCHK(typecode, NPY_LIST_PICKLE)) {
@@ -1760,7 +1765,7 @@ array_setstate(PyArrayObject *self, PyObject *args)
             return NULL;
         }
 
-        if ((len != (PyArray_DESCR(self)->elsize * size))) {
+        if (len != nbytes) {
             PyErr_SetString(PyExc_ValueError,
                             "buffer size does not"  \
                             " match array size");
@@ -1822,7 +1827,7 @@ array_setstate(PyArrayObject *self, PyObject *args)
             }
             if (swap) {
                 /* byte-swap on pickle-read */
-                npy_intp numels = num / PyArray_DESCR(self)->elsize;
+                npy_intp numels = PyArray_SIZE(self);
                 PyArray_DESCR(self)->f->copyswapn(PyArray_DATA(self),
                                         PyArray_DESCR(self)->elsize,
                                         datastr, PyArray_DESCR(self)->elsize,
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
index 8bb91619e..9cad2e9db 100644
--- a/numpy/core/tests/test_multiarray.py
+++ b/numpy/core/tests/test_multiarray.py
@@ -1262,6 +1262,15 @@ class TestZeroSizeFlexible(object):
             # viewing as any non-empty type gives an empty result
             assert_equal(zs.view((dt, 1)).shape, (0,))
 
+    def test_pickle(self):
+        import pickle
+        for dt in [bytes, np.void, unicode]:
+            zs = self._zeros(10, dt)
+            p = pickle.dumps(zs)
+            zs2 = pickle.loads(p)
+
+            assert_equal(zs.dtype, zs2.dtype)
+
 
 class TestMethods(object):
     def test_compress(self):