gh-129813: Add PyBytesWriter C API (version 2)

Add functions:

* PyBytesWriter_Create()
* PyBytesWriter_Discard()
* PyBytesWriter_Finish()
* PyBytesWriter_Alloc()
* PyBytesWriter_Extend()
* PyBytesWriter_Truncate()
* PyBytesWriter_WriteBytes()

Author: vstinner

Include/cpython/bytesobject.h

@@ -40,3 +40,31 @@ _PyBytes_Join(PyObject *sep, PyObject *iterable)
 {
     return PyBytes_Join(sep, iterable);
 }
+
+
+// --- PyBytesWriter API -----------------------------------------------------
+
+typedef struct PyBytesWriter PyBytesWriter;
+
+PyAPI_FUNC(PyBytesWriter *) PyBytesWriter_Create(
+    Py_ssize_t prealloc);
+PyAPI_FUNC(void) PyBytesWriter_Discard(
+    PyBytesWriter *writer);
+PyAPI_FUNC(PyObject*) PyBytesWriter_Finish(
+    PyBytesWriter *writer);
+
+PyAPI_FUNC(void*) PyBytesWriter_Alloc(
+    PyBytesWriter *writer,
+    Py_ssize_t alloc);
+PyAPI_FUNC(void*) PyBytesWriter_Extend(
+    PyBytesWriter *writer,
+    void *buf,
+    Py_ssize_t extend);
+PyAPI_FUNC(int) PyBytesWriter_Truncate(
+    PyBytesWriter *writer,
+    void *buf);
+
+PyAPI_FUNC(int) PyBytesWriter_WriteBytes(
+    PyBytesWriter *writer,
+    const void *bytes,
+    Py_ssize_t size);

Include/internal/pycore_freelist_state.h

