gh-157242: Fix PyBytesWriter_Resize() on MemoryError (#157243)

If PyBytesWriter_Resize() fails, leave the writer unchanged.

Add a new internal _PyBytes_ResizeKeepOnError() function similar to
_PyBytes_Resize() but leaves the bytes object unchanged on error.

Add BaseWriterTest for BytesWriterTest and ByteArrayWriterTest.

Add _testcapi.PyBytesWriter_small_buffer constant.

Co-authored-by: Maurycy Pawłowski-Wieroński <maurycy@maurycy.com>
(cherry picked from commit f5dd52d)

Author: vstinner

Include/internal/pycore_bytesobject.h

@@ -73,6 +73,8 @@ _PyBytes_Repeat(char* dest, Py_ssize_t len_dest,
 */
 #define _PyBytesObject_SIZE (offsetof(PyBytesObject, ob_sval) + 1)
 
+extern int _PyBytes_ResizeKeepOnError(PyObject **pv, Py_ssize_t newsize);
+
 
 // Return the cached hash value, or -1 if not cached yet.
 static inline Py_hash_t

Lib/test/test_capi/test_bytes.py

@@ -316,11 +316,11 @@ def test_join(self):
             bytes_join(b'', NULL)
 
 
-class BytesWriterTest(unittest.TestCase):
-    RESULT_TYPE = bytes
+class BaseWriterTest:
+    RESULT_TYPE = NotImplementedError
 
     def create_writer(self, alloc=0, string=b''):
-        return _testcapi.PyBytesWriter(alloc, string, 0)
+        raise NotImplementedError
 
     def bytes_equal(self, result, expected):
         # Similar to assertEqual(), but check also that the result type
@@ -420,6 +420,13 @@ def test_format_s_memory_error(self):
         writer.write_bytes(b'after', -1)
         self.bytes_equal(writer.finish(), b'after')
 
+
+class BytesWriterTest(BaseWriterTest, unittest.TestCase):
+    RESULT_TYPE = bytes
+
+    def create_writer(self, alloc=0, string=b''):
+        return _testcapi.PyBytesWriter(alloc, string, 0)
+
     def test_example_abc(self):
         self.assertEqual(_testcapi.byteswriter_abc(), b'abc')
 
@@ -429,8 +436,26 @@ def test_example_resize(self):
     def test_example_highlevel(self):
         self.assertEqual(_testcapi.byteswriter_highlevel(), b'Hello World!')
 
-
-class ByteArrayWriterTest(BytesWriterTest):
+    def test_resize_error(self):
+        small_buffer = _testcapi.PyBytesWriter_small_buffer
+        init = b'x' * (small_buffer * 2)
+        writer = self.create_writer(len(init), init)
+        size = len(init) + 100
+        try:
+            with self.assertRaises(MemoryError):
+                _testcapi.set_nomemory(0)
+                writer.resize(size, b'')
+        finally:
+            _testcapi.remove_mem_hooks()
+        suffix = b'still working'
+        writer.write_bytes(suffix, -1)
+        self.bytes_equal(writer.finish(), init + suffix)
+
+        # Note: PyBytesWriter_Resize() leaves the buffer unchanged (no resize)
+        # if the new size is smaller than the allocated size
+
+
+class ByteArrayWriterTest(BaseWriterTest, unittest.TestCase):
     RESULT_TYPE = bytearray
 
     def create_writer(self, alloc=0, string=b''):

Misc/NEWS.d/next/C_API/2026-09-10-02-51-26.gh-issue-157242.tfXIsf.rst

@@ -0,0 +1,3 @@
+Fix :c:func:`PyBytesWriter_Resize` to handle properly memory allocation
+failure. Leave the writer unchanged on error, instead of leaving it in an
+inconsistent state. Patch by Victor Stinner.

Modules/_testcapi/bytes.c

@@ -4,6 +4,8 @@
 #include "parts.h"
 #include "util.h"
 
+#include <stddef.h>               // offsetof()
+
 #include "pycore_bytesobject.h"   // _PyBytesWriter_CreateByteArray()
 
 
@@ -373,6 +375,14 @@ byteswriter_highlevel(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
 }
 
 
+static size_t
+pybyteswriter_small_buffer_size(void)
+{
+    // PyBytesWriter.obj is the second member, small_buffer is the first member
+    return offsetof(PyBytesWriter, obj);
+}
+
+
 static PyMethodDef test_methods[] = {
     {"bytes_resize", bytes_resize, METH_VARARGS},
     {"bytes_join", bytes_join, METH_VARARGS},
@@ -399,5 +409,10 @@ _PyTestCapi_Init_Bytes(PyObject *m)
     }
     Py_DECREF(writer_type);
 
+    long size = (long)pybyteswriter_small_buffer_size();
+    if (PyModule_AddIntConstant(m, "PyBytesWriter_small_buffer", size) < 0) {
+        return -1;
+    }
+
     return 0;
 }

Modules/_testcapi/mem.c

@@ -155,18 +155,29 @@ fm_remove_hooks(void)
     }
 }
 
