gh-121710: Add PyBytesWriter API

python · vstinner · Jul 13, 2024 · Jul 13, 2024 · Jul 14, 2024 · Jul 14, 2024
commit fe9667d149ae1fc1380683b7496a115efea81e76
@@ -201,3 +201,52 @@ called with a non-bytes parameter.
    reallocation fails, the original bytes object at *\*bytes* is deallocated,
    *\*bytes* is set to ``NULL``, :exc:`MemoryError` is set, and ``-1`` is
    returned.
+
+PyBytesWriter
+^^^^^^^^^^^^^
+
+The :c:type:`PyBytesWriter` API can be used to create a Python :class:`bytes`
+object.
+
+.. versionadded:: 3.14
+
+.. c:type:: PyBytesWriter
+
+   A bytes writer instance.
+
+   The instance must be destroyed by :c:func:`PyBytesWriter_Finish` on
+   success, or :c:func:`PyBytesWriter_Discard` on error.
+
+.. c:function:: PyBytesWriter* PyBytesWriter_Create(Py_ssize_t size, char **str)
+
+   Create a bytes writer instance.
+   Preallocate *size* bytes.
+
+   On success, set *\*str* and return a new writer.
+   On error, set an exception and return ``NULL``.
+
+.. c:function:: PyObject* PyBytesWriter_Finish(PyBytesWriter *writer, char *str)
+
+   Return the final Python :class:`bytes` object and destroy the writer
+   instance.
+
+   On success, return a bytes object.
+   On error, set an exception and return ``NULL``.
+
+.. c:function:: void PyBytesWriter_Discard(PyBytesWriter *writer)
+
+   Discard the internal bytes buffer and destroy the writer instance.
+
+.. c:function:: int PyBytesWriter_Prepare(PyBytesWriter *writer, char **str, Py_ssize_t size)
+
+   Allocate *size* bytes to prepare writing *size* bytes into *writer*.
+
+   On success, update *\*str* and return ``0``.
+   On error, set an exception and return ``-1``.
+
+.. c:function:: int PyBytesWriter_WriteBytes(PyBytesWriter *writer, char **str, const void *bytes, Py_ssize_t size)
+
+   Write a the bytes string *bytes* of *size* bytes into *writer*.
+
+   On success, update *\*str* and return ``0``.
+   On error, set an exception and return ``-1``.
@@ -1546,6 +1546,7 @@ object.
 .. c:function:: PyUnicodeWriter* PyUnicodeWriter_Create(Py_ssize_t length)
 
    Create a Unicode writer instance.
+   Preallocate *length* characters.
 
    Set an exception and return ``NULL`` on error.
 

@@ -405,6 +405,17 @@ New Features
 
   (Contributed by Victor Stinner in :gh:`119182`.)
 
+* Add a new :c:type:`PyBytesWriter` API to create a Python :class:`bytes`
+  object:
+
+  * :c:func:`PyBytesWriter_Create`;
+  * :c:func:`PyBytesWriter_Finish`;
+  * :c:func:`PyBytesWriter_Discard`;
+  * :c:func:`PyBytesWriter_Prepare`;
+  * :c:func:`PyBytesWriter_WriteBytes`.
+
+  (Contributed by Victor Stinner in :gh:`121710`.)
+
 Porting to Python 3.14
 ----------------------
 

diff --git a/Include/cpython/bytesobject.h b/Include/cpython/bytesobject.h
@@ -35,3 +35,25 @@ static inline Py_ssize_t PyBytes_GET_SIZE(PyObject *op) {
 /* _PyBytes_Join(sep, x) is like sep.join(x).  sep must be PyBytesObject*,
    x must be an iterable object. */
 PyAPI_FUNC(PyObject*) _PyBytes_Join(PyObject *sep, PyObject *x);
+
+/* --- PyBytesWriter ------------------------------------------------------ */
+
+typedef struct PyBytesWriter PyBytesWriter;
+
+PyAPI_FUNC(PyBytesWriter*) PyBytesWriter_Create(
+    Py_ssize_t size,
+    char **str);
+PyAPI_FUNC(PyObject *) PyBytesWriter_Finish(
+    PyBytesWriter *writer,
+    char *str);
+PyAPI_FUNC(void) PyBytesWriter_Discard(PyBytesWriter *writer);
+
+PyAPI_FUNC(int) PyBytesWriter_Prepare(
+    PyBytesWriter *writer,
+    char **str,
+    Py_ssize_t size);
+PyAPI_FUNC(int) PyBytesWriter_WriteBytes(
+    PyBytesWriter *writer,
+    char **str,
+    const void *bytes,
+    Py_ssize_t size);
diff --git a/Misc/NEWS.d/next/C_API/2024-07-13-21-49-12.gh-issue-121710.j-9Vhk.rst b/Misc/NEWS.d/next/C_API/2024-07-13-21-49-12.gh-issue-121710.j-9Vhk.rst
@@ -0,0 +1,10 @@
+Add a new :c:type:`PyBytesWriter` API to create a Python :class:`bytes`
+object:
+
+* :c:func:`PyBytesWriter_Create`;
+* :c:func:`PyBytesWriter_Finish`;
+* :c:func:`PyBytesWriter_Discard`;
+* :c:func:`PyBytesWriter_Prepare`;
+* :c:func:`PyBytesWriter_WriteBytes`.
+
+Patch by Victor Stinner.
diff --git a/Modules/_testcapi/bytes.c b/Modules/_testcapi/bytes.c
@@ -37,8 +37,140 @@ bytes_resize(PyObject *Py_UNUSED(module), PyObject *args)
 }
 
 
