gh-157242: Allow negative grow in PyBytesWriter_Grow() (#157375)

PyBytesWriter_Resize() doesn't need to call byteswriter_resize() if
the buffer is already large enough.

Author: vstinner

Doc/c-api/bytes.rst

@@ -429,7 +429,7 @@ Low-level API
    On success, return ``0``.
    On error, set an exception and return ``-1``.
 
-   *size* can be negative to shrink the writer.
+   *grow* can be negative to shrink the writer.
 
 .. c:function:: void* PyBytesWriter_GrowAndUpdatePointer(PyBytesWriter *writer, Py_ssize_t size, void *buf)
 

Lib/test/test_capi/test_bytes.py

@@ -461,6 +461,24 @@ def test_grow(self):
         writer.grow(0)  # noop
         self.assertEqual(writer.finish(), b'number=123')
 
+        for size in (self.SMALL_BUFFER, self.LARGE_BUFFER):
+            with self.subTest(size=size):
+                # Truncate the last byte
+                data = b'x' * size
+                writer = self.create_writer(size)
+                writer.write(0, data)
+                self.assertEqual(writer.get_data(), data)
+                writer.grow(-1)
+                self.assertEqual(writer.get_data(), data[:-1])
+                self.assertEqual(writer.finish(),  data[:-1])
+
+                # Make the buffer empty
+                writer = self.create_writer(size)
+                writer.write(0, data)
+                writer.grow(-size)
+                self.assertEqual(writer.get_data(), b'')
+                self.assertEqual(writer.finish(),  b'')
+
         # Switch from small buffer to large buffer
         writer = self.create_writer()
         small, large = self.SMALL_BUFFER, self.LARGE_BUFFER
@@ -476,8 +494,8 @@ def test_grow(self):
             with self.subTest(size=size):
                 writer = self.create_writer()
                 writer.write_bytes(b'x' * size, -1)
-                with self.assertRaisesRegex(ValueError, 'size must be >= 0'):
-                    writer.grow(-1)
+                with self.assertRaisesRegex(ValueError, 'invalid size'):
+                    writer.grow(-size - 1)
                 with self.assertRaises(MemoryError):
                     writer.grow(_testcapi.PY_SSIZE_T_MAX)
                 self.assertEqual(writer.finish(), b'x' * size)

Objects/bytesobject.c

@@ -3845,8 +3845,13 @@ PyBytesWriter_Resize(PyBytesWriter *writer, Py_ssize_t size)
         PyErr_SetString(PyExc_ValueError, "size must be >= 0");
         return -1;
     }
-    if (byteswriter_resize(writer, size, 1) < 0) {
-        return -1;
+    if (writer->size < size) {
+        if (byteswriter_resize(writer, size, 1) < 0) {
+            return -1;
+        }
+    }
+    else {
+        // The buffer is already large enough. Never shrink the buffer.
     }
     writer->size = size;
     return 0;
@@ -3866,22 +3871,26 @@ _PyBytesWriter_ResizeAndUpdatePointer(PyBytesWriter *writer, Py_ssize_t size,
 
 
 int
-PyBytesWriter_Grow(PyBytesWriter *writer, Py_ssize_t size)
+PyBytesWriter_Grow(PyBytesWriter *writer, Py_ssize_t grow)
 {
-    if (size < 0) {
-        PyErr_SetString(PyExc_ValueError, "size must be >= 0");
-        return -1;
-    }
-    if (size == 0) {
+    if (grow == 0) {
         // Nothing to do
         return 0;
     }
 
-    if (size > PY_SSIZE_T_MAX - writer->size) {
-        PyErr_NoMemory();
-        return -1;
+    if (grow >= 0) {
+        if (grow > PY_SSIZE_T_MAX - writer->size) {
+            PyErr_NoMemory();
+            return -1;
+        }
+    }
+    else {
+        if (writer->size + grow < 0) {
+            PyErr_SetString(PyExc_ValueError, "invalid size");
+            return -1;
+        }
     }
-    size = writer->size + size;
+    Py_ssize_t size = writer->size + grow;
 
     if (byteswriter_resize(writer, size, 1) < 0) {
         return -1;