[3.15] gh-157242: Add more PyBytesWriter tests

Doc/c-api/bytes.rst

@@ -299,8 +299,8 @@ object.
 
    A bytes writer instance.
 
-   The API is **not thread safe**: a writer should only be used by a single
-   thread at the same time.
+   The API is **not thread safe**. To share a writer with multiple threads, a
+   critical section or a lock is needed.
 
    The instance must be destroyed by :c:func:`PyBytesWriter_Finish` on
    success, or :c:func:`PyBytesWriter_Discard` on error.
@@ -436,7 +436,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

@@ -1,7 +1,9 @@
 import sys
+import threading
 import unittest
 from test import support
 from test.support import import_helper
+from test.support import threading_helper
 
 _testlimitedcapi = import_helper.import_module('_testlimitedcapi')
 _testcapi = import_helper.import_module('_testcapi')
@@ -318,6 +320,10 @@ def test_join(self):
 
 class BaseWriterTest:
     RESULT_TYPE = NotImplementedError
+    SMALL_BUFFER = 11  # bytes
+    assert SMALL_BUFFER < _testcapi.PyBytesWriter_small_buffer
+    LARGE_BUFFER = _testcapi.PyBytesWriter_small_buffer + 17  # bytes
+    NEW_BYTE = b'\xff'
 
     def create_writer(self, alloc=0, string=b''):
         raise NotImplementedError
@@ -334,23 +340,60 @@ def test_create(self):
         self.assertEqual(writer.get_size(), 0)
         self.bytes_equal(writer.finish(), b'')
 
-        writer = self.create_writer(3, b'abc')
+        writer = self.create_writer(3)
+        writer.write(0, b'abc')
         self.assertEqual(writer.get_size(), 3)
         self.bytes_equal(writer.finish(), b'abc')
 
+    @unittest.skipUnless(support.Py_DEBUG, 'need Py_DEBUG')
+    def test_get_data(self):
+        # Test PyBytesWriter_GetData()
+        writer = self.create_writer(6)
+        NEW_BYTE = self.NEW_BYTE
+        self.assertEqual(writer.get_data(), NEW_BYTE * 6)
+        writer.write(0, b'abc')
+        self.assertEqual(writer.get_data(), b'abc' + NEW_BYTE * 3)
+        writer.write(3, b'123')
+        self.assertEqual(writer.get_data(), b'abc123')
+
+        # Switch from small buffer to large buffer
+        small, large = self.SMALL_BUFFER, self.LARGE_BUFFER
+        writer = self.create_writer(small)
+        self.assertEqual(writer.get_data(), NEW_BYTE * small)
+        writer.write(0, b's' * small)
+        self.assertEqual(writer.get_data(), b's' * small)
+        writer.resize(large)
+        self.assertEqual(writer.get_data(), b's' * small + NEW_BYTE * (large - small))
+        writer.write(small, b'L' * (large - small))
+        self.assertEqual(writer.get_data(), b's' * small + b'L' * (large - small))
+
+        # Resize large buffer
+        small, large = self.SMALL_BUFFER, self.LARGE_BUFFER
+        writer = self.create_writer(large)
+        self.assertEqual(writer.get_data(), NEW_BYTE * large)
+        writer.write(0, b'L' * large)
+        self.assertEqual(writer.get_data(), b'L' * large)
+        writer.resize(large + 10)
+        self.assertEqual(writer.get_data(), b'L' * large + NEW_BYTE * 10)
+        writer.write(large, b'#' * 10)
+        self.assertEqual(writer.get_data(), b'L' * large + b'#' * 10)
+
     def test_finish_with_size(self):
         # Test PyBytesWriter_FinishWithSize()
-        writer = self.create_writer(10, b'abcdef')
+        writer = self.create_writer(10)
+        writer.write(0, b'abc123')
         self.assertEqual(writer.get_size(), 10)
         self.bytes_equal(writer.finish_with_size(3), b'abc')
 
         # Error if the size is negative
-        writer = self.create_writer(3, b'abc')
+        writer = self.create_writer(3, )
+        writer.write(0, b'abc')
         with self.assertRaises(ValueError):
             writer.finish_with_size(-3)
 
         # Error if the requested size is larger than the allocated size
