[WIP] gh-129813, PEP 782: Add PyBytesWriter C API

Include/cpython/bytesobject.h

@@ -40,3 +40,46 @@ _PyBytes_Join(PyObject *sep, PyObject *iterable)
 {
     return PyBytes_Join(sep, iterable);
 }
+
+
+// --- PyBytesWriter API -----------------------------------------------------
+
+typedef struct PyBytesWriter PyBytesWriter;
+
+PyAPI_FUNC(PyBytesWriter *) PyBytesWriter_Create(
+    Py_ssize_t size);
+PyAPI_FUNC(void) PyBytesWriter_Discard(
+    PyBytesWriter *writer);
+PyAPI_FUNC(PyObject*) PyBytesWriter_Finish(
+    PyBytesWriter *writer);
+PyAPI_FUNC(PyObject*) PyBytesWriter_FinishWithSize(
+    PyBytesWriter *writer,
+    Py_ssize_t size);
+PyAPI_FUNC(PyObject*) PyBytesWriter_FinishWithPointer(
+    PyBytesWriter *writer,
+    void *buf);
+
+PyAPI_FUNC(void*) PyBytesWriter_GetData(
+    PyBytesWriter *writer);
+PyAPI_FUNC(Py_ssize_t) PyBytesWriter_GetSize(
+    PyBytesWriter *writer);
+
+PyAPI_FUNC(int) PyBytesWriter_WriteBytes(
+    PyBytesWriter *writer,
+    const void *bytes,
+    Py_ssize_t size);
+PyAPI_FUNC(int) PyBytesWriter_Format(
+    PyBytesWriter *writer,
+    const char *format,
+    ...);
+
+PyAPI_FUNC(int) PyBytesWriter_Resize(
+    PyBytesWriter *writer,
+    Py_ssize_t size);
+PyAPI_FUNC(int) PyBytesWriter_Grow(
+    PyBytesWriter *writer,
+    Py_ssize_t size);
+PyAPI_FUNC(void*) PyBytesWriter_GrowAndUpdatePointer(
+    PyBytesWriter *writer,
+    Py_ssize_t size,
+    void *buf);

Include/internal/pycore_bytesobject.h

@@ -142,6 +142,10 @@ PyAPI_FUNC(void*) _PyBytesWriter_WriteBytes(_PyBytesWriter *writer,
     const void *bytes,
     Py_ssize_t size);
 
+// Export for '_testcapi' shared extension.
+PyAPI_FUNC(PyBytesWriter*) _PyBytesWriter_CreateByteArray(
+    Py_ssize_t size);
+
 #ifdef __cplusplus
 }
 #endif

Include/internal/pycore_freelist_state.h

