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>

Lib/test/test_capi/test_bytes.py

@@ -356,18 +356,22 @@ class BaseWriterTest:
     def create_writer(self, alloc=0, string=b''):
         raise NotImplementedError
 
+    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(), b'')
+        self.bytes_equal(writer.finish(), b'')
 
         writer = self.create_writer(3)
         writer.write(0, b'abc')
         self.assertEqual(writer.get_size(), 3)
-        result = writer.finish()
-        self.assertEqual(result, b'abc')
-        self.assertEqual(type(result), self.RESULT_TYPE)
+        self.bytes_equal(writer.finish(), b'abc')
 
     @unittest.skipUnless(support.Py_DEBUG, 'need Py_DEBUG')
     def test_get_data(self):
@@ -408,9 +412,7 @@ class BaseWriterTest:
         writer = self.create_writer(10)
         writer.write(0, b'abc123')
         self.assertEqual(writer.get_size(), 10)
-        result = writer.finish_with_size(3)
-        self.assertEqual(result, b'abc')
-        self.assertEqual(type(result), self.RESULT_TYPE)
+        self.bytes_equal(writer.finish_with_size(3), b'abc')
 
         # Error if the size is negative
         writer = self.create_writer(3, )
@@ -428,19 +430,19 @@ class BaseWriterTest:
         # Test PyBytesWriter_WriteBytes()
         writer = self.create_writer()
         writer.write_bytes(b'Hello World!', -1)
-        self.assertEqual(writer.finish(), b'Hello World!')
+        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(), b'Hello World!')
+        self.bytes_equal(writer.finish(), b'Hello World!')
 
     def test_resize(self):
         # Test PyBytesWriter_Resize()
         writer = self.create_writer()
         writer.resize(len(b'hello'))
         writer.write(0, b'hello')
-        self.assertEqual(writer.finish(), b'hello')
+        self.bytes_equal(writer.finish(), b'hello')
 
         writer = self.create_writer()
         writer.resize(0)  # noop
@@ -451,7 +453,7 @@ class BaseWriterTest:
         writer.resize(len(b'number=123'), )
         writer.write(len(b'number='), b'123')
         writer.resize(len(b'number=123'))  # noop
-        self.assertEqual(writer.finish(), b'number=123')
+        self.bytes_equal(writer.finish(), b'number=123')
 
         # Switch from small buffer to large buffer
         writer = self.create_writer()
@@ -460,7 +462,7 @@ class BaseWriterTest:
         writer.write(0, b's' * small)
         writer.resize(large)
         writer.write(small, b'L' * (large - small))
