Fix resize() when shrinking the bytearray

Author: vstinner

Lib/test/test_bytes.py

@@ -1563,24 +1563,36 @@ def test_resize_error(self):
         # Simple bytearray
         data = b'some data'
         ba = bytearray(data)
-        try:
-            with self.assertRaises(MemoryError):
+        with self.assertRaises(MemoryError):
+            try:
                 _testcapi.set_nomemory(0)
                 ba.resize(1024)
-        finally:
-            _testcapi.remove_mem_hooks()
+            finally:
+                _testcapi.remove_mem_hooks()
         self.assertEqual(ba, bytearray(data))
 
-        # bytearray with non-zero logical start
+        # growing bytearray with non-zero logical start
         ba = bytearray(b'0123456789')
         expected = ba[3:]
         del ba[:3]
-        try:
-            with self.assertRaises(MemoryError):
+        with self.assertRaises(MemoryError):
+            try:
                 _testcapi.set_nomemory(0)
                 ba.resize(1024)
-        finally:
-            _testcapi.remove_mem_hooks()
+            finally:
+                _testcapi.remove_mem_hooks()
+        self.assertEqual(ba, expected)
+
+        # shrink bytearray with non-zero logical start
+        ba = bytearray(b'0123456789')
+        expected = ba[3:]
+        del ba[:3]
+        with self.assertRaises(MemoryError):
+            try:
+                _testcapi.set_nomemory(0)
+                ba.resize(1)
+            finally:
+                _testcapi.remove_mem_hooks()
         self.assertEqual(ba, expected)
 
     def test_take_bytes(self):

Objects/bytearrayobject.c

@@ -281,19 +281,31 @@ bytearray_resize_lock_held(PyObject *self, Py_ssize_t requested_size)
         return -1;
     }
 
-    /* Re-align data to the start of the allocation. */
-    if (logical_offset > 0) {
-        /* optimization tradeoff: This is faster than a new allocation when
-           the number of bytes being removed in a resize is small; for large
-           size changes it may be better to just make a new bytes object as
-           _PyBytes_Resize will do a malloc + memcpy internally. */
-        memmove(obj->ob_bytes, obj->ob_start,
-                Py_MIN(requested_size, Py_SIZE(self)));
-        obj->ob_start = obj->ob_bytes;
-    }
-
-    if (_PyBytes_ResizeKeepOnError(&obj->ob_bytes_object, alloc) < 0) {
-        return -1;
+    if (logical_offset == 0 || requested_size >= Py_SIZE(self)) {
+        /* Re-align data to the start of the allocation. */
+        if (logical_offset > 0) {
+            /* optimization tradeoff: This is faster than a new allocation when
+               the number of bytes being removed in a resize is small; for large
+               size changes it may be better to just make a new bytes object as
+               _PyBytes_Resize will do a malloc + memcpy internally. */
+            memmove(obj->ob_bytes, obj->ob_start, Py_SIZE(self));
+            obj->ob_start = obj->ob_bytes;
+        }
+
+        if (_PyBytes_ResizeKeepOnError(&obj->ob_bytes_object, alloc) < 0) {
+            return -1;
+        }
+    }
+    else {
+        // Using memmove() would be unsafe, since _PyBytes_ResizeKeepOnError()
+        // failure code path would be unable to restore the bytearray to its
+        // previous state.
+        PyObject *resized = PyBytes_FromStringAndSize(NULL, requested_size);
+        if (resized == NULL) {
+            return -1;
+        }
+        memcpy(PyBytes_AS_STRING(resized), obj->ob_start, requested_size);
+        Py_SETREF(obj->ob_bytes_object, resized);
     }
 
     bytearray_reinit_from_bytes(obj, size, alloc);