Add _PyBytesWriter_Resize() function

This function gives a control to the buffer size without using min_size.

2 files changed, 48 insertions, 19 deletions

diff --git a/Include/bytesobject.h b/Include/bytesobject.h
index 4046c1c..8469112 100644
--- a/Include/bytesobject.h
+++ b/Include/bytesobject.h
@@ -178,7 +178,9 @@ PyAPI_FUNC(void) _PyBytesWriter_Dealloc(_PyBytesWriter *writer);
 PyAPI_FUNC(void*) _PyBytesWriter_Alloc(_PyBytesWriter *writer,
     Py_ssize_t size);
 
-/* Add *size* bytes to the buffer.
+/* Ensure that the buffer is large enough to write *size* bytes.
+   Add size to the writer minimum size (min_size attribute).
+
    str is the current pointer inside the buffer.
    Return the updated current pointer inside the buffer.
    Raise an exception and return NULL on error. */
@@ -186,6 +188,21 @@ PyAPI_FUNC(void*) _PyBytesWriter_Prepare(_PyBytesWriter *writer,
     void *str,
     Py_ssize_t size);
 
+/* Resize the buffer to make it larger.
+   The new buffer may be larger than size bytes because of overallocation.
+   Return the updated current pointer inside the buffer.
+   Raise an exception and return NULL on error.
+
+   Note: size must be greater than the number of allocated bytes in the writer.
+
+   This function doesn't use the writer minimum size (min_size attribute).
+
+   See also _PyBytesWriter_Prepare().
+   */
+PyAPI_FUNC(void*) _PyBytesWriter_Resize(_PyBytesWriter *writer,
+    void *str,
+    Py_ssize_t size);
+
 /* Write bytes.
    Raise an exception and return NULL on error. */
 PyAPI_FUNC(void*) _PyBytesWriter_WriteBytes(_PyBytesWriter *writer,
diff --git a/Objects/bytesobject.c b/Objects/bytesobject.c
index e54f299..ae7b1ea 100644
--- a/Objects/bytesobject.c
+++ b/Objects/bytesobject.c
@@ -3997,29 +3997,14 @@ _PyBytesWriter_CheckConsistency(_PyBytesWriter *writer, char *str)
 }
 
 void*
-_PyBytesWriter_Prepare(_PyBytesWriter *writer, void *str, Py_ssize_t size)
+_PyBytesWriter_Resize(_PyBytesWriter *writer, void *str, Py_ssize_t size)
 {
     Py_ssize_t allocated, pos;
 
     _PyBytesWriter_CheckConsistency(writer, str);
-    assert(size >= 0);
-
-    if (size == 0) {
-        /* nothing to do */
-        return str;
-    }
-
-    if (writer->min_size > PY_SSIZE_T_MAX - size) {
-        PyErr_NoMemory();
-        goto error;
-    }
-    writer->min_size += size;
-
-    allocated = writer->allocated;
-    if (writer->min_size <= allocated)
-        return str;
+    assert(writer->allocated < size);
 
-    allocated = writer->min_size;
+    allocated = size;
     if (writer->overallocate
         && allocated <= (PY_SSIZE_T_MAX - allocated / OVERALLOCATE_FACTOR)) {
         /* overallocate to limit the number of realloc() */
@@ -4080,6 +4065,33 @@ error:
     return NULL;
 }
 
+void*
+_PyBytesWriter_Prepare(_PyBytesWriter *writer, void *str, Py_ssize_t size)
+{
+    Py_ssize_t new_min_size;
+
+    _PyBytesWriter_CheckConsistency(writer, str);
+    assert(size >= 0);
+
+    if (size == 0) {
+        /* nothing to do */
+        return str;
+    }
+
+    if (writer->min_size > PY_SSIZE_T_MAX - size) {
+        PyErr_NoMemory();
+        _PyBytesWriter_Dealloc(writer);
+        return NULL;
+    }
+    new_min_size = writer->min_size + size;
+
+    if (new_min_size > writer->allocated)
+        str = _PyBytesWriter_Resize(writer, str, new_min_size);
+
+    writer->min_size = new_min_size;
+    return str;
+}
+
 /* Allocate the buffer to write size bytes.
    Return the pointer to the beginning of buffer data.
    Raise an exception and return NULL on error. */