pythongh-155742: Get singleton in PyBytesWriter_FinishWithSize() (python#155795)

If the result size of 1 byte, return the singleton rather than
creating a new bytes string.

Author: vstinner

Lib/test/test_capi/test_bytes.py

@@ -299,11 +299,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 test_create(self):
         # Test PyBytesWriter_Create()
@@ -388,10 +388,48 @@ def test_example_highlevel(self):
         self.assertEqual(_testcapi.byteswriter_highlevel(), b'Hello World!')
 
 
-class ByteArrayWriterTest(BytesWriterTest):
+class BytesWriterTest(BaseWriterTest, unittest.TestCase):
+    result_type = bytes
+
+    def create_writer(self, alloc=0, string=b''):
+        # Test PyBytesWriter_Create()
+        return _testcapi.PyBytesWriter(alloc, string, 0)
+
+    # Only PyBytesWriter_Create() returns singletons
+    def test_singletons(self):
+        empty = b''
+        singletons = {ch: bytes((ch,)) for ch in range(256)}
+        small_buffer = _testcapi.PyBytesWriter_small_buffer
+
+        writer = self.create_writer()
+        self.assertIs(writer.finish(), empty)
+
+        # Test writer larger than small_buffer
+        writer = self.create_writer()
+        unused_text = b'x' * (small_buffer * 2)
+        writer.write_bytes(unused_text, len(unused_text))
+        self.assertIs(writer.finish_with_size(0), empty)
+
+        for ch in range(256):
+            text = bytes((ch,))
+
+            writer = self.create_writer()
+            writer.write_bytes(text, 1)
+            self.assertIs(writer.finish(), singletons[ch])
+
+            # Test writer larger than small_buffer
+            writer = self.create_writer()
+            writer.write_bytes(text, 1)
+            unused_text = b'x' * (small_buffer * 2)
+            writer.write_bytes(unused_text, len(unused_text))
+            self.assertIs(writer.finish_with_size(1), singletons[ch])
+
+
+class ByteArrayWriterTest(BaseWriterTest, unittest.TestCase):
     result_type = bytearray
 
     def create_writer(self, alloc=0, string=b''):
+        # Test private _PyBytesWriter_CreateByteArray()
         return _testcapi.PyBytesWriter(alloc, string, 1)
 
 

Misc/NEWS.d/next/C_API/2026-08-14-14-47-08.gh-issue-155742.UGI3Pn.rst

@@ -0,0 +1,2 @@
+:c:func:`PyBytesWriter_FinishWithSize` now returns single byte singletons if
+*size* equals to ``1``. 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()
 
 
@@ -150,8 +152,8 @@ writer_write_bytes(PyObject *self_raw, PyObject *args)
     }
 
     char *bytes;
-    Py_ssize_t size;
-    if (!PyArg_ParseTuple(args, "yn", &bytes, &size)) {
+    Py_ssize_t unused_size, size;
+    if (!PyArg_ParseTuple(args, "y#n", &bytes, &unused_size, &size)) {
         return NULL;
     }
 
@@ -377,5 +379,12 @@ _PyTestCapi_Init_Bytes(PyObject *m)
     }
     Py_DECREF(writer_type);
 
+    // PyBytesWriter.obj is the second member, small_buffer is the first member
+    long size = (long)offsetof(PyBytesWriter, obj);
+    if (PyModule_AddIntConstant(m, "PyBytesWriter_small_buffer", size) < 0) {
+        Py_DECREF(writer_type);
+        return -1;
+    }
+
     return 0;
 }

Objects/bytesobject.c

@@ -3766,13 +3766,23 @@ PyBytesWriter_FinishWithSize(PyBytesWriter *writer, Py_ssize_t size)
                 }
             }
         }
+
         result = writer->obj;
         writer->obj = NULL;
+
+        if (size == 1 && !writer->use_bytearray) {
+            // Get the single byte singleton
+            unsigned char ch = PyBytes_AS_STRING(result)[0];
+            PyObject *op = (PyObject*)CHARACTER(ch);
+            assert(_Py_IsImmortal(op));
+            Py_SETREF(result, op);
+        }
     }
     else if (writer->use_bytearray) {
         result = PyByteArray_FromStringAndSize(writer->small_buffer, size);
     }
     else {
+        // The function returns single byte singleton if size equals 1
         result = PyBytes_FromStringAndSize(writer->small_buffer, size);
     }
     PyBytesWriter_Discard(writer);