@@ -26,6 +26,7 @@ extern "C" {
 #  define Py_futureiters_MAXFREELIST 255
 #  define Py_object_stack_chunks_MAXFREELIST 4
 #  define Py_unicode_writers_MAXFREELIST 1
+#  define Py_bytes_writers_MAXFREELIST 1
 #  define Py_pycfunctionobject_MAXFREELIST 16
 #  define Py_pycmethodobject_MAXFREELIST 16
 #  define Py_pymethodobjects_MAXFREELIST 20
@@ -59,6 +60,7 @@ struct _Py_freelists {
     struct _Py_freelist futureiters;
     struct _Py_freelist object_stack_chunks;
     struct _Py_freelist unicode_writers;
+    struct _Py_freelist bytes_writers;
     struct _Py_freelist pycfunctionobject;
     struct _Py_freelist pycmethodobject;
     struct _Py_freelist pymethodobjects;

Include/internal/pycore_long.h

@@ -135,7 +135,7 @@ extern int _PyLong_FormatWriter(
     int alternate);
 
 extern char* _PyLong_FormatBytesWriter(
-    _PyBytesWriter *writer,
+    PyBytesWriter *writer,
     char *str,
     PyObject *obj,
     int base,

Lib/test/test_capi/test_bytes.py

@@ -291,5 +291,94 @@ def test_join(self):
             bytes_join(b'', NULL)
 
 
+class BytesWriterTest(unittest.TestCase):
+    result_type = bytes
+
+    def create_writer(self, alloc=0, string=b''):
+        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.finish(), self.result_type(b''))
+
+        writer = self.create_writer(3, b'abc')
+        self.assertEqual(writer.get_size(), 3)
+        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.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(), 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(), 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(), 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(), 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(), 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(), 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(), 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(), 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(), self.result_type(b'x=123, y=456'))
+
+    def test_example_abc(self):
+        self.assertEqual(_testcapi.byteswriter_abc(), b'abc')
+
+    def test_example_resize(self):
+        self.assertEqual(_testcapi.byteswriter_resize(), b'Hello World')
+
+    def test_example_highlevel(self):
+        self.assertEqual(_testcapi.byteswriter_highlevel(), b'Hello World!')
+
+
+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()

Modules/_bz2module.c

@@ -668,9 +668,7 @@ _bz2_BZ2Decompressor_impl(PyTypeObject *type)
     self->bzs_avail_in_real = 0;
     self->input_buffer = NULL;
     self->input_buffer_size = 0;
-    self->unused_data = PyBytes_FromStringAndSize(NULL, 0);
-    if (self->unused_data == NULL)
-        goto error;
+    self->unused_data = Py_GetConstant(Py_CONSTANT_EMPTY_BYTES);
 
     bzerror = BZ2_bzDecompressInit(&self->bzs, 0, 0);
     if (catch_bz2_error(bzerror))

Modules/_codecsmodule.c

@@ -201,52 +201,45 @@ _codecs_escape_encode_impl(PyObject *module, PyObject *data,
                            const char *errors)
 /*[clinic end generated code: output=4af1d477834bab34 input=8f4b144799a94245]*/
 {
-    Py_ssize_t size;
-    Py_ssize_t newsize;
-    PyObject *v;
-
-    size = PyBytes_GET_SIZE(data);
+    Py_ssize_t size = PyBytes_GET_SIZE(data);
     if (size > PY_SSIZE_T_MAX / 4) {
         PyErr_SetString(PyExc_OverflowError,
             "string is too large to encode");
             return NULL;
     }
-    newsize = 4*size;
-    v = PyBytes_FromStringAndSize(NULL, newsize);
+    Py_ssize_t alloc_size = 4*size;
 
-    if (v == NULL) {
+    PyBytesWriter *writer = PyBytesWriter_Create(alloc_size);
+    if (writer == NULL) {
         return NULL;
     }
-    else {
-        Py_ssize_t i;
-        char c;
-        char *p = PyBytes_AS_STRING(v);
-
-        for (i = 0; i < size; i++) {
-            /* There's at least enough room for a hex escape */
-            assert(newsize - (p - PyBytes_AS_STRING(v)) >= 4);
-            c = PyBytes_AS_STRING(data)[i];
-            if (c == '\'' || c == '\\')
-                *p++ = '\\', *p++ = c;
-            else if (c == '\t')
-                *p++ = '\\', *p++ = 't';
-            else if (c == '\n')
-                *p++ = '\\', *p++ = 'n';
-            else if (c == '\r')
-                *p++ = '\\', *p++ = 'r';
-            else if (c < ' ' || c >= 0x7f) {
-                *p++ = '\\';
-                *p++ = 'x';
-                *p++ = Py_hexdigits[(c & 0xf0) >> 4];
-                *p++ = Py_hexdigits[c & 0xf];
-            }
-            else
-                *p++ = c;
-        }
-        *p = '\0';
-        if (_PyBytes_Resize(&v, (p - PyBytes_AS_STRING(v)))) {
-            return NULL;
+    char *p = PyBytesWriter_GetData(writer);
+
+    for (Py_ssize_t i = 0; i < size; i++) {
+        /* There's at least enough room for a hex escape */
+        assert(alloc_size - (p - (char*)PyBytesWriter_GetData(writer)) >= 4);
+        char c = PyBytes_AS_STRING(data)[i];
+        if (c == '\'' || c == '\\')
+            *p++ = '\\', *p++ = c;
+        else if (c == '\t')
+            *p++ = '\\', *p++ = 't';
+        else if (c == '\n')
+            *p++ = '\\', *p++ = 'n';
+        else if (c == '\r')
+            *p++ = '\\', *p++ = 'r';
+        else if (c < ' ' || c >= 0x7f) {
+            *p++ = '\\';
+            *p++ = 'x';
+            *p++ = Py_hexdigits[(c & 0xf0) >> 4];
+            *p++ = Py_hexdigits[c & 0xf];
         }
+        else
+            *p++ = c;
+    }
+
+    PyObject *v = PyBytesWriter_FinishWithPointer(writer, p);
+    if (v == NULL) {
+        return NULL;
     }
 
     return codec_tuple(v, size);

Modules/_dbmmodule.c

@@ -401,10 +401,7 @@ _dbm_dbm_setdefault_impl(dbmobject *self, PyTypeObject *cls, const char *key,
         return PyBytes_FromStringAndSize(val.dptr, val.dsize);
     }
     if (default_value == NULL) {
-        default_value = PyBytes_FromStringAndSize(NULL, 0);
-        if (default_value == NULL) {
-            return NULL;
-        }
+        default_value = Py_GetConstant(Py_CONSTANT_EMPTY_BYTES);
         val.dptr = NULL;
         val.dsize = 0;
     }

Modules/_hashopenssl.c

@@ -806,15 +806,15 @@ EVPXOF_digest_impl(EVPobject *self, Py_ssize_t length)
 /*[clinic end generated code: output=ef9320c23280efad input=816a6537cea3d1db]*/
 {
     EVP_MD_CTX *temp_ctx;
-    PyObject *retval = PyBytes_FromStringAndSize(NULL, length);
 
-    if (retval == NULL) {
+    PyBytesWriter *writer = PyBytesWriter_Create(length);
+    if (writer == NULL) {
         return NULL;
     }
 
     temp_ctx = EVP_MD_CTX_new();
     if (temp_ctx == NULL) {
-        Py_DECREF(retval);
+        PyBytesWriter_Discard(writer);
         PyErr_NoMemory();
         return NULL;
     }
@@ -823,17 +823,17 @@ EVPXOF_digest_impl(EVPobject *self, Py_ssize_t length)
         goto error;
     }
     if (!EVP_DigestFinalXOF(temp_ctx,
-                            (unsigned char*)PyBytes_AS_STRING(retval),
+                            (unsigned char*)PyBytesWriter_GetData(writer),
                             length))
     {
         goto error;
     }
 
     EVP_MD_CTX_free(temp_ctx);
-    return retval;
+    return PyBytesWriter_Finish(writer);
 
 error:
-    Py_DECREF(retval);
+    PyBytesWriter_Discard(writer);
     EVP_MD_CTX_free(temp_ctx);
     notify_ssl_error_occurred();
     return NULL;
@@ -1414,8 +1414,6 @@ _hashlib_scrypt_impl(PyObject *module, Py_buffer *password, Py_buffer *salt,
                      long maxmem, long dklen)
 /*[clinic end generated code: output=14849e2aa2b7b46c input=48a7d63bf3f75c42]*/
 {
-    PyObject *key_obj = NULL;
-    char *key;
     int retval;
     unsigned long n, r, p;
 
@@ -1486,27 +1484,27 @@ _hashlib_scrypt_impl(PyObject *module, Py_buffer *password, Py_buffer *salt,
         return NULL;
    }
 
-    key_obj = PyBytes_FromStringAndSize(NULL, dklen);
-    if (key_obj == NULL) {
+    PyBytesWriter *writer = PyBytesWriter_Create(dklen);
+    if (writer == NULL) {
         return NULL;
     }
-    key = PyBytes_AS_STRING(key_obj);
+    unsigned char *key = PyBytesWriter_GetData(writer);
 
     Py_BEGIN_ALLOW_THREADS
     retval = EVP_PBE_scrypt(
         (const char*)password->buf, (size_t)password->len,
         (const unsigned char *)salt->buf, (size_t)salt->len,
         n, r, p, maxmem,
-        (unsigned char *)key, (size_t)dklen
+        key, (size_t)dklen
     );
     Py_END_ALLOW_THREADS
 
     if (!retval) {
-        Py_CLEAR(key_obj);
+        PyBytesWriter_Discard(writer);
         notify_ssl_error_occurred();
         return NULL;
     }
-    return key_obj;
+    return PyBytesWriter_Finish(writer);
 }
 #endif  /* PY_OPENSSL_HAS_SCRYPT */