-        writer = self.create_writer(3, b'abc')
+        writer = self.create_writer(3)
+        writer.write(0, b'abc')
         with self.assertRaises(ValueError):
             writer.finish_with_size(4)
 
@@ -368,44 +411,118 @@ def test_write_bytes(self):
     def test_resize(self):
         # Test PyBytesWriter_Resize()
         writer = self.create_writer()
-        writer.resize(len(b'number=123456'), b'number=123456')
-        writer.resize(len(b'number=123456'), b'')
-        self.assertEqual(writer.get_size(), len(b'number=123456'))
-        self.bytes_equal(writer.finish(), b'number=123456')
+        writer.resize(len(b'hello'))
+        writer.write(0, b'hello')
+        self.assertEqual(writer.finish(), b'hello')
 
         writer = self.create_writer()
-        writer.resize(0, b'')
-        writer.resize(len(b'number=123456'), b'number=123456')
-        self.bytes_equal(writer.finish(), b'number=123456')
-
+        writer.resize(0)  # noop
+        writer.resize(len(b'number'))
+        writer.write(0, b'number')
+        writer.resize(len(b'number='))
+        writer.write(len(b'number'), b'=')
+        writer.resize(len(b'number=123'), )
+        writer.write(len(b'number='), b'123')
+        writer.resize(len(b'number=123'))  # noop
+        self.assertEqual(writer.finish(), b'number=123')
+
+        # Switch from small buffer to large buffer
         writer = self.create_writer()
-        writer.resize(len(b'number='), b'number=')
-        writer.resize(len(b'number=123456'), b'123456')
-        self.bytes_equal(writer.finish(), b'number=123456')
-
+        small, large = self.SMALL_BUFFER, self.LARGE_BUFFER
+        writer.resize(small)
+        writer.write(0, b's' * small)
+        writer.resize(large)
+        writer.write(small, b'L' * (large - small))
+        self.assertEqual(writer.finish(),
+                         b's' * small + b'L' * (large - small))
+
+        # Make sure that it's possible to write after a resize to zero
+        # when a bytes/bytearray object is allocated.
         writer = self.create_writer()
-        writer.resize(len(b'number='), b'number=')
-        writer.resize(len(b'number='), b'')
-        writer.resize(len(b'number=123456'), b'123456')
-        self.bytes_equal(writer.finish(), b'number=123456')
+        writer.resize(self.LARGE_BUFFER)
+        writer.resize(0)
+        writer.write_bytes(b'abc', 3)
+        self.assertEqual(writer.finish(), b'abc')
+
+        # invalid size
+        for size in (self.SMALL_BUFFER, self.LARGE_BUFFER):
+            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.resize(-1)
+                with self.assertRaises((MemoryError, OverflowError)):
+                    writer.resize(_testcapi.PY_SSIZE_T_MAX)
+                self.assertEqual(writer.finish(), b'x' * size)
+
+    def test_grow(self):
+        # Test PyBytesWriter_Grow()
+        writer = self.create_writer(0)
+        writer.grow(len(b'number=123'))
+        writer.write(0, b'number=123')
+        self.assertEqual(writer.finish(), b'number=123')
 
         writer = self.create_writer()
-        writer.resize(len(b'number'), b'number')
-        writer.resize(len(b'number='), b'=')
-        writer.resize(len(b'number=123'), b'123')
-        writer.resize(len(b'number=123456'), b'456')
-        self.bytes_equal(writer.finish(), b'number=123456')
+        writer.grow(0)  # noop
+        writer.grow(len(b'number'))
+        writer.write(0, b'number')
+        writer.grow(len(b'='))
+        writer.write(len(b'number'), b'=')
+        writer.grow(len(b'123'), )
+        writer.write(len(b'number='), b'123')
+        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
+        writer.grow(small)
+        writer.write(0, b's' * small)
+        writer.grow(large - small)
+        writer.write(small, b'L' * (large - small))
+        self.assertEqual(writer.finish(),
+                         b's' * small + b'L' * (large - small))
+
+        # invalid size
+        for size in (self.SMALL_BUFFER, self.LARGE_BUFFER):
+            with self.subTest(size=size):
+                writer = self.create_writer()
+                writer.write_bytes(b'x' * size, -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)
 
     @support.nomemtest
     def test_resize_error(self):
