[3.15] gh-156939: Fix two bugs in PyBytesWriter (#158404) (#158409)

* gh-156939: Fix two bugs in PyBytesWriter (#158404)

Fix PyBytesWriter_Format() error handling: restore the writer to its
previous state on error. Add an unit test injecting MemoryError.

Fix PyBytesWriter_FinishWithSize() to return an empty bytearray when
use_bytearray is set and size is 0, instead of always returning an
empty bytes object.

Fix test_capi.test_bytes: always check writer finish result type, to
make sure that _PyBytesWriter_CreateByteArray() always create
bytearray objects.

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
(cherry picked from commit 24ea85a)

Author: vstinner

Lib/test/test_capi/test_bytes.py

@@ -1,6 +1,7 @@
 import sys
 import unittest
 from test.support import import_helper
+from test import support
 
 _testlimitedcapi = import_helper.import_module('_testlimitedcapi')
 _testcapi = import_helper.import_module('_testcapi')
@@ -316,83 +317,108 @@ def test_join(self):
 
 
 class BytesWriterTest(unittest.TestCase):
-    result_type = bytes
+    RESULT_TYPE = bytes
 
     def create_writer(self, alloc=0, string=b''):
         return _testcapi.PyBytesWriter(alloc, string, 0)
 
+    def bytes_equal(self, result, expected):
+        # Similar to assertEqual(), but check also that the result type
+        # is RESULT_TYPE.
+        self.assertEqual(result, expected)
+        self.assertEqual(type(result), self.RESULT_TYPE)
+
     def test_create(self):
         # Test PyBytesWriter_Create()
         writer = self.create_writer()
         self.assertEqual(writer.get_size(), 0)
-        self.assertEqual(writer.finish(), self.result_type(b''))
+        self.bytes_equal(writer.finish(), b'')
 
         writer = self.create_writer(3, b'abc')
         self.assertEqual(writer.get_size(), 3)
-        self.assertEqual(writer.finish(), self.result_type(b'abc'))
+        self.bytes_equal(writer.finish(), b'abc')
 
     def test_finish_with_size(self):
         # Test PyBytesWriter_FinishWithSize()
         writer = self.create_writer(10, b'abc')
         self.assertEqual(writer.get_size(), 10)
-        self.assertEqual(writer.finish_with_size(3), self.result_type(b'abc'))
+        self.bytes_equal(writer.finish_with_size(3), b'abc')
 
         writer = self.create_writer(3, b'abc')
         with self.assertRaises(SystemError):
             writer.finish_with_size(-3)
 
     def test_write_bytes(self):
-         # Test PyBytesWriter_WriteBytes()
-         writer = self.create_writer()
-         writer.write_bytes(b'Hello World!', -1)
-         self.assertEqual(writer.finish(), self.result_type(b'Hello World!'))
+        # Test PyBytesWriter_WriteBytes()
+        writer = self.create_writer()
+        writer.write_bytes(b'Hello World!', -1)
+        self.bytes_equal(writer.finish(), b'Hello World!')
 
-         writer = self.create_writer()
-         writer.write_bytes(b'Hello ', -1)
-         writer.write_bytes(b'World! <truncated>', 6)
-         self.assertEqual(writer.finish(), self.result_type(b'Hello World!'))
+        writer = self.create_writer()
+        writer.write_bytes(b'Hello ', -1)
+        writer.write_bytes(b'World! <truncated>', 6)
+        self.bytes_equal(writer.finish(), b'Hello World!')
 
     def test_resize(self):
         # Test PyBytesWriter_Resize()
         writer = self.create_writer()
         writer.resize(len(b'number=123456'), b'number=123456')
         writer.resize(len(b'number=123456'), b'')
         self.assertEqual(writer.get_size(), len(b'number=123456'))
-        self.assertEqual(writer.finish(), self.result_type(b'number=123456'))
+        self.bytes_equal(writer.finish(), b'number=123456')
 
         writer = self.create_writer()
         writer.resize(0, b'')
         writer.resize(len(b'number=123456'), b'number=123456')
-        self.assertEqual(writer.finish(), self.result_type(b'number=123456'))
+        self.bytes_equal(writer.finish(), b'number=123456')
 
         writer = self.create_writer()
         writer.resize(len(b'number='), b'number=')
         writer.resize(len(b'number=123456'), b'123456')
-        self.assertEqual(writer.finish(), self.result_type(b'number=123456'))
+        self.bytes_equal(writer.finish(), b'number=123456')
 
         writer = self.create_writer()
         writer.resize(len(b'number='), b'number=')
         writer.resize(len(b'number='), b'')
         writer.resize(len(b'number=123456'), b'123456')
-        self.assertEqual(writer.finish(), self.result_type(b'number=123456'))
+        self.bytes_equal(writer.finish(), b'number=123456')
 
         writer = self.create_writer()
         writer.resize(len(b'number'), b'number')
         writer.resize(len(b'number='), b'=')
         writer.resize(len(b'number=123'), b'123')
         writer.resize(len(b'number=123456'), b'456')
-        self.assertEqual(writer.finish(), self.result_type(b'number=123456'))
+        self.bytes_equal(writer.finish(), b'number=123456')
 
     def test_format_i(self):
         # Test PyBytesWriter_Format()
         writer = self.create_writer()
         writer.format_i(b'x=%i', 123456)
-        self.assertEqual(writer.finish(), self.result_type(b'x=123456'))
+        self.bytes_equal(writer.finish(), b'x=123456')
 
         writer = self.create_writer()
         writer.format_i(b'x=%i, ', 123)
         writer.format_i(b'y=%i', 456)
-        self.assertEqual(writer.finish(), self.result_type(b'x=123, y=456'))
+        self.bytes_equal(writer.finish(), b'x=123, y=456')
+
+    def test_format_s(self):
+        # Test PyBytesWriter_Format()
+        writer = self.create_writer()
+        writer.format_s(b's=%s', b'Hello World')
+        self.bytes_equal(writer.finish(), b's=Hello World')
+
+    @support.nomemtest
+    def test_format_s_memory_error(self):
+        writer = self.create_writer()
+        s = b'x' * 500
+        with self.assertRaises(MemoryError):
+            try:
+                _testcapi.set_nomemory(0, 0)
+                writer.format_s(b's=%s', s)
+            finally:
+                _testcapi.remove_mem_hooks()
+        writer.write_bytes(b'after', -1)
+        self.bytes_equal(writer.finish(), b'after')
 
     def test_example_abc(self):
         self.assertEqual(_testcapi.byteswriter_abc(), b'abc')
@@ -405,7 +431,7 @@ def test_example_highlevel(self):
 
 
 class ByteArrayWriterTest(BytesWriterTest):
-    result_type = bytearray
+    RESULT_TYPE = bytearray
 
     def create_writer(self, alloc=0, string=b''):
         return _testcapi.PyBytesWriter(alloc, string, 1)

Misc/NEWS.d/next/C_API/2026-09-29-15-40-41.gh-issue-156939.Y5X3Wm.rst

@@ -0,0 +1,2 @@
+Fix :c:func:`PyBytesWriter_Format` error handling: restore the writer to its
+previous state on error. Patch by Victor Stinner.

Misc/NEWS.d/next/C_API/2026-09-29-21-37-08.gh-issue-156939.P_UD2P.rst

@@ -0,0 +1,3 @@
+Fix :c:func:`PyBytesWriter_FinishWithSize` to return an empty bytearray when
+:c:func:`!_PyBytesWriter_CreateByteArray` is used, instead of returning an
+empty bytes string. Patch by Victor Stinner.

Modules/_testcapi/bytes.c

@@ -183,6 +183,27 @@ writer_format_i(PyObject *self_raw, PyObject *args)
 }
 
 
+static PyObject*
+writer_format_s(PyObject *self_raw, PyObject *args)
+{
+    WriterObject *self = (WriterObject *)self_raw;
+    if (writer_check(self) < 0) {
+        return NULL;
+    }
+
+    char *format, *str;
+    if (!PyArg_ParseTuple(args, "yy", &format, &str)) {
+        return NULL;
+    }
+
+    if (PyBytesWriter_Format(self->writer, format, str) < 0) {
+        return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
+
+// PyBytesWriter_Resize
 static PyObject*
 writer_resize(PyObject *self_raw, PyObject *args)
 {
@@ -262,6 +283,7 @@ writer_finish_with_size(PyObject *self_raw, PyObject *args)
 static PyMethodDef writer_methods[] = {
     {"write_bytes", _PyCFunction_CAST(writer_write_bytes), METH_VARARGS},
     {"format_i", _PyCFunction_CAST(writer_format_i), METH_VARARGS},
+    {"format_s", _PyCFunction_CAST(writer_format_s), METH_VARARGS},
     {"resize", _PyCFunction_CAST(writer_resize), METH_VARARGS},
     {"get_size", _PyCFunction_CAST(writer_get_size), METH_NOARGS},
     {"finish", _PyCFunction_CAST(writer_finish), METH_NOARGS},

Objects/bytesobject.c

@@ -218,7 +218,7 @@ bytes_fromformat(PyBytesWriter *writer, Py_ssize_t writer_pos,
         size_t len = (len_expr); \
         s = PyBytesWriter_GrowAndUpdatePointer(writer, len, s); \
         if (s == NULL) { \
-            goto error; \
+            return NULL; \
         } \
         memcpy(s, (str), len); \
         s += len; \
@@ -272,7 +272,7 @@ bytes_fromformat(PyBytesWriter *writer, Py_ssize_t writer_pos,
                 PyErr_SetString(PyExc_OverflowError,
                                 "PyBytes_FromFormatV(): %c format "
                                 "expects an integer in range [0; 255]");
-                goto error;
+                return NULL;
             }
             *s++ = (unsigned char)c;
             break;
@@ -365,9 +365,6 @@ bytes_fromformat(PyBytesWriter *writer, Py_ssize_t writer_pos,
 #undef WRITE_BYTES_LEN
 
     return s;
-
- error:
-    return NULL;
 }
 
 
@@ -3746,7 +3743,7 @@ PyObject*
 PyBytesWriter_FinishWithSize(PyBytesWriter *writer, Py_ssize_t size)
 {
     PyObject *result;
-    if (size == 0) {
+    if (size == 0 && !writer->use_bytearray) {
         result = bytes_get_empty();
     }
     else if (writer->obj != NULL) {
@@ -3910,6 +3907,7 @@ int
 PyBytesWriter_Format(PyBytesWriter *writer, const char *format, ...)
 {
     Py_ssize_t pos = writer->size;
+    Py_ssize_t old_pos = pos;
     if (PyBytesWriter_Grow(writer, strlen(format)) < 0) {
         return -1;
     }
@@ -3919,6 +3917,12 @@ PyBytesWriter_Format(PyBytesWriter *writer, const char *format, ...)
     char *buf = bytes_fromformat(writer, pos, format, vargs);
     va_end(vargs);
 
+    if (buf == NULL) {
+        // On error, reset the writer to its previous state (undo any write)
+        writer->size = old_pos;
+        return -1;
+    }
+
     Py_ssize_t size = buf - byteswriter_data(writer);
     return PyBytesWriter_Resize(writer, size);
 }