+static void
+fm_set_nomemory(int start, int stop)
+{
+    /* Memory allocation fails after 'start' allocation requests, and until
+     * 'stop' allocation requests except when 'stop' is negative or equal
+     * to 0 (default) in which case allocation failures never stop. */
+    FmData.start = start;
+    FmData.stop = stop;
+    FmData.count = 0;
+    fm_setup_hooks();
+}
+
 static PyObject *
 set_nomemory(PyObject *self, PyObject *args)
 {
     /* Memory allocation fails after 'start' allocation requests, and until
      * 'stop' allocation requests except when 'stop' is negative or equal
      * to 0 (default) in which case allocation failures never stop. */
-    FmData.count = 0;
-    FmData.stop = 0;
-    if (!PyArg_ParseTuple(args, "i|i", &FmData.start, &FmData.stop)) {
+    int start, stop = 0;
+    if (!PyArg_ParseTuple(args, "i|i", &start, &stop)) {
         return NULL;
     }
-    fm_setup_hooks();
+    fm_set_nomemory(start, stop);
     Py_RETURN_NONE;
 }
 
@@ -731,6 +742,129 @@ tracemalloc_track_race(PyObject *self, PyObject *args)
 }
 
 
+struct bytes_resize_tracer {
+    PyObject *create;
+    PyObject *destroy;
+};
+
+
+static int
+bytes_resize_tracer(PyObject *obj, PyRefTracerEvent event, void* data)
+{
+    if (event != PyRefTracer_CREATE && event != PyRefTracer_DESTROY) {
+        return 0;
+    }
+
+    struct bytes_resize_tracer *tracer = (struct bytes_resize_tracer*)data;
+    if (!PyBytes_Check(obj)) {
+        return 0;
+    }
+
+    switch (event) {
+        case PyRefTracer_CREATE:
+            tracer->create = obj;
+            break;
+        case PyRefTracer_DESTROY:
+            tracer->destroy = obj;
+            break;
+        default:
+            break;
+    }
+    return 0;
+}
+
+
+// When _PyBytes_Resize() resizes a bytes object in-place, check that
+// PyRefTracer_DESTROY and PyRefTracer_CREATE events are emitted.
+// If no_memory is non-zero, inject MemoryError.
+static int
+check_bytes_resize_tracer(int no_memory)
+{
+    PyObject *bytes = NULL;
+    PyRefTracer old_tracer = NULL;
+    void *old_tracer_data = NULL;
+    int restore_tracer = 0;
+
+    bytes = PyBytes_FromString("hello");
+    if (bytes == NULL) {
+        goto error;
+    }
+    assert(PyUnstable_Object_IsUniquelyReferenced(bytes));
+
+    old_tracer = PyRefTracer_GetTracer(&old_tracer_data);
+    restore_tracer = 1;
+
+    struct bytes_resize_tracer tracer = {0};
+    if (PyRefTracer_SetTracer(bytes_resize_tracer, &tracer) != 0) {
+        goto error;
+    }
+
+    PyObject *old_bytes = bytes;  // borrowed reference
+    if (no_memory) {
+        fm_set_nomemory(0, 0);
+        int res = _PyBytes_Resize(&bytes, 100);
+        assert(res < 0);
+        assert(bytes == NULL);
+        fm_remove_hooks();
+
+        assert(PyErr_ExceptionMatches(PyExc_MemoryError));
+        PyErr_Clear();
+    }
+    else {
+        if (_PyBytes_Resize(&bytes, 100) < 0) {
+            assert(bytes == NULL);
+            goto error;
+        }
+    }
+
+    if (tracer.destroy != old_bytes) {
+        PyErr_SetString(PyExc_AssertionError, "PyRefTracer_DESTROY not seen");
+        goto error;
+    }
+
+    int seen_create;
+    if (no_memory) {
+        seen_create = (tracer.create == old_bytes);
+    }
+    else {
+        seen_create = (tracer.create == bytes);
+    }
+    if (!seen_create) {
+        PyErr_SetString(PyExc_AssertionError, "PyRefTracer_CREATE not seen");
+        goto error;
+    }
+
+    Py_CLEAR(bytes);
+    if (PyRefTracer_SetTracer(old_tracer, old_tracer_data) != 0) {
+        restore_tracer = 0;
+        goto error;
+    }
+    return 0;
+
+error:
+    Py_XDECREF(bytes);
+    if (restore_tracer) {
+        if (PyRefTracer_SetTracer(old_tracer, old_tracer_data) != 0) {
+            return -1;
+        }
+    }
+    return -1;
+}
+
+
+static PyObject*
+test_bytes_resize_tracer(PyObject *self, PyObject *Py_UNUSED(ignored))
+{
+    if (check_bytes_resize_tracer(0) < 0) {
+        return NULL;
+    }
+    if (check_bytes_resize_tracer(1) < 0) {
+        return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
+
 static PyMethodDef test_methods[] = {
     {"pymem_api_misuse",              pymem_api_misuse,              METH_NOARGS},
     {"pymem_buffer_overflow",         pymem_buffer_overflow,         METH_NOARGS},
@@ -745,6 +879,7 @@ static PyMethodDef test_methods[] = {
     {"test_pymem_setrawallocators",   test_pymem_setrawallocators,   METH_NOARGS},
     {"test_pyobject_new",             test_pyobject_new,             METH_NOARGS},
     {"test_pyobject_setallocators",   test_pyobject_setallocators,   METH_NOARGS},
+    {"test_bytes_resize_tracer",      test_bytes_resize_tracer,      METH_NOARGS},
 
     // Tracemalloc tests
     {"tracemalloc_track",             tracemalloc_track,             METH_VARARGS},