+static int
+bytes_equal(PyObject *obj, const char *str)
+{
+    return (PyBytes_Size(obj) == (Py_ssize_t)strlen(str)
+            && strcmp(PyBytes_AsString(obj), str) == 0);
+}
+
+
+/* Test PyBytesWriter API */
+static PyObject *
+test_byteswriter(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
+{
+    char *str;
+    PyBytesWriter *writer = PyBytesWriter_Create(3, &str);
+    if (writer == NULL) {
+        return NULL;
+    }
+
+    if (PyBytesWriter_WriteBytes(writer, &str, "abc", 3) < 0) {
+        goto error;
+    }
+
+    // write empty string
+    if (PyBytesWriter_WriteBytes(writer, &str, "", 0) < 0) {
+        goto error;
+    }
+
+    PyObject *obj = PyBytesWriter_Finish(writer, str);
+    if (obj == NULL) {
+        return NULL;
+    }
+
+    assert(bytes_equal(obj, "abc"));
+    Py_DECREF(obj);
+
+    Py_RETURN_NONE;
+
+error:
+    PyBytesWriter_Discard(writer);
+    return NULL;
+}
+
+
+/* Test PyBytesWriter_Discard() */
+static PyObject *
+test_byteswriter_discard(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
+{
+    char *str;
+    PyBytesWriter *writer = PyBytesWriter_Create(3, &str);
+    if (writer == NULL) {
+        return NULL;
+    }
+    assert(PyBytesWriter_WriteBytes(writer, &str, "abc", 3) == 0);
+
+    PyBytesWriter_Discard(writer);
+    Py_RETURN_NONE;
+}
+
+
+/* Test PyBytesWriter_WriteBytes() */
+static PyObject *
+test_byteswriter_writebytes(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
+{
+    char *str;
+    PyBytesWriter *writer = PyBytesWriter_Create(0, &str);
+    if (writer == NULL) {
+        return NULL;
+    }
+
+    if (PyBytesWriter_WriteBytes(writer, &str, "abc", 3) < 0) {
+        goto error;
+    }
+    if (PyBytesWriter_WriteBytes(writer, &str, "def", 3) < 0) {
+        goto error;
+    }
+
+    PyObject *obj = PyBytesWriter_Finish(writer, str);
+    if (obj == NULL) {
+        return NULL;
+    }
+
+    assert(bytes_equal(obj, "abcdef"));
+    Py_DECREF(obj);
+
+    Py_RETURN_NONE;
+
+error:
+    PyBytesWriter_Discard(writer);
+    return NULL;
+}
+
+
+/* Test PyBytesWriter_Prepare() */
+static PyObject *
+test_byteswriter_prepare(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
+{
+    char *str;
+    PyBytesWriter *writer = PyBytesWriter_Create(0, &str);
+    if (writer == NULL) {
+        return NULL;
+    }
+
+    // test error on purpose (negative size)
+    assert(PyBytesWriter_Prepare(writer, &str, -3) < 0);
+    assert(PyErr_ExceptionMatches(PyExc_ValueError));
+    PyErr_Clear();
+
+    if (PyBytesWriter_Prepare(writer, &str, 3) < 0) {
+        PyBytesWriter_Discard(writer);
+        return NULL;
+    }
+
+    // Write "abc"
+    memcpy(str, "abc", 3);
+    str += 3;
+
+    PyObject *obj = PyBytesWriter_Finish(writer, str);
+    if (obj == NULL) {
+        return NULL;
+    }
+
+    assert(bytes_equal(obj, "abc"));
+    Py_DECREF(obj);
+
+    Py_RETURN_NONE;
+}
+
+
 static PyMethodDef test_methods[] = {
     {"bytes_resize", bytes_resize, METH_VARARGS},
+    {"test_byteswriter", test_byteswriter, METH_NOARGS},
+    {"test_byteswriter_discard", test_byteswriter_discard, METH_NOARGS},
+    {"test_byteswriter_writebytes", test_byteswriter_writebytes, METH_NOARGS},
+    {"test_byteswriter_prepare", test_byteswriter_prepare, METH_NOARGS},
     {NULL},
 };