Add test_grow_error()

PyBytesWriter_Grow() no longer calls byteswriter_resize() if grow is
smaller than 0.

Author: vstinner

Lib/test/test_capi/test_bytes.py

@@ -473,6 +473,26 @@ def test_resize_canary(self):
                 self.assertEqual(writer.get_data(), b'')
                 self.assertEqual(writer.finish(),  b'')
 
+    @support.nomemtest
+    def test_resize_error(self):
+        # Test PyBytesWriter_Resize() error
+        init = b'x' * self.LARGE_BUFFER
+        writer = self.create_writer(len(init))
+        writer.write(0, init)
+        size = len(init) + 100
+        try:
+            with self.assertRaises(MemoryError):
+                _testcapi.set_nomemory(0)
+                writer.resize(size)
+        finally:
+            _testcapi.remove_mem_hooks()
+        suffix = b'still working'
+        writer.write_bytes(suffix, -1)
+        self.assertEqual(writer.finish(), init + suffix)
+
+        # Note: PyBytesWriter_Resize() leaves the buffer unchanged (no resize)
+        # if the new size is smaller than the allocated size
+
     def test_grow(self):
         # Test PyBytesWriter_Grow()
         writer = self.create_writer(0)
@@ -533,26 +553,24 @@ def test_grow_canary(self):
                 self.assertEqual(writer.get_data(), b'')
                 self.assertEqual(writer.finish(),  b'')
 
-
     @support.nomemtest
-    def test_resize_error(self):
-        # Test PyBytesWriter_Resize() error
+    def test_grow_error(self):
+        # Test PyBytesWriter_Grow() error
         init = b'x' * self.LARGE_BUFFER
         writer = self.create_writer(len(init))
         writer.write(0, init)
-        size = len(init) + 100
         try:
             with self.assertRaises(MemoryError):
                 _testcapi.set_nomemory(0)
-                writer.resize(size)
+                writer.grow(100)
         finally:
             _testcapi.remove_mem_hooks()
         suffix = b'still working'
         writer.write_bytes(suffix, -1)
         self.assertEqual(writer.finish(), init + suffix)
 
-        # Note: PyBytesWriter_Resize() leaves the buffer unchanged (no resize)
-        # if the new size is smaller than the allocated size
+        # Note: PyBytesWriter_Grow() leaves the buffer unchanged (no resize)
+        # if grow is negative.
 
     def test_format_i(self):
         # Test PyBytesWriter_Format()

Objects/bytesobject.c

@@ -3883,21 +3883,21 @@ PyBytesWriter_GetSize(PyBytesWriter *writer)
 
 
 int
-PyBytesWriter_Resize(PyBytesWriter *writer, Py_ssize_t size)
+PyBytesWriter_Resize(PyBytesWriter *writer, Py_ssize_t new_size)
 {
-    if (size < 0) {
+    if (new_size < 0) {
         PyErr_SetString(PyExc_ValueError, "size must be >= 0");
         return -1;
     }
-    if (writer->size < size) {
-        if (byteswriter_resize(writer, size, 1) < 0) {
+    if (writer->size < new_size) {
+        if (byteswriter_resize(writer, new_size, 1) < 0) {
             return -1;
         }
     }
     else {
         // The buffer is already large enough. Never shrink the buffer.
     }
-    writer->size = size;
+    writer->size = new_size;
 #ifdef Py_DEBUG
     byteswriter_write_canary_byte(writer);
 #endif
@@ -3925,24 +3925,27 @@ PyBytesWriter_Grow(PyBytesWriter *writer, Py_ssize_t grow)
         return 0;
     }
 
-    if (grow >= 0) {
+    if (grow > 0) {
         if (grow > PY_SSIZE_T_MAX - writer->size) {
             PyErr_NoMemory();
             return -1;
         }
+        Py_ssize_t new_size = writer->size + grow;
+
+        if (byteswriter_resize(writer, new_size, 1) < 0) {
+            return -1;
+        }
+        writer->size = new_size;
     }
     else {
         if (writer->size + grow < 0) {
             PyErr_SetString(PyExc_ValueError, "invalid size");
             return -1;
         }
+        // The buffer is already large enough. Never shrink the buffer.
+        writer->size = writer->size + grow;
     }
-    Py_ssize_t size = writer->size + grow;
 
-    if (byteswriter_resize(writer, size, 1) < 0) {
-        return -1;
-    }
-    writer->size = size;
 #ifdef Py_DEBUG
     byteswriter_write_canary_byte(writer);
 #endif