gh-129813: Add PyBytesWriter C API (with size flavor)

Add functions:

* PyBytesWriter_Create()
* PyBytesWriter_Discard()
* PyBytesWriter_Finish()
* PyBytesWriter_FinishWithSize()
* PyBytesWriter_FinishWithEndPointer()
* PyBytesWriter_Data()
* PyBytesWriter_Allocated()
* PyBytesWriter_SetSize()
* PyBytesWriter_Resize()

Author: vstinner

Include/cpython/bytesobject.h

@@ -40,3 +40,33 @@ _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 alloc);
+PyAPI_FUNC(void) PyBytesWriter_Discard(
+    PyBytesWriter *writer);
+PyAPI_FUNC(PyObject*) PyBytesWriter_Finish(
+    PyBytesWriter *writer);
+PyAPI_FUNC(PyObject*) PyBytesWriter_FinishWithSize(
+    PyBytesWriter *writer,
+    Py_ssize_t size);
+PyAPI_FUNC(PyObject*) PyBytesWriter_FinishWithEndPointer(
+    PyBytesWriter *writer,
+    void *data);
+
+PyAPI_FUNC(void*) PyBytesWriter_Data(
+    PyBytesWriter *writer);
+PyAPI_FUNC(Py_ssize_t) PyBytesWriter_Allocated(
+    PyBytesWriter *writer);
+
+PyAPI_FUNC(int) PyBytesWriter_SetSize(
+    PyBytesWriter *writer,
+    Py_ssize_t size);
+PyAPI_FUNC(int) PyBytesWriter_Resize(
+    PyBytesWriter *writer,
+    Py_ssize_t alloc);

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;
 };
 

Modules/_pickle.c

@@ -2615,31 +2615,31 @@ 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)
-        goto error;
-    writer.overallocate = 1;
+    Py_ssize_t alloc = size;
+    PyBytesWriter *writer = PyBytesWriter_Create(alloc);
+    if (writer == NULL) {
+        return NULL;
+    }
+    char *p = PyBytesWriter_Data(writer);
 
-    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);
-            if (p == NULL)
+            alloc += 10-1;
+            Py_ssize_t pos = p - (char*)PyBytesWriter_Data(writer);
+            if (PyBytesWriter_SetSize(writer, pos) < 0) {
                 goto error;
+            }
+            if (PyBytesWriter_Resize(writer, alloc) < 0) {
+                goto error;
+            }
+            p = (char*)PyBytesWriter_Data(writer) + pos;
 
             *p++ = '\\';
             *p++ = 'U';
@@ -2658,9 +2658,15 @@ raw_unicode_escape(PyObject *obj)
                  ch == 0x1a)
         {
             /* -1: subtract 1 preallocated byte */
-            p = _PyBytesWriter_Prepare(&writer, p, 6-1);
-            if (p == NULL)
+            alloc += 6-1;
+            Py_ssize_t pos = p - (char*)PyBytesWriter_Data(writer);
+            if (PyBytesWriter_SetSize(writer, pos) < 0) {
                 goto error;
+            }
+            if (PyBytesWriter_Resize(writer, alloc) < 0) {
+                goto error;
+            }
+            p = (char*)PyBytesWriter_Data(writer) + pos;
 
             *p++ = '\\';
             *p++ = 'u';
@@ -2674,10 +2680,10 @@ raw_unicode_escape(PyObject *obj)
             *p++ = (char) ch;
     }
 
-    return _PyBytesWriter_Finish(&writer, p);
+    return PyBytesWriter_FinishWithEndPointer(writer, p);
 
 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,19 @@ 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);
-    if (buf == NULL) {
-        _PyBytesWriter_Dealloc(&writer);
+    PyBytesWriter *writer = PyBytesWriter_Create(soself->s_size);
+    if (writer == NULL) {
         return NULL;
     }
+    char *buf = PyBytesWriter_Data(writer);
 
     /* 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_FinishWithSize(writer, soself->s_size);
 }
 
 PyDoc_STRVAR(s_pack_into__doc__,

Modules/binascii.c

@@ -302,16 +302,13 @@ static PyObject *
 binascii_b2a_uu_impl(PyObject *module, Py_buffer *data, int backtick)
 /*[clinic end generated code: output=b1b99de62d9bbeb8 input=beb27822241095cd]*/
 {
-    unsigned char *ascii_data;
     const unsigned char *bin_data;
     int leftbits = 0;
     unsigned char this_ch;
     unsigned int leftchar = 0;
     binascii_state *state;
-    Py_ssize_t bin_len, out_len;
-    _PyBytesWriter writer;
+    Py_ssize_t bin_len;
 
-    _PyBytesWriter_Init(&writer);
     bin_data = data->buf;
     bin_len = data->len;
     if ( bin_len > 45 ) {
@@ -325,10 +322,12 @@ binascii_b2a_uu_impl(PyObject *module, Py_buffer *data, int backtick)
     }
 
     /* We're lazy and allocate to much (fixed up later) */
-    out_len = 2 + (bin_len + 2) / 3 * 4;
-    ascii_data = _PyBytesWriter_Alloc(&writer, out_len);
-    if (ascii_data == NULL)
+    Py_ssize_t out_len = 2 + (bin_len + 2) / 3 * 4;
+    PyBytesWriter *writer = PyBytesWriter_Create(out_len);
+    if (writer == NULL) {
         return NULL;
+    }
+    unsigned char *ascii_data = PyBytesWriter_Data(writer);
 
     /* Store the length */
     if (backtick && !bin_len)
@@ -356,7 +355,7 @@ binascii_b2a_uu_impl(PyObject *module, Py_buffer *data, int backtick)
     }
     *ascii_data++ = '\n';       /* Append a courtesy newline */
 
-    return _PyBytesWriter_Finish(&writer, ascii_data);
+    return PyBytesWriter_FinishWithEndPointer(writer, ascii_data);
 }
 
 /*[clinic input]
@@ -387,12 +386,11 @@ binascii_a2b_base64_impl(PyObject *module, Py_buffer *data, int strict_mode)
 
     /* Allocate the buffer */
     Py_ssize_t bin_len = ((ascii_len+3)/4)*3; /* Upper bound, corrected later */
