Add _PyBytesWriter_CreateByteArray()

vstinner · vstinner · commit 62a15be6569a · 2025-03-25T18:08:08.000+01:00
Convert _PyBytes_FromHex().
diff --git a/Include/cpython/bytesobject.h b/Include/cpython/bytesobject.h
@@ -48,6 +48,8 @@ typedef struct PyBytesWriter PyBytesWriter;
 
 PyAPI_FUNC(PyBytesWriter *) PyBytesWriter_Create(
     Py_ssize_t size);
+PyAPI_FUNC(PyBytesWriter*) _PyBytesWriter_CreateByteArray(
+    Py_ssize_t size);
 PyAPI_FUNC(void) PyBytesWriter_Discard(
     PyBytesWriter *writer);
 PyAPI_FUNC(PyObject*) PyBytesWriter_Finish(
diff --git a/Lib/test/test_capi/test_bytes.py b/Lib/test/test_capi/test_bytes.py
@@ -291,82 +291,89 @@ def test_join(self):
             bytes_join(b'', NULL)
 
 
-class PyBytesWriterTest(unittest.TestCase):
+class BytesWriterTest(unittest.TestCase):
     SMALL_BUFFER = 256  # bytes
+    result_type = bytes
 
     def create_writer(self, alloc=0, string=b''):
-        return _testcapi.PyBytesWriter(alloc, string)
+        return _testcapi.PyBytesWriter(alloc, string, 0)
 
     def test_create(self):
         # Test PyBytesWriter_Create()
         writer = self.create_writer()
         self.assertEqual(writer.get_size(), 0)
         self.assertEqual(writer.get_allocated(), self.SMALL_BUFFER)
-        self.assertEqual(writer.finish(), b'')
+        self.assertEqual(writer.finish(), self.result_type(b''))
 
         writer = self.create_writer(3, b'abc')
         self.assertEqual(writer.get_size(), 3)
         self.assertEqual(writer.get_allocated(), self.SMALL_BUFFER)
-        self.assertEqual(writer.finish(), b'abc')
+        self.assertEqual(writer.finish(), self.result_type(b'abc'))
 
         writer = self.create_writer(10, b'abc')
         self.assertEqual(writer.get_size(), 10)
         self.assertEqual(writer.get_allocated(), self.SMALL_BUFFER)
-        self.assertEqual(writer.finish_with_size(3), b'abc')
+        self.assertEqual(writer.finish_with_size(3), self.result_type(b'abc'))
 
     def test_write_bytes(self):
          # Test PyBytesWriter_WriteBytes()
          writer = self.create_writer()
          writer.write_bytes(b'Hello World!', -1)
-         self.assertEqual(writer.finish(), b'Hello World!')
+         self.assertEqual(writer.finish(), self.result_type(b'Hello World!'))
 
          writer = self.create_writer()
          writer.write_bytes(b'Hello ', -1)
          writer.write_bytes(b'World! <truncated>', 6)
-         self.assertEqual(writer.finish(), b'Hello World!')
+         self.assertEqual(writer.finish(), self.result_type(b'Hello World!'))
 
     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.assertEqual(writer.finish(), b'number=123456')
+        self.assertEqual(writer.finish(), self.result_type(b'number=123456'))
 
         writer = self.create_writer()
         writer.resize(0, b'')
         writer.resize(len(b'number=123456'), b'number=123456')
-        self.assertEqual(writer.finish(), b'number=123456')
+        self.assertEqual(writer.finish(), self.result_type(b'number=123456'))
 
         writer = self.create_writer()
         writer.resize(len(b'number='), b'number=')
         writer.resize(len(b'number=123456'), b'123456')
-        self.assertEqual(writer.finish(), b'number=123456')
+        self.assertEqual(writer.finish(), self.result_type(b'number=123456'))
 
         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.assertEqual(writer.finish(), b'number=123456')
+        self.assertEqual(writer.finish(), self.result_type(b'number=123456'))
 
         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.assertEqual(writer.finish(), b'number=123456')
+        self.assertEqual(writer.finish(), self.result_type(b'number=123456'))
 
     def test_format_i(self):
         # Test PyBytesWriter_Format()
         writer = self.create_writer()
         writer.format_i(b'x=%i', 123456)
-        self.assertEqual(writer.finish(), b'x=123456')
+        self.assertEqual(writer.finish(), self.result_type(b'x=123456'))
 
         writer = self.create_writer()
         writer.format_i(b'x=%i, ', 123)
         writer.format_i(b'y=%i', 456)
-        self.assertEqual(writer.finish(), b'x=123, y=456')
+        self.assertEqual(writer.finish(), self.result_type(b'x=123, y=456'))
 
 
+class ByteArrayWriterTest(BytesWriterTest):
+    result_type = bytearray
+
+    def create_writer(self, alloc=0, string=b''):
+        return _testcapi.PyBytesWriter(alloc, string, 1)
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/Modules/_testcapi/bytes.c b/Modules/_testcapi/bytes.c
@@ -88,11 +88,18 @@ writer_init(PyObject *self_raw, PyObject *args, PyObject *kwargs)
     Py_ssize_t alloc;
     char *str;
     Py_ssize_t str_size;
-    if (!PyArg_ParseTuple(args, "ny#", &alloc, &str, &str_size)) {
+    int use_bytearray;
+    if (!PyArg_ParseTuple(args, "ny#i",
+                          &alloc, &str, &str_size, &use_bytearray)) {
         return -1;
     }
 
-    self->writer = PyBytesWriter_Create(alloc);
+    if (use_bytearray) {
+        self->writer = _PyBytesWriter_CreateByteArray(alloc);
+    }
+    else {
+        self->writer = PyBytesWriter_Create(alloc);
+    }
     if (self->writer == NULL) {
         return -1;
     }
diff --git a/Objects/bytesobject.c b/Objects/bytesobject.c
@@ -2515,17 +2515,13 @@ bytes_fromhex_impl(PyTypeObject *type, PyObject *string)
 PyObject*
 _PyBytes_FromHex(PyObject *string, int use_bytearray)
 {
-    char *buf;
     Py_ssize_t hexlen, invalid_char;
     unsigned int top, bot;
     const Py_UCS1 *str, *start, *end;
-    _PyBytesWriter writer;
+    PyBytesWriter *writer = NULL;
     Py_buffer view;
     view.obj = NULL;
 
-    _PyBytesWriter_Init(&writer);
-    writer.use_bytearray = use_bytearray;
-
     if (PyUnicode_Check(string)) {
         hexlen = PyUnicode_GET_LENGTH(string);
 
@@ -2561,10 +2557,11 @@ _PyBytes_FromHex(PyObject *string, int use_bytearray)
     }
 
     /* This overestimates if there are spaces */
-    buf = _PyBytesWriter_Alloc(&writer, hexlen / 2);
-    if (buf == NULL) {
+    writer = _PyBytesWriter_CreateByteArray(hexlen / 2);
+    if (writer == NULL) {
         goto release_buffer;
     }
+    char *buf = PyBytesWriter_GetData(writer);
 
     start = str;
     end = str + hexlen;
@@ -2603,7 +2600,7 @@ _PyBytes_FromHex(PyObject *string, int use_bytearray)
     if (view.obj != NULL) {
        PyBuffer_Release(&view);
     }
-    return _PyBytesWriter_Finish(&writer, buf);
+    return PyBytesWriter_FinishWithEndPointer(writer, buf);
 
   error:
     if (invalid_char == -1) {
@@ -2614,7 +2611,7 @@ _PyBytes_FromHex(PyObject *string, int use_bytearray)
                      "non-hexadecimal number found in "
                      "fromhex() arg at position %zd", invalid_char);
     }
-    _PyBytesWriter_Dealloc(&writer);
+    PyBytesWriter_Discard(writer);
 
   release_buffer:
     if (view.obj != NULL) {
@@ -3737,6 +3734,7 @@ struct PyBytesWriter {
     char small_buffer[256];
     PyObject *obj;
     Py_ssize_t size;
+    int use_bytearray;
 };
 
 
@@ -3758,6 +3756,9 @@ byteswriter_allocated(PyBytesWriter *writer)
     if (writer->obj == NULL) {
         return sizeof(writer->small_buffer);
     }
+    else if (writer->use_bytearray) {
+        return PyByteArray_GET_SIZE(writer->obj);
+    }
     else {
         return PyBytes_GET_SIZE(writer->obj);
     }
@@ -3778,15 +3779,8 @@ byteswriter_resize(PyBytesWriter *writer, Py_ssize_t size, int overallocate)
 {
     assert(size >= 0);
 
-    if (writer->obj == NULL) {
-        if ((size_t)size <= sizeof(writer->small_buffer)) {
-            return 0;
-        }
-    }
-    else {
-        if (size <= PyBytes_GET_SIZE(writer->obj)) {
-            return 0;
-        }
+    if (size <= byteswriter_allocated(writer)) {
+        return 0;
     }
 
     if (overallocate) {
@@ -3796,11 +3790,28 @@ byteswriter_resize(PyBytesWriter *writer, Py_ssize_t size, int overallocate)
     }
 
     if (writer->obj != NULL) {
-        if (_PyBytes_Resize(&writer->obj, size)) {
-            return -1;
+        if (writer->use_bytearray) {
+            if (PyByteArray_Resize(writer->obj, size)) {
+                return -1;
+            }
+        }
+        else {
+            if (_PyBytes_Resize(&writer->obj, size)) {
+                return -1;
+            }
         }
         assert(writer->obj != NULL);
     }
+    else if (writer->use_bytearray) {
+        writer->obj = PyByteArray_FromStringAndSize(NULL, size);
+        if (writer->obj == NULL) {
+            return -1;
+        }
+        assert((size_t)size > sizeof(writer->small_buffer));
+        memcpy(PyByteArray_AS_STRING(writer->obj),
+               writer->small_buffer,
+               sizeof(writer->small_buffer));
+    }
     else {
         writer->obj = PyBytes_FromStringAndSize(NULL, size);
         if (writer->obj == NULL) {
@@ -3815,8 +3826,8 @@ byteswriter_resize(PyBytesWriter *writer, Py_ssize_t size, int overallocate)
 }
 
 
-PyBytesWriter*
-PyBytesWriter_Create(Py_ssize_t size)
+static PyBytesWriter*
+byteswriter_create(Py_ssize_t size, int use_bytearray)
 {
     if (size < 0) {
         PyErr_SetString(PyExc_ValueError, "size must be >= 0");
@@ -3833,6 +3844,7 @@ PyBytesWriter_Create(Py_ssize_t size)
     }
     writer->obj = NULL;
     writer->size = 0;
+    writer->use_bytearray = use_bytearray;
 
     if (size >= 1) {
         if (byteswriter_resize(writer, size, 0) < 0) {
@@ -3844,6 +3856,18 @@ PyBytesWriter_Create(Py_ssize_t size)
     return writer;
 }
 
+PyBytesWriter*
+PyBytesWriter_Create(Py_ssize_t size)
+{
+    return byteswriter_create(size, 0);
+}
+
+PyBytesWriter*
+_PyBytesWriter_CreateByteArray(Py_ssize_t size)
+{
+    return byteswriter_create(size, 1);
+}
+
 
 void
 PyBytesWriter_Discard(PyBytesWriter *writer)
@@ -3865,14 +3889,26 @@ PyBytesWriter_FinishWithSize(PyBytesWriter *writer, Py_ssize_t size)
         result = bytes_get_empty();
     }
     else if (writer->obj != NULL) {
-        if (size != PyBytes_GET_SIZE(writer->obj)) {
-            if (_PyBytes_Resize(&writer->obj, size)) {
-                goto error;
+        if (writer->use_bytearray) {
+            if (size != PyByteArray_GET_SIZE(writer->obj)) {
+                if (PyByteArray_Resize(writer->obj, size)) {
+                    goto error;
+                }
+            }
+        }
+        else {
+            if (size != PyBytes_GET_SIZE(writer->obj)) {
+                if (_PyBytes_Resize(&writer->obj, size)) {
+                    goto error;
+                }
             }
         }
         result = writer->obj;
         writer->obj = NULL;
     }
+    else if (writer->use_bytearray) {
+        result = PyByteArray_FromStringAndSize(writer->small_buffer, size);
+    }
     else {
         result = PyBytes_FromStringAndSize(writer->small_buffer, size);
     }
