Convert _PyBytes_FormatEx()

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

@@ -431,7 +431,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;
@@ -459,7 +459,8 @@ formatfloat(PyObject *v, int flags, int prec, int type,
 
     len = strlen(p);
     if (writer != NULL) {
-        str = _PyBytesWriter_Prepare(writer, str, len);
+        Py_ssize_t resize = PyBytesWriter_GetSize(writer) + len;
+        str = PyBytesWriter_ResizeAndUpdatePointer(writer, resize, str);
         if (str == NULL) {
             PyMem_Free(p);
             return NULL;
@@ -612,12 +613,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();
@@ -626,14 +625,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);
@@ -836,11 +838,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) {
@@ -901,8 +898,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;
@@ -930,8 +926,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;
@@ -987,9 +982,11 @@ _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)
+                Py_ssize_t resize = PyBytesWriter_GetSize(writer) + alloc - 2;
+                res = PyBytesWriter_ResizeAndUpdatePointer(writer, resize, res);
+                if (res == NULL) {
                     goto error;
+                }
             }
 #ifndef NDEBUG
             char *before = res;
@@ -1062,10 +1059,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) {
@@ -1077,10 +1070,10 @@ _PyBytes_FormatEx(const char *format, Py_ssize_t format_len,
     if (args_owned) {
         Py_DECREF(args);
     }
-    return _PyBytesWriter_Finish(&writer, res);
+    return PyBytesWriter_FinishWithEndPointer(writer, res);
 
  error:
-    _PyBytesWriter_Dealloc(&writer);
+    PyBytesWriter_Discard(writer);
     if (args_owned) {
         Py_DECREF(args);
     }

Objects/longobject.c

@@ -2016,7 +2016,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;
@@ -2047,7 +2047,9 @@ 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);
+        Py_ssize_t resize = PyBytesWriter_GetSize(bytes_writer) + size;
+        *bytes_str = PyBytesWriter_ResizeAndUpdatePointer(bytes_writer, resize,
+                                                          *bytes_str);
         if (*bytes_str == NULL) {
             goto error;
         }
@@ -2084,7 +2086,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;
@@ -2210,7 +2212,9 @@ long_to_decimal_string_internal(PyObject *aa,
         }
     }
     else if (bytes_writer) {
-        *bytes_str = _PyBytesWriter_Prepare(bytes_writer, *bytes_str, strlen);
+        Py_ssize_t resize = PyBytesWriter_GetSize(bytes_writer) + strlen;
+        *bytes_str = PyBytesWriter_ResizeAndUpdatePointer(bytes_writer, resize,
+                                                          *bytes_str);
         if (*bytes_str == NULL) {
             Py_DECREF(scratch);
             return -1;
@@ -2320,7 +2324,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;
@@ -2381,7 +2385,9 @@ long_format_binary(PyObject *aa, int base, int alternate,
             return -1;
     }
     else if (bytes_writer) {
-        *bytes_str = _PyBytesWriter_Prepare(bytes_writer, *bytes_str, sz);
+        Py_ssize_t resize = PyBytesWriter_GetSize(bytes_writer) + sz;
+        *bytes_str = PyBytesWriter_ResizeAndUpdatePointer(bytes_writer, resize,
+                                                          *bytes_str);
         if (*bytes_str == NULL)
             return -1;
     }
@@ -2510,7 +2516,7 @@ _PyLong_FormatWriter(_PyUnicodeWriter *writer,
 }
 
 char*
-_PyLong_FormatBytesWriter(_PyBytesWriter *writer, char *str,
+_PyLong_FormatBytesWriter(PyBytesWriter *writer, char *str,
                           PyObject *obj,
                           int base, int alternate)
 {