Add PyBytesWriter_ResizeAndUpdatePointer() function

Author: vstinner

Include/cpython/bytesobject.h

@@ -78,3 +78,7 @@ PyAPI_FUNC(int) PyBytesWriter_Format(
 PyAPI_FUNC(int) PyBytesWriter_Resize(
     PyBytesWriter *writer,
     Py_ssize_t size);
+PyAPI_FUNC(void*) PyBytesWriter_ResizeAndUpdatePointer(
+    PyBytesWriter *writer,
+    Py_ssize_t size,
+    void *data);

Modules/_pickle.c

@@ -2632,11 +2632,10 @@ raw_unicode_escape(PyObject *obj)
         if (ch >= 0x10000) {
             /* -1: subtract 1 preallocated byte */
             alloc += 10-1;
-            Py_ssize_t pos = p - (char*)PyBytesWriter_GetData(writer);
-            if (PyBytesWriter_Resize(writer, alloc) < 0) {
+            p = PyBytesWriter_ResizeAndUpdatePointer(writer, alloc, p);
+            if (p == NULL) {
                 goto error;
             }
-            p = (char*)PyBytesWriter_GetData(writer) + pos;
 
             *p++ = '\\';
             *p++ = 'U';
@@ -2656,11 +2655,10 @@ raw_unicode_escape(PyObject *obj)
         {
             /* -1: subtract 1 preallocated byte */
             alloc += 6-1;
-            Py_ssize_t pos = p - (char*)PyBytesWriter_GetData(writer);
-            if (PyBytesWriter_Resize(writer, alloc) < 0) {
+            p = PyBytesWriter_ResizeAndUpdatePointer(writer, alloc, p);
+            if (p == NULL) {
                 goto error;
             }
-            p = (char*)PyBytesWriter_GetData(writer) + pos;
 
             *p++ = '\\';
             *p++ = 'u';

Objects/bytesobject.c

@@ -222,11 +222,10 @@ bytes_fromformat(PyBytesWriter *writer, Py_ssize_t writer_pos,
     do { \
         size_t len = (len_expr); \
         alloc += len; \
-        Py_ssize_t pos = s - (char*)PyBytesWriter_GetData(writer); \
-        if (PyBytesWriter_Resize(writer, alloc) < 0) { \
+        s = PyBytesWriter_ResizeAndUpdatePointer(writer, alloc, s); \
+        if (s == NULL) { \
             goto error; \
         } \
-        s = (char*)PyBytesWriter_GetData(writer) + pos; \
         memcpy(s, (str), len); \
         s += len; \
     } while (0)
@@ -2889,11 +2888,10 @@ _PyBytes_FromList(PyObject *x)
         }
 
         if (i >= size) {
-            Py_ssize_t pos = str - (char*)PyBytesWriter_GetData(writer);
-            if (PyBytesWriter_Resize(writer, size + 1) < 0) {
+            str = PyBytesWriter_ResizeAndUpdatePointer(writer, size + 1, str);
+            if (str == NULL) {
                 goto error;
             }
-            str = (char*)PyBytesWriter_GetData(writer) + pos;
             size = PyBytesWriter_GetAllocated(writer);
         }
         *str++ = (char) value;
@@ -3992,3 +3990,15 @@ PyBytesWriter_Format(PyBytesWriter *writer, const char *format, ...)
     Py_ssize_t size = buf - byteswriter_data(writer);
     return PyBytesWriter_Resize(writer, size);
 }
+
+
+void*
+PyBytesWriter_ResizeAndUpdatePointer(PyBytesWriter *writer, Py_ssize_t size,
+                                     void *data)
+{
+    Py_ssize_t pos = (char*)data - byteswriter_data(writer);
+    if (PyBytesWriter_Resize(writer, size) < 0) {
+        return NULL;
+    }
+    return byteswriter_data(writer) + pos;
+}