@@ -24,6 +24,7 @@ extern "C" {
 #  define Py_futureiters_MAXFREELIST 255
 #  define Py_object_stack_chunks_MAXFREELIST 4
 #  define Py_unicode_writers_MAXFREELIST 1
+#  define Py_bytes_writers_MAXFREELIST 1
 #  define Py_pymethodobjects_MAXFREELIST 20
 
 // A generic freelist of either PyObjects or other data structures.
@@ -53,6 +54,7 @@ struct _Py_freelists {
     struct _Py_freelist futureiters;
     struct _Py_freelist object_stack_chunks;
     struct _Py_freelist unicode_writers;
+    struct _Py_freelist bytes_writers;
     struct _Py_freelist pymethodobjects;
 };
 

Lib/test/test_capi/test_bytes.py

@@ -291,5 +291,70 @@ def test_join(self):
             bytes_join(b'', NULL)
 
 
+class PyBytesWriterTest(unittest.TestCase):
+    def create_writer(self, prealloc):
+        return _testcapi.PyBytesWriter(prealloc)
+
+    def test_empty(self):
+        # Test PyBytesWriter_Create()
+        writer = self.create_writer(0)
+        self.assertEqual(writer.finish(), b'')
+
+    def test_write_bytes(self):
+        # Test PyBytesWriter_WriteBytes()
+
+        writer = self.create_writer(0)
+        writer.write_bytes(b'Hello World!', -1)
+        self.assertEqual(writer.finish(), b'Hello World!')
+
+        writer = self.create_writer(0)
+        writer.write_bytes(b'Hello ', -1)
+        writer.write_bytes(b'World! <truncated>', 6)
+        self.assertEqual(writer.finish(), b'Hello World!')
+
+    def test_extend(self):
+        # Test PyBytesWriter_Extend() and PyBytesWriter_SetSizeFromBuf()
+
+        writer = self.create_writer(0)
+        writer.extend(13, b'number=123456')
+        writer.extend(0, b'')
+        self.assertEqual(writer.finish(), b'number=123456')
+
+        writer = self.create_writer(0)
+        writer.extend(0, b'')
+        writer.extend(13, b'number=123456')
+        self.assertEqual(writer.finish(), b'number=123456')
+
+        writer = self.create_writer(0)
+        writer.extend(10, b'number=')
+        writer.extend(10, b'123456')
+        self.assertEqual(writer.finish(), b'number=123456')
+
+        writer = self.create_writer(0)
+        writer.extend(7, b'number=')
+        writer.extend(0, b'')
+        writer.extend(6, b'123456')
+        self.assertEqual(writer.finish(), b'number=123456')
+
+        writer = self.create_writer(0)
+        writer.extend(6, b'number')
+        writer.extend(1, b'=')
+        writer.extend(3, b'123')
+        writer.extend(3, b'456')
+        self.assertEqual(writer.finish(), b'number=123456')
+
+    def test_hello_world(self):
+        self.assertEqual(_testcapi.byteswriter_hello_world(),
+                         b'Hello World')
+
+    def test_alloc(self):
+        self.assertEqual(_testcapi.byteswriter_alloc(),
+                         b'abc')
+
+    def test_extend(self):
+        self.assertEqual(_testcapi.byteswriter_extend(),
+                         b'Hello World')
+
+
 if __name__ == "__main__":
     unittest.main()

Modules/_pickle.c

@@ -2615,29 +2615,25 @@ save_picklebuffer(PickleState *st, PicklerObject *self, PyObject *obj)
 static PyObject *
 raw_unicode_escape(PyObject *obj)
 {
-    char *p;
-    Py_ssize_t i, size;
-    const void *data;
-    int kind;
-    _PyBytesWriter writer;
+    Py_ssize_t size = PyUnicode_GET_LENGTH(obj);
+    const void *data = PyUnicode_DATA(obj);
+    int kind = PyUnicode_KIND(obj);
 
-    _PyBytesWriter_Init(&writer);
-
-    size = PyUnicode_GET_LENGTH(obj);
-    data = PyUnicode_DATA(obj);
-    kind = PyUnicode_KIND(obj);
-
-    p = _PyBytesWriter_Alloc(&writer, size);
-    if (p == NULL)
+    PyBytesWriter *writer = PyBytesWriter_Create(size);
+    if (writer == NULL) {
+        return NULL;
+    }
+    char *p = PyBytesWriter_Alloc(writer, size);
+    if (p == NULL) {
         goto error;
-    writer.overallocate = 1;
+    }
 
-    for (i=0; i < size; i++) {
+    for (Py_ssize_t i=0; i < size; i++) {
         Py_UCS4 ch = PyUnicode_READ(kind, data, i);
         /* Map 32-bit characters to '\Uxxxxxxxx' */
         if (ch >= 0x10000) {
             /* -1: subtract 1 preallocated byte */
-            p = _PyBytesWriter_Prepare(&writer, p, 10-1);
+            p = PyBytesWriter_Extend(writer, p, 10-1);
             if (p == NULL)
                 goto error;
 
@@ -2658,7 +2654,7 @@ raw_unicode_escape(PyObject *obj)
                  ch == 0x1a)
         {
             /* -1: subtract 1 preallocated byte */
-            p = _PyBytesWriter_Prepare(&writer, p, 6-1);
+            p = PyBytesWriter_Extend(writer, p, 6-1);
             if (p == NULL)
                 goto error;
 
@@ -2674,10 +2670,13 @@ raw_unicode_escape(PyObject *obj)
             *p++ = (char) ch;
     }
 
-    return _PyBytesWriter_Finish(&writer, p);
+    if (PyBytesWriter_Truncate(writer, p) < 0) {
+        goto error;
+    }
+    return PyBytesWriter_Finish(writer);
 
 error:
-    _PyBytesWriter_Dealloc(&writer);
+    PyBytesWriter_Discard(writer);
     return NULL;
 }
 

Modules/_struct.c

@@ -2272,7 +2272,6 @@ strings.");
 static PyObject *
 s_pack(PyObject *self, PyObject *const *args, Py_ssize_t nargs)
 {
-    char *buf;
     PyStructObject *soself;
     _structmodulestate *state = get_struct_state_structinst(self);
 
@@ -2288,21 +2287,23 @@ s_pack(PyObject *self, PyObject *const *args, Py_ssize_t nargs)
     }
 
     /* Allocate a new string */
-    _PyBytesWriter writer;
-    _PyBytesWriter_Init(&writer);
-    buf = _PyBytesWriter_Alloc(&writer, soself->s_size);
+    PyBytesWriter *writer = PyBytesWriter_Create(soself->s_size);
+    if (writer == NULL) {
+        return NULL;
+    }
+    char *buf = PyBytesWriter_Alloc(writer, soself->s_size);
     if (buf == NULL) {
-        _PyBytesWriter_Dealloc(&writer);
+        PyBytesWriter_Discard(writer);
         return NULL;
     }
 
     /* Call the guts */
     if ( s_pack_internal(soself, args, 0, buf, state) != 0 ) {
-        _PyBytesWriter_Dealloc(&writer);
+        PyBytesWriter_Discard(writer);
         return NULL;
     }
 
-    return _PyBytesWriter_Finish(&writer, buf + soself->s_size);
+    return PyBytesWriter_Finish(writer);
 }
 
 PyDoc_STRVAR(s_pack_into__doc__,