Fix take_bytes()

Author: vstinner

Lib/test/test_bytes.py

@@ -1659,6 +1659,25 @@ def test_take_bytes(self):
         self.assertEqual(ba, bytearray(b'A'))
         self.assertEqual(ord(b'c'), ord('c'))
 
+    def test_take_bytes_error(self):
+        # gh-157242: If bytearray.take_bytes() fails (memory allocation
+        # failure), the bytearray must be left unchanged.
+        _testcapi = import_helper.import_module('_testcapi')
+
+        for mem_error in (0, 1):
+            for to_take in (5, None):
+                with self.subTest(mem_error=mem_error, to_take=to_take):
+                    ba = bytearray(b'0123456789')
+                    expected = ba[3:]
+                    del ba[:3]
+                    with self.assertRaises(MemoryError):
+                        try:
+                            _testcapi.set_nomemory(mem_error)
+                            ba.take_bytes(5)
+                        finally:
+                            _testcapi.remove_mem_hooks()
+                    self.assertEqual(ba, expected)
+
     @support.cpython_only  # tests an implementation detail
     def test_take_bytes_optimization(self):
         # Validate optimization around taking lots of little chunks out of a

Objects/bytearrayobject.c

@@ -1621,6 +1621,7 @@ bytearray_take_bytes_impl(PyByteArrayObject *self, PyObject *n)
     }
 
     Py_ssize_t remaining_length = size - to_take;
+
     // optimization: If taking less than leaving, just copy the small to_take
     // portion out and move ob_start.
     if (to_take < remaining_length) {
@@ -1642,15 +1643,30 @@ bytearray_take_bytes_impl(PyByteArrayObject *self, PyObject *n)
     memcpy(PyBytes_AS_STRING(remaining), self->ob_start + to_take,
            remaining_length);
 
-    // If the bytes are offset inside the buffer must first align.
-    if (self->ob_start != self->ob_bytes) {
-        memmove(self->ob_bytes, self->ob_start, to_take);
-        self->ob_start = self->ob_bytes;
-    }
+    size_t logical_offset = (size_t) (self->ob_start - self->ob_bytes);
+    if (logical_offset == 0 || remaining_length == 0) {
+        // If the bytes are offset inside the buffer must first align.
+        if (logical_offset != 0) {
+            memmove(self->ob_bytes, self->ob_start, to_take);
+            self->ob_start = self->ob_bytes;
+        }
 
-    if (_PyBytes_ResizeKeepOnError(&self->ob_bytes_object, to_take) == -1) {
-        Py_DECREF(remaining);
-        return NULL;
+        if (_PyBytes_ResizeKeepOnError(&self->ob_bytes_object, to_take) == -1) {
+            Py_DECREF(remaining);
+            return NULL;
+        }
+    }
+    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, to_take);
+        if (resized == NULL) {
+            Py_DECREF(remaining);
+            return NULL;
+        }
+        memcpy(PyBytes_AS_STRING(resized), self->ob_start, to_take);
+        Py_SETREF(self->ob_bytes_object, resized);
     }
 
     // Point the bytearray towards the buffer with the remaining data.