Add PyBytesWriter_WriteBytes()

Author: vstinner

Include/cpython/bytesobject.h

@@ -66,6 +66,11 @@ PyAPI_FUNC(Py_ssize_t) PyBytesWriter_GetSize(
 PyAPI_FUNC(Py_ssize_t) PyBytesWriter_GetAllocated(
     PyBytesWriter *writer);
 
+PyAPI_FUNC(int) PyBytesWriter_WriteBytes(
+    PyBytesWriter *writer,
+    const void *bytes,
+    Py_ssize_t size);
+
 PyAPI_FUNC(int) PyBytesWriter_Resize(
     PyBytesWriter *writer,
     Py_ssize_t size);

Lib/test/test_capi/test_bytes.py

@@ -297,15 +297,13 @@ class PyBytesWriterTest(unittest.TestCase):
     def create_writer(self, alloc=0, string=b''):
         return _testcapi.PyBytesWriter(alloc, string)
 
-    def test_empty(self):
+    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'')
 
-    def test_abc(self):
-        # Test PyBytesWriter_Create()
         writer = self.create_writer(3, b'abc')
         self.assertEqual(writer.get_size(), 3)
         self.assertEqual(writer.get_allocated(), self.SMALL_BUFFER)
@@ -316,21 +314,19 @@ def test_abc(self):
         self.assertEqual(writer.get_allocated(), self.SMALL_BUFFER)
         self.assertEqual(writer.finish_with_size(3), b'abc')
 
-    # def test_write_bytes(self):
-    #     # Test PyBytesWriter_WriteBytes()
+    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!')
 
-    #     writer = self.create_writer()
-    #     writer.write_bytes(b'Hello World!', -1)
-    #     self.assertEqual(writer.finish(), 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!')
+         writer = self.create_writer()
+         writer.write_bytes(b'Hello ', -1)
+         writer.write_bytes(b'World! <truncated>', 6)
+         self.assertEqual(writer.finish(), b'Hello World!')
 
     def test_resize(self):
-        # Test PyBytesWriter_Extend()
-
+        # Test PyBytesWriter_Resize()
         writer = self.create_writer()
         writer.resize(len(b'number=123456'), b'number=123456')
         writer.resize(len(b'number=123456'), b'')

Modules/_testcapi/bytes.c

@@ -130,6 +130,27 @@ writer_check(WriterObject *self)
 }
 
 
+static PyObject*
+writer_write_bytes(PyObject *self_raw, PyObject *args)
+{
+    WriterObject *self = (WriterObject *)self_raw;
+    if (writer_check(self) < 0) {
+        return NULL;
+    }
+
+    char *bytes;
+    Py_ssize_t size;
+    if (!PyArg_ParseTuple(args, "yn", &bytes, &size)) {
+        return NULL;
+    }
+
+    if (PyBytesWriter_WriteBytes(self->writer, bytes, size) < 0) {
+        return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
+
 static PyObject*
 writer_resize(PyObject *self_raw, PyObject *args)
 {
@@ -220,6 +241,7 @@ writer_finish_with_size(PyObject *self_raw, PyObject *args)
 
 
 static PyMethodDef writer_methods[] = {
+    {"write_bytes", _PyCFunction_CAST(writer_write_bytes), METH_VARARGS},
     {"resize", _PyCFunction_CAST(writer_resize), METH_VARARGS},
     {"get_size", _PyCFunction_CAST(writer_get_size), METH_NOARGS},
     {"get_allocated", _PyCFunction_CAST(writer_get_allocated), METH_NOARGS},

Objects/bytesobject.c

@@ -3925,3 +3925,27 @@ PyBytesWriter_Resize(PyBytesWriter *writer, Py_ssize_t size)
     writer->size = size;
     return 0;
 }
+
+
+int
+PyBytesWriter_WriteBytes(PyBytesWriter *writer,
+                         const void *bytes, Py_ssize_t size)
+{
+    if (size < 0) {
+        size = strlen(bytes);
+    }
+
+    Py_ssize_t pos = writer->size;
+    if (size > PY_SSIZE_T_MAX - pos) {
+        PyErr_NoMemory();
+        return -1;
+    }
+    Py_ssize_t total = pos + size;
+
+    if (PyBytesWriter_Resize(writer, total) < 0) {
+        return -1;
+    }
+    char *buf = byteswriter_data(writer);
+    memcpy(buf + pos, bytes, size);
+    return 0;
+}