diff --git a/Doc/c-api/bytes.rst b/Doc/c-api/bytes.rst index d47beee68eaa33..0a3cbaa62712f8 100644 --- a/Doc/c-api/bytes.rst +++ b/Doc/c-api/bytes.rst @@ -219,3 +219,153 @@ 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:: next + +.. c:type:: PyBytesWriter + + 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 instance must be destroyed by :c:func:`PyBytesWriter_Finish` on + success, or :c:func:`PyBytesWriter_Discard` on error. + + +Create, Finish, Discard +^^^^^^^^^^^^^^^^^^^^^^^ + +.. c:function:: PyBytesWriter* PyBytesWriter_Create(Py_ssize_t size) + + Create a :c:type:`PyBytesWriter` to write *size* bytes. + + If *size* is greater than zero, allocate *size* bytes, and set the + writer size to *size*. The caller is responsible to write *size* + bytes using :c:func:`PyBytesWriter_GetData`. + + On error, set an exception and return NULL. + + *size* must be positive or zero. + +.. c:function:: PyObject* PyBytesWriter_Finish(PyBytesWriter *writer) + + Finish a :c:type:`PyBytesWriter` created by + :c:func:`PyBytesWriter_Create`. + + On success, return a Python :class:`bytes` object. + On error, set an exception and return ``NULL``. + + The writer instance is invalid after the call in any case. + No API can be called on the writer after :c:func:`PyBytesWriter_Finish`. + +.. c:function:: PyObject* PyBytesWriter_FinishWithSize(PyBytesWriter *writer, Py_ssize_t size) + + Similar to :c:func:`PyBytesWriter_Finish`, but resize the writer + to *size* bytes before creating the :class:`bytes` object. + +.. c:function:: PyObject* PyBytesWriter_FinishWithPointer(PyBytesWriter *writer, void *buf) + + Similar to :c:func:`PyBytesWriter_Finish`, but resize the writer + using *buf* pointer before creating the :class:`bytes` object. + + Set an exception and return ``NULL`` if *buf* pointer is outside the + internal buffer bounds. + + Function pseudo-code:: + + Py_ssize_t size = (char*)buf - (char*)PyBytesWriter_GetData(writer); + return PyBytesWriter_FinishWithSize(writer, size); + +.. c:function:: void PyBytesWriter_Discard(PyBytesWriter *writer) + + Discard a :c:type:`PyBytesWriter` created by :c:func:`PyBytesWriter_Create`. + + Do nothing if *writer* is ``NULL``. + + The writer instance is invalid after the call. + No API can be called on the writer after :c:func:`PyBytesWriter_Discard`. + +High-level API +^^^^^^^^^^^^^^ + +.. c:function:: int PyBytesWriter_WriteBytes(PyBytesWriter *writer, const void *bytes, Py_ssize_t size) + + Grow the *writer* internal buffer by *size* bytes, + write *size* bytes of *bytes* at the *writer* end, + and add *size* to the *writer* size. + + If *size* is equal to ``-1``, call ``strlen(bytes)`` to get the + string length. + + On success, return ``0``. + On error, set an exception and return ``-1``. + + +Getters +^^^^^^^ + +.. c:function:: Py_ssize_t PyBytesWriter_GetSize(PyBytesWriter *writer) + + Get the writer size. + +.. c:function:: void* PyBytesWriter_GetData(PyBytesWriter *writer) + + Get the writer data: start of the internal buffer. + + The pointer is valid until :c:func:`PyBytesWriter_Finish` or + :c:func:`PyBytesWriter_Discard` is called on *writer*. + + +Low-level API +^^^^^^^^^^^^^ + +.. c:function:: int PyBytesWriter_Resize(PyBytesWriter *writer, Py_ssize_t size) + + Resize the writer to *size* bytes. It can be used to enlarge or to + shrink the writer. + + Newly allocated bytes are left uninitialized. + + On success, return ``0``. + On error, set an exception and return ``-1``. + + *size* must be positive or zero. + +.. c:function:: int PyBytesWriter_Grow(PyBytesWriter *writer, Py_ssize_t grow) + + Resize the writer by adding *grow* bytes to the current writer size. + + Newly allocated bytes are left uninitialized. + + On success, return ``0``. + On error, set an exception and return ``-1``. + + *size* can be negative to shrink the writer. + +.. c:function:: void* PyBytesWriter_GrowAndUpdatePointer(PyBytesWriter *writer, Py_ssize_t size, void *buf) + + Similar to :c:func:`PyBytesWriter_Grow`, but update also the *buf* + pointer. + + The *buf* pointer is moved if the internal buffer is moved in memory. + The *buf* relative position within the internal buffer is left + unchanged. + + On error, set an exception and return ``NULL``. + + *buf* must not be ``NULL``. + + Function pseudo-code:: + + Py_ssize_t pos = (char*)buf - (char*)PyBytesWriter_GetData(writer); + if (PyBytesWriter_Grow(writer, size) < 0) { + return NULL; + } + return (char*)PyBytesWriter_GetData(writer) + pos; diff --git a/Doc/whatsnew/3.15.rst b/Doc/whatsnew/3.15.rst index 56ef80c068634a..02d5d4e6d5194d 100644 --- a/Doc/whatsnew/3.15.rst +++ b/Doc/whatsnew/3.15.rst @@ -701,6 +701,22 @@ New features and :c:data:`Py_mod_abi`. (Contributed by Petr Viktorin in :gh:`137210`.) +* Implement :pep:`782`, the :c:type:`PyBytesWriter` API. Add functions: + + * :c:func:`PyBytesWriter_Create` + * :c:func:`PyBytesWriter_Discard` + * :c:func:`PyBytesWriter_FinishWithPointer` + * :c:func:`PyBytesWriter_FinishWithSize` + * :c:func:`PyBytesWriter_Finish` + * :c:func:`PyBytesWriter_GetData` + * :c:func:`PyBytesWriter_GetSize` + * :c:func:`PyBytesWriter_GrowAndUpdatePointer` + * :c:func:`PyBytesWriter_Grow` + * :c:func:`PyBytesWriter_Resize` + * :c:func:`PyBytesWriter_WriteBytes` + + (Contributed by Victor Stinner in :gh:`129813`.) + Porting to Python 3.15 ---------------------- diff --git a/Include/cpython/bytesobject.h b/Include/cpython/bytesobject.h index 71c133f173f157..23a652cedf32cd 100644 --- a/Include/cpython/bytesobject.h +++ b/Include/cpython/bytesobject.h @@ -40,3 +40,42 @@ _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_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); diff --git a/Include/internal/pycore_bytesobject.h b/Include/internal/pycore_bytesobject.h index 8ea9b3ebb88454..9f519d3ca95e92 100644 --- a/Include/internal/pycore_bytesobject.h +++ b/Include/internal/pycore_bytesobject.h @@ -143,6 +143,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 diff --git a/Include/internal/pycore_freelist_state.h b/Include/internal/pycore_freelist_state.h index 59beb92f3f7b9c..46e2a82ea03456 100644 --- a/Include/internal/pycore_freelist_state.h +++ b/Include/internal/pycore_freelist_state.h @@ -27,6 +27,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 @@ -61,6 +62,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; diff --git a/Lib/test/test_capi/test_bytes.py b/Lib/test/test_capi/test_bytes.py index bc820bd68d9e21..cc6c932c7d9072 100644 --- a/Lib/test/test_capi/test_bytes.py +++ b/Lib/test/test_capi/test_bytes.py @@ -299,5 +299,80 @@ 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! ', 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_example_abc(self): + self.assertEqual(_testcapi.byteswriter_abc(), b'abc') + + def test_example_resize(self): + self.assertEqual(_testcapi.byteswriter_resize(), 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() diff --git a/Misc/NEWS.d/next/C_API/2025-09-12-13-05-20.gh-issue-129813.dJZpME.rst b/Misc/NEWS.d/next/C_API/2025-09-12-13-05-20.gh-issue-129813.dJZpME.rst new file mode 100644 index 00000000000000..30d07279918045 --- /dev/null +++ b/Misc/NEWS.d/next/C_API/2025-09-12-13-05-20.gh-issue-129813.dJZpME.rst @@ -0,0 +1,15 @@ +Implement :pep:`782`, the :c:type:`PyBytesWriter` API. Add functions: + +* :c:func:`PyBytesWriter_Create` +* :c:func:`PyBytesWriter_Discard` +* :c:func:`PyBytesWriter_FinishWithPointer` +* :c:func:`PyBytesWriter_FinishWithSize` +* :c:func:`PyBytesWriter_Finish` +* :c:func:`PyBytesWriter_GetData` +* :c:func:`PyBytesWriter_GetSize` +* :c:func:`PyBytesWriter_GrowAndUpdatePointer` +* :c:func:`PyBytesWriter_Grow` +* :c:func:`PyBytesWriter_Resize` +* :c:func:`PyBytesWriter_WriteBytes` + +Patch by Victor Stinner. diff --git a/Modules/_testcapi/bytes.c b/Modules/_testcapi/bytes.c index 33903de14ba68d..3530b6a4a42b44 100644 --- a/Modules/_testcapi/bytes.c +++ b/Modules/_testcapi/bytes.c @@ -1,6 +1,11 @@ +// Use pycore_bytes.h +#define PYTESTCAPI_NEED_INTERNAL_API + #include "parts.h" #include "util.h" +#include "pycore_bytesobject.h" // _PyBytesWriter_CreateByteArray() + /* Test _PyBytes_Resize() */ static PyObject * @@ -51,9 +56,264 @@ bytes_join(PyObject *Py_UNUSED(module), PyObject *args) } +// --- PyBytesWriter type --------------------------------------------------- + +typedef struct { + PyObject_HEAD + PyBytesWriter *writer; +} WriterObject; + + +static PyObject * +writer_new(PyTypeObject *type, PyObject *args, PyObject *kwargs) +{ + WriterObject *self = (WriterObject *)type->tp_alloc(type, 0); + if (!self) { + return NULL; + } + self->writer = NULL; + return (PyObject*)self; +} + + +static int +writer_init(PyObject *self_raw, PyObject *args, PyObject *kwargs) +{ + WriterObject *self = (WriterObject *)self_raw; + if (self->writer) { + PyBytesWriter_Discard(self->writer); + } + + if (kwargs && PyDict_GET_SIZE(kwargs)) { + PyErr_Format(PyExc_TypeError, + "PyBytesWriter() takes exactly no keyword arguments"); + return -1; + } + + Py_ssize_t alloc; + char *str; + Py_ssize_t str_size; + int use_bytearray; + if (!PyArg_ParseTuple(args, "ny#i", + &alloc, &str, &str_size, &use_bytearray)) { + return -1; + } + + if (use_bytearray) { + self->writer = _PyBytesWriter_CreateByteArray(alloc); + } + else { + self->writer = PyBytesWriter_Create(alloc); + } + if (self->writer == NULL) { + return -1; + } + + if (str_size) { + char *buf = PyBytesWriter_GetData(self->writer); + memcpy(buf, str, str_size); + } + + return 0; +} + + +static void +writer_dealloc(PyObject *self_raw) +{ + WriterObject *self = (WriterObject *)self_raw; + PyTypeObject *tp = Py_TYPE(self); + if (self->writer) { + PyBytesWriter_Discard(self->writer); + } + tp->tp_free(self); + Py_DECREF(tp); +} + + +static inline int +writer_check(WriterObject *self) +{ + if (self->writer == NULL) { + PyErr_SetString(PyExc_ValueError, "operation on finished writer"); + return -1; + } + return 0; +} + + +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) +{ + WriterObject *self = (WriterObject *)self_raw; + if (writer_check(self) < 0) { + return NULL; + } + + Py_ssize_t size; + char *str; + Py_ssize_t str_size; + if (!PyArg_ParseTuple(args, + "ny#", + &size, &str, &str_size)) { + return NULL; + } + assert(size >= str_size); + + Py_ssize_t pos = PyBytesWriter_GetSize(self->writer); + if (PyBytesWriter_Resize(self->writer, size) < 0) { + return NULL; + } + + char *buf = PyBytesWriter_GetData(self->writer); + memcpy(buf + pos, str, str_size); + + Py_RETURN_NONE; +} + + +static PyObject* +writer_get_size(PyObject *self_raw, PyObject *Py_UNUSED(args)) +{ + WriterObject *self = (WriterObject *)self_raw; + if (writer_check(self) < 0) { + return NULL; + } + + Py_ssize_t alloc = PyBytesWriter_GetSize(self->writer); + return PyLong_FromSsize_t(alloc); +} + + +static PyObject* +writer_finish(PyObject *self_raw, PyObject *Py_UNUSED(args)) +{ + WriterObject *self = (WriterObject *)self_raw; + if (writer_check(self) < 0) { + return NULL; + } + + PyObject *str = PyBytesWriter_Finish(self->writer); + self->writer = NULL; + return str; +} + + +static PyObject* +writer_finish_with_size(PyObject *self_raw, PyObject *args) +{ + WriterObject *self = (WriterObject *)self_raw; + if (writer_check(self) < 0) { + return NULL; + } + + Py_ssize_t size; + if (!PyArg_ParseTuple(args, "n", &size)) { + return NULL; + } + + PyObject *str = PyBytesWriter_FinishWithSize(self->writer, size); + self->writer = NULL; + return str; +} + + +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}, + {"finish", _PyCFunction_CAST(writer_finish), METH_NOARGS}, + {"finish_with_size", _PyCFunction_CAST(writer_finish_with_size), METH_VARARGS}, + {NULL, NULL} /* sentinel */ +}; + +static PyType_Slot Writer_Type_slots[] = { + {Py_tp_new, writer_new}, + {Py_tp_init, writer_init}, + {Py_tp_dealloc, writer_dealloc}, + {Py_tp_methods, writer_methods}, + {0, 0}, /* sentinel */ +}; + +static PyType_Spec Writer_spec = { + .name = "_testcapi.PyBytesWriter", + .basicsize = sizeof(WriterObject), + .flags = Py_TPFLAGS_DEFAULT, + .slots = Writer_Type_slots, +}; + + +static PyObject * +byteswriter_abc(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args)) +{ + PyBytesWriter *writer = PyBytesWriter_Create(3); + if (writer == NULL) { + return NULL; + } + + char *str = PyBytesWriter_GetData(writer); + memcpy(str, "abc", 3); + + return PyBytesWriter_Finish(writer); +} + + +static PyObject * +byteswriter_resize(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args)) +{ + // Allocate 10 bytes + PyBytesWriter *writer = PyBytesWriter_Create(10); + if (writer == NULL) { + return NULL; + } + char *buf = PyBytesWriter_GetData(writer); + + // Write some bytes + memcpy(buf, "Hello ", strlen("Hello ")); + buf += strlen("Hello "); + + // Allocate 10 more bytes + buf = PyBytesWriter_GrowAndUpdatePointer(writer, 10, buf); + if (buf == NULL) { + PyBytesWriter_Discard(writer); + return NULL; + } + + // Write more bytes + memcpy(buf, "World", strlen("World")); + buf += strlen("World"); + + // Truncate to the exact size and create a bytes object + return PyBytesWriter_FinishWithPointer(writer, buf); +} + + static PyMethodDef test_methods[] = { {"bytes_resize", bytes_resize, METH_VARARGS}, {"bytes_join", bytes_join, METH_VARARGS}, + {"byteswriter_abc", byteswriter_abc, METH_NOARGS}, + {"byteswriter_resize", byteswriter_resize, METH_NOARGS}, {NULL}, }; @@ -64,5 +324,15 @@ _PyTestCapi_Init_Bytes(PyObject *m) return -1; } + PyTypeObject *writer_type = (PyTypeObject *)PyType_FromSpec(&Writer_spec); + if (writer_type == NULL) { + return -1; + } + if (PyModule_AddType(m, writer_type) < 0) { + Py_DECREF(writer_type); + return -1; + } + Py_DECREF(writer_type); + return 0; } diff --git a/Objects/bytesobject.c b/Objects/bytesobject.c index db82f7eb684f30..aa099af8cf17f3 100644 --- a/Objects/bytesobject.c +++ b/Objects/bytesobject.c @@ -7,6 +7,7 @@ #include "pycore_call.h" // _PyObject_CallNoArgs() #include "pycore_ceval.h" // _PyEval_GetBuiltin() #include "pycore_format.h" // F_LJUST +#include "pycore_freelist.h" // _Py_FREELIST_FREE() #include "pycore_global_objects.h"// _Py_GET_GLOBAL_OBJECT() #include "pycore_initconfig.h" // _PyStatus_OK() #include "pycore_long.h" // _PyLong_DigitValue @@ -3747,3 +3748,303 @@ _PyBytes_Repeat(char* dest, Py_ssize_t len_dest, } } + +// --- PyBytesWriter API ----------------------------------------------------- + +struct PyBytesWriter { + char small_buffer[256]; + PyObject *obj; + Py_ssize_t size; + int use_bytearray; +}; + + +static inline char* +byteswriter_data(PyBytesWriter *writer) +{ + if (writer->obj == NULL) { + return writer->small_buffer; + } + else if (writer->use_bytearray) { + return PyByteArray_AS_STRING(writer->obj); + } + else { + return PyBytes_AS_STRING(writer->obj); + } +} + + +static inline Py_ssize_t +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); + } +} + + +#ifdef MS_WINDOWS + /* On Windows, overallocate by 50% is the best factor */ +# define OVERALLOCATE_FACTOR 2 +#else + /* On Linux, overallocate by 25% is the best factor */ +# define OVERALLOCATE_FACTOR 4 +#endif + + +static inline int +byteswriter_resize(PyBytesWriter *writer, Py_ssize_t size, int overallocate) +{ + assert(size >= 0); + + if (size <= byteswriter_allocated(writer)) { + return 0; + } + + if (overallocate && !writer->use_bytearray) { + if (size <= (PY_SSIZE_T_MAX - size / OVERALLOCATE_FACTOR)) { + size += size / OVERALLOCATE_FACTOR; + } + } + + if (writer->obj != NULL) { + 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) { + return -1; + } + assert((size_t)size > sizeof(writer->small_buffer)); + memcpy(PyBytes_AS_STRING(writer->obj), + writer->small_buffer, + sizeof(writer->small_buffer)); + } + return 0; +} + + +static PyBytesWriter* +byteswriter_create(Py_ssize_t size, int use_bytearray) +{ + if (size < 0) { + PyErr_SetString(PyExc_ValueError, "size must be >= 0"); + return NULL; + } + + PyBytesWriter *writer = _Py_FREELIST_POP_MEM(bytes_writers); + if (writer == NULL) { + writer = (PyBytesWriter *)PyMem_Malloc(sizeof(PyBytesWriter)); + if (writer == NULL) { + PyErr_NoMemory(); + return NULL; + } + } + writer->obj = NULL; + writer->size = 0; + writer->use_bytearray = use_bytearray; + + if (size >= 1) { + if (byteswriter_resize(writer, size, 0) < 0) { + PyBytesWriter_Discard(writer); + return NULL; + } + writer->size = 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) +{ + if (writer == NULL) { + return; + } + + Py_XDECREF(writer->obj); + _Py_FREELIST_FREE(bytes_writers, writer, PyMem_Free); +} + + +PyObject* +PyBytesWriter_FinishWithSize(PyBytesWriter *writer, Py_ssize_t size) +{ + PyObject *result; + if (size == 0) { + result = bytes_get_empty(); + } + else if (writer->obj != NULL) { + 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); + } + PyBytesWriter_Discard(writer); + return result; + +error: + PyBytesWriter_Discard(writer); + return NULL; +} + +PyObject* +PyBytesWriter_Finish(PyBytesWriter *writer) +{ + return PyBytesWriter_FinishWithSize(writer, writer->size); +} + + +PyObject* +PyBytesWriter_FinishWithPointer(PyBytesWriter *writer, void *buf) +{ + Py_ssize_t size = (char*)buf - byteswriter_data(writer); + if (size < 0 || size > byteswriter_allocated(writer)) { + PyBytesWriter_Discard(writer); + PyErr_SetString(PyExc_ValueError, "invalid end pointer"); + return NULL; + } + + return PyBytesWriter_FinishWithSize(writer, size); +} + + +void* +PyBytesWriter_GetData(PyBytesWriter *writer) +{ + return byteswriter_data(writer); +} + + +Py_ssize_t +PyBytesWriter_GetSize(PyBytesWriter *writer) +{ + return writer->size; +} + + +int +PyBytesWriter_Resize(PyBytesWriter *writer, Py_ssize_t size) +{ + if (size < 0) { + PyErr_SetString(PyExc_ValueError, "size must be >= 0"); + return -1; + } + if (byteswriter_resize(writer, size, 1) < 0) { + return -1; + } + writer->size = size; + return 0; +} + + +int +PyBytesWriter_Grow(PyBytesWriter *writer, Py_ssize_t size) +{ + if (size < 0 && writer->size + size < 0) { + PyErr_SetString(PyExc_ValueError, "invalid size"); + return -1; + } + if (size > PY_SSIZE_T_MAX - writer->size) { + PyErr_NoMemory(); + return -1; + } + size = writer->size + size; + + if (byteswriter_resize(writer, size, 1) < 0) { + return -1; + } + writer->size = size; + return 0; +} + + +void* +PyBytesWriter_GrowAndUpdatePointer(PyBytesWriter *writer, Py_ssize_t size, + void *buf) +{ + Py_ssize_t pos = (char*)buf - byteswriter_data(writer); + if (PyBytesWriter_Grow(writer, size) < 0) { + return NULL; + } + return byteswriter_data(writer) + pos; +} + + +int +PyBytesWriter_WriteBytes(PyBytesWriter *writer, + const void *bytes, Py_ssize_t size) +{ + if (size < 0) { + size_t len = strlen(bytes); + if (len > (size_t)PY_SSIZE_T_MAX) { + PyErr_NoMemory(); + return -1; + } + size = (Py_ssize_t)len; + } + + Py_ssize_t pos = writer->size; + if (PyBytesWriter_Grow(writer, size) < 0) { + return -1; + } + char *buf = byteswriter_data(writer); + memcpy(buf + pos, bytes, size); + return 0; +} diff --git a/Objects/object.c b/Objects/object.c index bd3ba02f8eb255..aaa3c0b338434e 100644 --- a/Objects/object.c +++ b/Objects/object.c @@ -945,6 +945,7 @@ _PyObject_ClearFreeLists(struct _Py_freelists *freelists, int is_finalization) clear_freelist(&freelists->object_stack_chunks, 1, PyMem_RawFree); } clear_freelist(&freelists->unicode_writers, is_finalization, PyMem_Free); + clear_freelist(&freelists->bytes_writers, is_finalization, PyMem_Free); clear_freelist(&freelists->ints, is_finalization, free_object); clear_freelist(&freelists->pycfunctionobject, is_finalization, PyObject_GC_Del); clear_freelist(&freelists->pycmethodobject, is_finalization, PyObject_GC_Del);