gh-129813, PEP 782: Add PyBytesWriter C API (#138822)

This commit is contained in:
Victor Stinner 2025-09-12 13:41:59 +02:00 committed by GitHub
parent 3d521a62e7
commit adb414044f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 873 additions and 0 deletions

View file

@ -219,3 +219,153 @@ called with a non-bytes parameter.
reallocation fails, the original bytes object at *\*bytes* is deallocated, reallocation fails, the original bytes object at *\*bytes* is deallocated,
*\*bytes* is set to ``NULL``, :exc:`MemoryError` is set, and ``-1`` is *\*bytes* is set to ``NULL``, :exc:`MemoryError` is set, and ``-1`` is
returned. 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;

View file

@ -707,6 +707,22 @@ New features
and :c:data:`Py_mod_abi`. and :c:data:`Py_mod_abi`.
(Contributed by Petr Viktorin in :gh:`137210`.) (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 Porting to Python 3.15
---------------------- ----------------------

View file

@ -40,3 +40,42 @@ _PyBytes_Join(PyObject *sep, PyObject *iterable)
{ {
return PyBytes_Join(sep, 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);

View file

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

View file

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

View file

@ -299,5 +299,80 @@ def test_join(self):
bytes_join(b'', NULL) 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_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__": if __name__ == "__main__":
unittest.main() unittest.main()

View file

@ -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.

View file

@ -1,6 +1,11 @@
// Use pycore_bytes.h
#define PYTESTCAPI_NEED_INTERNAL_API
#include "parts.h" #include "parts.h"
#include "util.h" #include "util.h"
#include "pycore_bytesobject.h" // _PyBytesWriter_CreateByteArray()
/* Test _PyBytes_Resize() */ /* Test _PyBytes_Resize() */
static PyObject * 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[] = { static PyMethodDef test_methods[] = {
{"bytes_resize", bytes_resize, METH_VARARGS}, {"bytes_resize", bytes_resize, METH_VARARGS},
{"bytes_join", bytes_join, METH_VARARGS}, {"bytes_join", bytes_join, METH_VARARGS},
{"byteswriter_abc", byteswriter_abc, METH_NOARGS},
{"byteswriter_resize", byteswriter_resize, METH_NOARGS},
{NULL}, {NULL},
}; };
@ -64,5 +324,15 @@ _PyTestCapi_Init_Bytes(PyObject *m)
return -1; 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; return 0;
} }

View file

@ -7,6 +7,7 @@
#include "pycore_call.h" // _PyObject_CallNoArgs() #include "pycore_call.h" // _PyObject_CallNoArgs()
#include "pycore_ceval.h" // _PyEval_GetBuiltin() #include "pycore_ceval.h" // _PyEval_GetBuiltin()
#include "pycore_format.h" // F_LJUST #include "pycore_format.h" // F_LJUST
#include "pycore_freelist.h" // _Py_FREELIST_FREE()
#include "pycore_global_objects.h"// _Py_GET_GLOBAL_OBJECT() #include "pycore_global_objects.h"// _Py_GET_GLOBAL_OBJECT()
#include "pycore_initconfig.h" // _PyStatus_OK() #include "pycore_initconfig.h" // _PyStatus_OK()
#include "pycore_long.h" // _PyLong_DigitValue #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;
}

View file

@ -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->object_stack_chunks, 1, PyMem_RawFree);
} }
clear_freelist(&freelists->unicode_writers, is_finalization, PyMem_Free); 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->ints, is_finalization, free_object);
clear_freelist(&freelists->pycfunctionobject, is_finalization, PyObject_GC_Del); clear_freelist(&freelists->pycfunctionobject, is_finalization, PyObject_GC_Del);
clear_freelist(&freelists->pycmethodobject, is_finalization, PyObject_GC_Del); clear_freelist(&freelists->pycmethodobject, is_finalization, PyObject_GC_Del);