Merge pull request #371 from rlamy/bytes-pickle

teoliphant · teoliphant · commit fd15162fbff5 · 2012-08-02T20:52:39.000-07:00
Copy bytes object for small arrays when unpickling an array
diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c
@@ -1587,8 +1587,9 @@ array_setstate(PyArrayObject *self, PyObject *args)
         /* Check that the string is not interned */
         if (!_IsAligned(self) || swap || PyString_CHECK_INTERNED(rawdata)) {
 #else
-        /* Bytes are never interned */
-        if (!_IsAligned(self) || swap) {
+        /* Bytes should always be considered immutable, but we just grab the
+         * pointer if they are large, to save memory. */
+        if (!_IsAligned(self) || swap || (len <= 1000)) {
 #endif
             npy_intp num = PyArray_NBYTES(self);
             fa->data = PyDataMem_NEW(num);
diff --git a/numpy/core/tests/test_regression.py b/numpy/core/tests/test_regression.py
@@ -1601,6 +1601,14 @@ def test_pickle_string_overwrite(self):
         s = re.sub("a(.)", "\x01\\1", "a_")
         assert_equal(s[0], "\x01")
 
+    def test_pickle_bytes_overwrite(self):
+        if sys.version_info[0] >= 3:
+            data = np.array([1], dtype='b')
+            data = pickle.loads(pickle.dumps(data))
+            data[0] = 0xdd
+            bytestring = "\x01  ".encode('ascii')
+            assert_equal(bytestring[0:1], '\x01'.encode('ascii'))
+
     def test_structured_type_to_object(self):
         a_rec = np.array([(0,1), (3,2)], dtype='i4,i8')
         a_obj = np.empty((2,), dtype=object)
