Add PyBytesWriter_GrowAndUpdatePointer()

Author: vstinner

Include/cpython/bytesobject.h

@@ -84,3 +84,10 @@ PyAPI_FUNC(void*) PyBytesWriter_ResizeAndUpdatePointer(
     PyBytesWriter *writer,
     Py_ssize_t size,
     void *data);
+PyAPI_FUNC(int) PyBytesWriter_Grow(
+    PyBytesWriter *writer,
+    Py_ssize_t size);
+PyAPI_FUNC(void*) PyBytesWriter_GrowAndUpdatePointer(
+    PyBytesWriter *writer,
+    Py_ssize_t size,
+    void *data);

Modules/_pickle.c

@@ -2631,8 +2631,7 @@ raw_unicode_escape(PyObject *obj)
         /* Map 32-bit characters to '\Uxxxxxxxx' */
         if (ch >= 0x10000) {
             /* -1: subtract 1 preallocated byte */
-            alloc += 10-1;
-            p = PyBytesWriter_ResizeAndUpdatePointer(writer, alloc, p);
+            p = PyBytesWriter_GrowAndUpdatePointer(writer, 10-1, p);
             if (p == NULL) {
                 goto error;
             }
@@ -2654,8 +2653,7 @@ raw_unicode_escape(PyObject *obj)
                  ch == 0x1a)
         {
             /* -1: subtract 1 preallocated byte */
-            alloc += 6-1;
-            p = PyBytesWriter_ResizeAndUpdatePointer(writer, alloc, p);
+            p = PyBytesWriter_GrowAndUpdatePointer(writer, 6-1, p);
             if (p == NULL) {
                 goto error;
             }

Objects/bytesobject.c

@@ -221,8 +221,7 @@ bytes_fromformat(PyBytesWriter *writer, Py_ssize_t writer_pos,
 #define WRITE_BYTES_LEN(str, len_expr) \
     do { \
         size_t len = (len_expr); \
-        alloc += len; \
-        s = PyBytesWriter_ResizeAndUpdatePointer(writer, alloc, s); \
+        s = PyBytesWriter_GrowAndUpdatePointer(writer, len, s); \
         if (s == NULL) { \
             goto error; \
         } \
@@ -459,8 +458,7 @@ formatfloat(PyObject *v, int flags, int prec, int type,
 
     len = strlen(p);
     if (writer != NULL) {
-        Py_ssize_t resize = PyBytesWriter_GetSize(writer) + len;
-        str = PyBytesWriter_ResizeAndUpdatePointer(writer, resize, str);
+        str = PyBytesWriter_GrowAndUpdatePointer(writer, len, str);
         if (str == NULL) {
             PyMem_Free(p);
             return NULL;
@@ -982,8 +980,7 @@ _PyBytes_FormatEx(const char *format, Py_ssize_t format_len,
                 alloc++;
             /* 2: size preallocated for %s */
             if (alloc > 2) {
-                Py_ssize_t resize = PyBytesWriter_GetSize(writer) + alloc - 2;
-                res = PyBytesWriter_ResizeAndUpdatePointer(writer, resize, res);
+                res = PyBytesWriter_GrowAndUpdatePointer(writer, alloc - 2, res);
                 if (res == NULL) {
                     goto error;
                 }
@@ -3971,41 +3968,50 @@ PyBytesWriter_Resize(PyBytesWriter *writer, Py_ssize_t size)
 
 
 int
-PyBytesWriter_WriteBytes(PyBytesWriter *writer,
-                         const void *bytes, Py_ssize_t size)
+PyBytesWriter_Grow(PyBytesWriter *writer, Py_ssize_t size)
 {
     if (size < 0) {
-        size = strlen(bytes);
+        PyErr_SetString(PyExc_ValueError, "size must be >= 0");
+        return -1;
     }
 
-    Py_ssize_t pos = writer->size;
-    if (size > PY_SSIZE_T_MAX - pos) {
+    if (size > PY_SSIZE_T_MAX - writer->size) {
         PyErr_NoMemory();
         return -1;
     }
-    Py_ssize_t total = pos + size;
+    size = writer->size + size;
 
-    if (PyBytesWriter_Resize(writer, total) < 0) {
+    if (byteswriter_resize(writer, size, 1) < 0) {
         return -1;
     }
-    char *buf = byteswriter_data(writer);
-    memcpy(buf + pos, bytes, size);
+    writer->size = size;
     return 0;
 }
 
 
 int
-PyBytesWriter_Format(PyBytesWriter *writer, const char *format, ...)
+PyBytesWriter_WriteBytes(PyBytesWriter *writer,
+                         const void *bytes, Py_ssize_t size)
 {
+    if (size < 0) {
+        size = strlen(bytes);
+    }
+
     Py_ssize_t pos = writer->size;
-    Py_ssize_t format_len = strlen(format);
-    if (format_len > PY_SSIZE_T_MAX - pos) {
-        PyErr_NoMemory();
+    if (PyBytesWriter_Grow(writer, size) < 0) {
         return -1;
     }
-    Py_ssize_t alloc = pos + format_len;
+    char *buf = byteswriter_data(writer);
+    memcpy(buf + pos, bytes, size);
+    return 0;
+}
+
 
-    if (PyBytesWriter_Resize(writer, alloc) < 0) {
+int
+PyBytesWriter_Format(PyBytesWriter *writer, const char *format, ...)
+{
+    Py_ssize_t pos = writer->size;
+    if (PyBytesWriter_Grow(writer, strlen(format)) < 0) {
         return -1;
     }
 
@@ -4029,3 +4035,15 @@ PyBytesWriter_ResizeAndUpdatePointer(PyBytesWriter *writer, Py_ssize_t size,
     }
     return byteswriter_data(writer) + pos;
 }
+
+
+void*
+PyBytesWriter_GrowAndUpdatePointer(PyBytesWriter *writer, Py_ssize_t size,
+                                   void *data)
+{
+    Py_ssize_t pos = (char*)data - byteswriter_data(writer);
+    if (PyBytesWriter_Grow(writer, size) < 0) {
+        return NULL;
+    }
+    return byteswriter_data(writer) + pos;
+}

Objects/longobject.c

@@ -2047,9 +2047,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);
-        Py_ssize_t resize = PyBytesWriter_GetSize(bytes_writer) + size;
-        *bytes_str = PyBytesWriter_ResizeAndUpdatePointer(bytes_writer, resize,
-                                                          *bytes_str);
+        *bytes_str = PyBytesWriter_GrowAndUpdatePointer(bytes_writer, size,
+                                                        *bytes_str);
         if (*bytes_str == NULL) {
             goto error;
         }
@@ -2212,9 +2211,8 @@ long_to_decimal_string_internal(PyObject *aa,
         }
     }
     else if (bytes_writer) {
-        Py_ssize_t resize = PyBytesWriter_GetSize(bytes_writer) + strlen;
-        *bytes_str = PyBytesWriter_ResizeAndUpdatePointer(bytes_writer, resize,
-                                                          *bytes_str);
+        *bytes_str = PyBytesWriter_GrowAndUpdatePointer(bytes_writer, strlen,
+                                                        *bytes_str);
         if (*bytes_str == NULL) {
             Py_DECREF(scratch);
             return -1;
@@ -2385,9 +2383,8 @@ long_format_binary(PyObject *aa, int base, int alternate,
             return -1;
     }
     else if (bytes_writer) {
-        Py_ssize_t resize = PyBytesWriter_GetSize(bytes_writer) + sz;
-        *bytes_str = PyBytesWriter_ResizeAndUpdatePointer(bytes_writer, resize,
-                                                          *bytes_str);
+        *bytes_str = PyBytesWriter_GrowAndUpdatePointer(bytes_writer, sz,
+                                                        *bytes_str);
         if (*bytes_str == NULL)
             return -1;
     }