gh-129813, PEP 782: Use PyBytesWriter in _PyBytes_FormatEx() (#138839)

Replace the private _PyBytesWriter API with the new public
PyBytesWriter API.

Author: vstinner

Include/internal/pycore_long.h

@@ -135,7 +135,7 @@ extern int _PyLong_FormatWriter(
     int alternate);
 
 extern char* _PyLong_FormatBytesWriter(
-    _PyBytesWriter *writer,
+    PyBytesWriter *writer,
     char *str,
     PyObject *obj,
     int base,

Objects/bytesobject.c

@@ -432,7 +432,7 @@ getnextarg(PyObject *args, Py_ssize_t arglen, Py_ssize_t *p_argidx)
 
 static char*
 formatfloat(PyObject *v, int flags, int prec, int type,
-            PyObject **p_result, _PyBytesWriter *writer, char *str)
+            PyObject **p_result, PyBytesWriter *writer, char *str)
 {
     char *p;
     PyObject *result;
@@ -460,7 +460,7 @@ formatfloat(PyObject *v, int flags, int prec, int type,
 
     len = strlen(p);
     if (writer != NULL) {
-        str = _PyBytesWriter_Prepare(writer, str, len);
+        str = PyBytesWriter_GrowAndUpdatePointer(writer, len, str);
         if (str == NULL) {
             PyMem_Free(p);
             return NULL;
@@ -611,12 +611,10 @@ _PyBytes_FormatEx(const char *format, Py_ssize_t format_len,
                   PyObject *args, int use_bytearray)
 {
     const char *fmt;
-    char *res;
     Py_ssize_t arglen, argidx;
     Py_ssize_t fmtcnt;
     int args_owned = 0;
     PyObject *dict = NULL;
-    _PyBytesWriter writer;
 
     if (args == NULL) {
         PyErr_BadInternalCall();
@@ -625,14 +623,17 @@ _PyBytes_FormatEx(const char *format, Py_ssize_t format_len,
     fmt = format;
     fmtcnt = format_len;
 
-    _PyBytesWriter_Init(&writer);
-    writer.use_bytearray = use_bytearray;
-
-    res = _PyBytesWriter_Alloc(&writer, fmtcnt);
-    if (res == NULL)
+    PyBytesWriter *writer;
+    if (use_bytearray) {
+        writer = _PyBytesWriter_CreateByteArray(fmtcnt);
+    }
+    else {
+        writer = PyBytesWriter_Create(fmtcnt);
+    }
+    if (writer == NULL) {
         return NULL;
-    if (!use_bytearray)
-        writer.overallocate = 1;
+    }
+    char *res = PyBytesWriter_GetData(writer);
 
     if (PyTuple_Check(args)) {
         arglen = PyTuple_GET_SIZE(args);
@@ -835,11 +836,6 @@ _PyBytes_FormatEx(const char *format, Py_ssize_t format_len,
             if (v == NULL)
                 goto error;
 
-            if (fmtcnt == 0) {
-                /* last write: disable writer overallocation */
-                writer.overallocate = 0;
-            }
-
             sign = 0;
             fill = ' ';
             switch (c) {
@@ -900,8 +896,7 @@ _PyBytes_FormatEx(const char *format, Py_ssize_t format_len,
                     }
 
                     /* Fast path */
-                    writer.min_size -= 2; /* size preallocated for "%d" */
-                    res = _PyLong_FormatBytesWriter(&writer, res,
+                    res = _PyLong_FormatBytesWriter(writer, res,
                                                     v, base, alternate);
                     if (res == NULL)
                         goto error;
@@ -929,8 +924,7 @@ _PyBytes_FormatEx(const char *format, Py_ssize_t format_len,
                     && !(flags & (F_SIGN | F_BLANK)))
                 {
                     /* Fast path */
-                    writer.min_size -= 2; /* size preallocated for "%f" */
-                    res = formatfloat(v, flags, prec, c, NULL, &writer, res);
+                    res = formatfloat(v, flags, prec, c, NULL, writer, res);
                     if (res == NULL)
                         goto error;
                     continue;
@@ -986,9 +980,10 @@ _PyBytes_FormatEx(const char *format, Py_ssize_t format_len,
                 alloc++;
             /* 2: size preallocated for %s */
             if (alloc > 2) {
-                res = _PyBytesWriter_Prepare(&writer, res, alloc - 2);
-                if (res == NULL)
+                res = PyBytesWriter_GrowAndUpdatePointer(writer, alloc - 2, res);
+                if (res == NULL) {
                     goto error;
+                }
             }
 #ifndef NDEBUG
             char *before = res;
@@ -1061,10 +1056,6 @@ _PyBytes_FormatEx(const char *format, Py_ssize_t format_len,
             assert((res - before) == alloc);
 #endif
         } /* '%' */
-
-        /* If overallocation was disabled, ensure that it was the last
-           write. Otherwise, we missed an optimization */
-        assert(writer.overallocate || fmtcnt == 0 || use_bytearray);
     } /* until end */
 
     if (argidx < arglen && !dict) {
@@ -1076,10 +1067,10 @@ _PyBytes_FormatEx(const char *format, Py_ssize_t format_len,
     if (args_owned) {
         Py_DECREF(args);
     }
-    return _PyBytesWriter_Finish(&writer, res);
+    return PyBytesWriter_FinishWithPointer(writer, res);
 
  error:
-    _PyBytesWriter_Dealloc(&writer);
+    PyBytesWriter_Discard(writer);
     if (args_owned) {
         Py_DECREF(args);
     }

Objects/longobject.c

@@ -2020,7 +2020,7 @@ static int
 pylong_int_to_decimal_string(PyObject *aa,
                              PyObject **p_output,
                              _PyUnicodeWriter *writer,
-                             _PyBytesWriter *bytes_writer,
+                             PyBytesWriter *bytes_writer,
                              char **bytes_str)
 {
     PyObject *s = NULL;
@@ -2051,7 +2051,8 @@ pylong_int_to_decimal_string(PyObject *aa,
         Py_ssize_t size = PyUnicode_GET_LENGTH(s);
         const void *data = PyUnicode_DATA(s);
         int kind = PyUnicode_KIND(s);
-        *bytes_str = _PyBytesWriter_Prepare(bytes_writer, *bytes_str, size);
+        *bytes_str = PyBytesWriter_GrowAndUpdatePointer(bytes_writer, size,
+                                                        *bytes_str);
         if (*bytes_str == NULL) {
             goto error;
         }
@@ -2088,7 +2089,7 @@ static int
 long_to_decimal_string_internal(PyObject *aa,
                                 PyObject **p_output,
                                 _PyUnicodeWriter *writer,
-                                _PyBytesWriter *bytes_writer,
+                                PyBytesWriter *bytes_writer,
                                 char **bytes_str)
 {
     PyLongObject *scratch, *a;
@@ -2214,7 +2215,8 @@ long_to_decimal_string_internal(PyObject *aa,
         }
     }
     else if (bytes_writer) {
-        *bytes_str = _PyBytesWriter_Prepare(bytes_writer, *bytes_str, strlen);
+        *bytes_str = PyBytesWriter_GrowAndUpdatePointer(bytes_writer, strlen,
+                                                        *bytes_str);
         if (*bytes_str == NULL) {
             Py_DECREF(scratch);
             return -1;
@@ -2324,7 +2326,7 @@ long_to_decimal_string(PyObject *aa)
 static int
 long_format_binary(PyObject *aa, int base, int alternate,
                    PyObject **p_output, _PyUnicodeWriter *writer,
-                   _PyBytesWriter *bytes_writer, char **bytes_str)
+                   PyBytesWriter *bytes_writer, char **bytes_str)
 {
     PyLongObject *a = (PyLongObject *)aa;
     PyObject *v = NULL;
@@ -2385,7 +2387,8 @@ long_format_binary(PyObject *aa, int base, int alternate,
             return -1;
     }
     else if (bytes_writer) {
-        *bytes_str = _PyBytesWriter_Prepare(bytes_writer, *bytes_str, sz);
+        *bytes_str = PyBytesWriter_GrowAndUpdatePointer(bytes_writer, sz,
+                                                        *bytes_str);
         if (*bytes_str == NULL)
             return -1;
     }
@@ -2514,7 +2517,7 @@ _PyLong_FormatWriter(_PyUnicodeWriter *writer,
 }
 
 char*
-_PyLong_FormatBytesWriter(_PyBytesWriter *writer, char *str,
+_PyLong_FormatBytesWriter(PyBytesWriter *writer, char *str,
                           PyObject *obj,
                           int base, int alternate)
 {
@@ -6403,8 +6406,6 @@ int_to_bytes_impl(PyObject *self, Py_ssize_t length, PyObject *byteorder,
 /*[clinic end generated code: output=89c801df114050a3 input=66f9d0c20529b44f]*/
 {
     int little_endian;
-    PyObject *bytes;
-
     if (byteorder == NULL)
         little_endian = 0;
     else if (_PyUnicode_Equal(byteorder, &_Py_ID(little)))
@@ -6417,18 +6418,19 @@ int_to_bytes_impl(PyObject *self, Py_ssize_t length, PyObject *byteorder,
         return NULL;
     }
 
-    bytes = PyBytes_FromStringAndSize(NULL, length);
-    if (bytes == NULL)
+    PyBytesWriter *writer = PyBytesWriter_Create(length);
+    if (writer == NULL) {
         return NULL;
+    }
 
     if (_PyLong_AsByteArray((PyLongObject *)self,
-                            (unsigned char *)PyBytes_AS_STRING(bytes),
+                            PyBytesWriter_GetData(writer),
                             length, little_endian, is_signed, 1) < 0) {
-        Py_DECREF(bytes);
+        PyBytesWriter_Discard(writer);
         return NULL;
     }
 
-    return bytes;
+    return PyBytesWriter_Finish(writer);
 }
 
 /*[clinic input]