-        small_buffer = _testcapi.PyBytesWriter_small_buffer
-        init = b'x' * (small_buffer * 2)
-        writer = self.create_writer(len(init), init)
+        # 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, b'')
+                writer.resize(size)
         finally:
             _testcapi.remove_mem_hooks()
         suffix = b'still working'
@@ -443,36 +560,71 @@ def test_format_s_memory_error(self):
             finally:
                 _testcapi.remove_mem_hooks()
         writer.write_bytes(b'after', -1)
-        self.bytes_equal(writer.finish(), b'after')
+
+    @threading_helper.requires_working_threading()
+    def test_thread(self):
+        # PyBytesWriter can be used by multiple threads: it's up to the caller
+        # to implement a lock to prevent concurrent accesses.
+        writer = self.create_writer(0)
+        size = None
+        data = None
+
+        def thread_func(writer, LARGE_BUFFER):
+            nonlocal size, data
+
+            # create a bytes object for the buffer
+            writer.write_bytes(b'x' * LARGE_BUFFER, LARGE_BUFFER)
+
+            # so we can check a write with a bytes object
+            writer.write_bytes(b'yz', 2)
+            writer.format_i(b'i=%i', 5)
+            writer.resize(10)
+            data = writer.get_data()
+            size = writer.get_size()
+
+        thread = threading.Thread(target=thread_func,
+                                  args=(writer, self.LARGE_BUFFER))
+        thread.start()
+        threading_helper.join_thread(thread)
+
+        self.assertEqual(size, 10)
+        self.assertEqual(data, b'x' * 10)
+        self.assertEqual(writer.finish(), b'x' * 10)
 
 
 class BytesWriterTest(BaseWriterTest, unittest.TestCase):
     RESULT_TYPE = bytes
 
-    def create_writer(self, alloc=0, string=b''):
+    def create_writer(self, size=0):
         # Test PyBytesWriter_Create()
-        return _testcapi.PyBytesWriter(alloc, string, 0)
+        return _testcapi.PyBytesWriter(size, 0)
 
     # Only PyBytesWriter_Create() returns singletons
     def test_singletons(self):
         empty = b''
         singletons = {ch: bytes((ch,)) for ch in range(256)}
-        small_buffer = _testcapi.PyBytesWriter_small_buffer
 
         writer = self.create_writer()
         self.assertIs(writer.finish(), empty)
 
-        # Test writer larger than small_buffer
+        # Large buffer
         writer = self.create_writer()
-        unused_text = b'x' * (small_buffer * 2)
+        unused_text = b'x' * self.LARGE_BUFFER
         writer.write_bytes(unused_text, len(unused_text))
         self.assertIs(writer.finish_with_size(0), empty)
 
+        # Large buffer with resize
+        writer = self.create_writer()
+        unused_text = b'x' * self.LARGE_BUFFER
+        writer.write_bytes(unused_text, len(unused_text))
+        writer.resize(0)
+        self.assertIs(writer.finish(), empty)
+
         for ch in range(256):
-            text = bytes((ch,))
+            byte = bytes((ch,))
 
             writer = self.create_writer()
-            writer.write_bytes(text, 1)
+            writer.write_bytes(byte, len(byte))
             self.assertIs(writer.finish(), singletons[ch])
 
     def test_example_abc(self):
@@ -488,9 +640,9 @@ def test_example_highlevel(self):
 class ByteArrayWriterTest(BaseWriterTest, unittest.TestCase):
     RESULT_TYPE = bytearray
 
-    def create_writer(self, alloc=0, string=b''):
+    def create_writer(self, size=0):
         # Test private _PyBytesWriter_CreateByteArray()
-        return _testcapi.PyBytesWriter(alloc, string, 1)
+        return _testcapi.PyBytesWriter(size, 1)
 
 
 if __name__ == "__main__":