-        self.assertEqual(writer.finish(),
+        self.bytes_equal(writer.finish(),
                          b's' * small + b'L' * (large - small))
 
         # Make sure that it's possible to write after a resize to zero
@@ -469,7 +471,7 @@ class BaseWriterTest:
         writer.resize(self.LARGE_BUFFER)
         writer.resize(0)
         writer.write_bytes(b'abc', 3)
-        self.assertEqual(writer.finish(), b'abc')
+        self.bytes_equal(writer.finish(), b'abc')
 
         # invalid size
         for size in (self.SMALL_BUFFER, self.LARGE_BUFFER):
@@ -480,7 +482,7 @@ class BaseWriterTest:
                     writer.resize(-1)
                 with self.assertRaises((MemoryError, OverflowError)):
                     writer.resize(_testcapi.PY_SSIZE_T_MAX)
-                self.assertEqual(writer.finish(), b'x' * size)
+                self.bytes_equal(writer.finish(), b'x' * size)
 
     @unittest.skipUnless(support.Py_DEBUG, 'need debug build')
     def test_resize_canary(self):
@@ -495,14 +497,14 @@ class BaseWriterTest:
                 self.assertEqual(get_data_canary(writer), data + CANARY_BYTE)
                 writer.resize(size - 1)
                 self.assertEqual(get_data_canary(writer), data[:-1] + CANARY_BYTE)
-                self.assertEqual(writer.finish(),  data[:-1])
+                self.bytes_equal(writer.finish(),  data[:-1])
 
                 # Make the buffer empty
                 writer = self.create_writer(size)
                 writer.write(0, data)
                 writer.resize(0)
                 self.assertEqual(get_data_canary(writer), CANARY_BYTE)
-                self.assertEqual(writer.finish(),  b'')
+                self.bytes_equal(writer.finish(),  b'')
 
     @support.nomemtest
     def test_resize_error(self):
@@ -516,7 +518,7 @@ class BaseWriterTest:
                 writer.resize(size)
         suffix = b'still working'
         writer.write_bytes(suffix, -1)
-        self.assertEqual(writer.finish(), init + suffix)
+        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
@@ -526,7 +528,7 @@ class BaseWriterTest:
         writer = self.create_writer(0)
         writer.grow(len(b'number=123'))
         writer.write(0, b'number=123')
-        self.assertEqual(writer.finish(), b'number=123')
+        self.bytes_equal(writer.finish(), b'number=123')
 
         writer = self.create_writer()
         writer.grow(0)  # noop
@@ -537,7 +539,7 @@ class BaseWriterTest:
         writer.grow(len(b'123'), )
         writer.write(len(b'number='), b'123')
         writer.grow(0)  # noop
-        self.assertEqual(writer.finish(), b'number=123')
+        self.bytes_equal(writer.finish(), b'number=123')
 
         # Switch from small buffer to large buffer
         writer = self.create_writer()
@@ -546,7 +548,7 @@ class BaseWriterTest:
         writer.write(0, b's' * small)
         writer.grow(large - small)
         writer.write(small, b'L' * (large - small))
-        self.assertEqual(writer.finish(),
+        self.bytes_equal(writer.finish(),
                          b's' * small + b'L' * (large - small))
 
         # invalid size
@@ -558,7 +560,7 @@ class BaseWriterTest:
                     writer.grow(-size - 1)
                 with self.assertRaises(MemoryError):
                     writer.grow(_testcapi.PY_SSIZE_T_MAX)
-                self.assertEqual(writer.finish(), b'x' * size)
+                self.bytes_equal(writer.finish(), b'x' * size)
 
     @unittest.skipUnless(support.Py_DEBUG, 'need debug build')
     def test_grow_canary(self):
@@ -572,14 +574,14 @@ class BaseWriterTest:
                 self.assertEqual(get_data_canary(writer), data + CANARY_BYTE)
                 writer.grow(-1)
                 self.assertEqual(get_data_canary(writer), data[:-1] + CANARY_BYTE)
-                self.assertEqual(writer.finish(),  data[:-1])
+                self.bytes_equal(writer.finish(),  data[:-1])
 
                 # Make the buffer empty
                 writer = self.create_writer(size)
                 writer.write(0, data)
                 writer.grow(-size)
                 self.assertEqual(writer.get_data(), b'')
-                self.assertEqual(writer.finish(),  b'')
+                self.bytes_equal(writer.finish(),  b'')
 
     @support.nomemtest
     def test_grow_error(self):
@@ -592,7 +594,7 @@ class BaseWriterTest:
                 writer.grow(100)
         suffix = b'still working'
         writer.write_bytes(suffix, -1)
-        self.assertEqual(writer.finish(), init + suffix)
+        self.bytes_equal(writer.finish(), init + suffix)
 
         # Note: PyBytesWriter_Grow() leaves the buffer unchanged (no resize)
         # if grow is negative.
@@ -601,12 +603,27 @@ class BaseWriterTest:
         # Test PyBytesWriter_Format()
         writer = self.create_writer()
         writer.format_i(b'x=%i', 123456)
-        self.assertEqual(writer.finish(), 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(), 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' * self.LARGE_BUFFER
+        with self.assertRaises(MemoryError):
+            _testcapi.call_with_nomemory(0, 0, writer.format_s, b's=%s', s)
+        writer.write_bytes(b'after', -1)
+        self.bytes_equal(writer.finish(), b'after')
 
     @unittest.skipUnless(support.Py_DEBUG, 'need a Python debug build')
     def test_canary_byte(self):
@@ -697,7 +714,7 @@ class BaseWriterTest:
 
         self.assertEqual(size, 10)
         self.assertEqual(data, b'x' * 10)
-        self.assertEqual(writer.finish(), b'x' * 10)
+        self.bytes_equal(writer.finish(), b'x' * 10)
 
 
 class BytesWriterTest(BaseWriterTest, unittest.TestCase):

Modules/_testcapi/bytes.c

@@ -209,6 +209,26 @@ 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)
@@ -333,6 +353,7 @@ static PyMethodDef writer_methods[] = {
     {"write", _PyCFunction_CAST(writer_write), METH_VARARGS | METH_KEYWORDS},
     {"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},
     {"grow", _PyCFunction_CAST(writer_grow), METH_VARARGS},
     {"get_data", _PyCFunction_CAST(writer_get_data), METH_VARARGS},

Objects/bytesobject.c

@@ -208,7 +208,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; \
@@ -262,7 +262,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;
@@ -355,9 +355,6 @@ bytes_fromformat(PyBytesWriter *writer, Py_ssize_t writer_pos,
 #undef WRITE_BYTES_LEN
 
     return s;
-
- error:
-    return NULL;
 }
 
 
@@ -3920,7 +3917,7 @@ 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) {
@@ -4030,6 +4027,7 @@ PyBytesWriter_Resize(PyBytesWriter *writer, Py_ssize_t new_size)
     else {
         // The buffer is already large enough. Never shrink the buffer.
     }
+
     writer->size = new_size;
 #ifdef Py_DEBUG
     byteswriter_write_canary_byte(writer);
@@ -4131,6 +4129,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;
     }
@@ -4140,6 +4139,15 @@ 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;
+#ifdef Py_DEBUG
+        byteswriter_write_canary_byte(writer);
+#endif
+        return -1;
+    }
+
     Py_ssize_t size = buf - byteswriter_data(writer);
     return PyBytesWriter_Resize(writer, size);
 }