-    _PyBytesWriter writer;
-    _PyBytesWriter_Init(&writer);
-    unsigned char *bin_data = _PyBytesWriter_Alloc(&writer, bin_len);
-    if (bin_data == NULL)
+    PyBytesWriter *writer = PyBytesWriter_Create(bin_len);
+    if (writer == NULL) {
         return NULL;
-    unsigned char *bin_data_start = bin_data;
+    }
+    unsigned char *bin_data = PyBytesWriter_Data(writer);
 
     if (strict_mode && ascii_len > 0 && ascii_data[0] == '=') {
         state = get_binascii_state(module);
@@ -488,12 +486,14 @@ binascii_a2b_base64_impl(PyObject *module, Py_buffer *data, int strict_mode)
         state = get_binascii_state(module);
         if (state == NULL) {
             /* error already set, from get_binascii_state */
+            assert(PyErr_Occurred());
         } else if (quad_pos == 1) {
             /*
             ** There is exactly one extra valid, non-padding, base64 character.
             ** This is an invalid length, as there is no possible input that
             ** could encoded into such a base64 string.
             */
+            unsigned char *bin_data_start = PyBytesWriter_Data(writer);
             PyErr_Format(state->Error,
                          "Invalid base64-encoded string: "
                          "number of data characters (%zd) cannot be 1 more "
@@ -502,13 +502,15 @@ binascii_a2b_base64_impl(PyObject *module, Py_buffer *data, int strict_mode)
         } else {
             PyErr_SetString(state->Error, "Incorrect padding");
         }
-        error_end:
-        _PyBytesWriter_Dealloc(&writer);
-        return NULL;
+        goto error_end;
     }
 
 done:
-    return _PyBytesWriter_Finish(&writer, bin_data);
+    return PyBytesWriter_FinishWithEndPointer(writer, bin_data);
+
+error_end:
+    PyBytesWriter_Discard(writer);
+    return NULL;
 }
 
 
@@ -527,18 +529,15 @@ static PyObject *
 binascii_b2a_base64_impl(PyObject *module, Py_buffer *data, int newline)
 /*[clinic end generated code: output=4ad62c8e8485d3b3 input=0e20ff59c5f2e3e1]*/
 {
-    unsigned char *ascii_data;
     const unsigned char *bin_data;
     int leftbits = 0;
     unsigned char this_ch;
     unsigned int leftchar = 0;
-    Py_ssize_t bin_len, out_len;
-    _PyBytesWriter writer;
+    Py_ssize_t bin_len;
     binascii_state *state;
 
     bin_data = data->buf;
     bin_len = data->len;
-    _PyBytesWriter_Init(&writer);
 
     assert(bin_len >= 0);
 
@@ -554,12 +553,15 @@ binascii_b2a_base64_impl(PyObject *module, Py_buffer *data, int newline)
     /* We're lazy and allocate too much (fixed up later).
        "+2" leaves room for up to two pad characters.
        Note that 'b' gets encoded as 'Yg==\n' (1 in, 5 out). */
-    out_len = bin_len*2 + 2;
-    if (newline)
+    Py_ssize_t out_len = bin_len*2 + 2;
+    if (newline) {
         out_len++;
-    ascii_data = _PyBytesWriter_Alloc(&writer, out_len);
-    if (ascii_data == NULL)
+    }
+    PyBytesWriter *writer = PyBytesWriter_Create(out_len);
+    if (writer == NULL) {
         return NULL;
+    }
+    unsigned char *ascii_data = PyBytesWriter_Data(writer);
 
     for( ; bin_len > 0 ; bin_len--, bin_data++ ) {
         /* Shift the data into our buffer */
@@ -584,7 +586,7 @@ binascii_b2a_base64_impl(PyObject *module, Py_buffer *data, int newline)
     if (newline)
         *ascii_data++ = '\n';       /* Append a courtesy newline */
 
-    return _PyBytesWriter_Finish(&writer, ascii_data);
+    return PyBytesWriter_FinishWithEndPointer(writer, ascii_data);
 }