Skip to content

Commit 8091566

Browse files
encukouencukou
committed
Alternate PoC implementation of PyBytesWriter
1 parent 3712eb5 commit 8091566

3 files changed

Lines changed: 94 additions & 17 deletions

File tree

‎Include/cpython/bytesobject.h‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,11 +49,11 @@ PyAPI_FUNC(PyObject *) PyBytesWriter_Finish(
4949
PyAPI_FUNC(void) PyBytesWriter_Discard(PyBytesWriter *writer);
5050

5151
PyAPI_FUNC(int) PyBytesWriter_Prepare(
52-
PyBytesWriter *writer,
52+
PyBytesWriter **writer,
5353
char **str,
5454
Py_ssize_t size);
5555
PyAPI_FUNC(int) PyBytesWriter_WriteBytes(
56-
PyBytesWriter *writer,
56+
PyBytesWriter **writer,
5757
char **str,
5858
const void *bytes,
5959
Py_ssize_t size);

‎Modules/_testcapi/bytes.c‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -55,12 +55,12 @@ test_byteswriter(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
5555
return NULL;
5656
}
5757

58-
if (PyBytesWriter_WriteBytes(writer, &str, "abc", 3) < 0) {
58+
if (PyBytesWriter_WriteBytes(&writer, &str, "abc", 3) < 0) {
5959
goto error;
6060
}
6161

6262
// write empty string
63-
if (PyBytesWriter_WriteBytes(writer, &str, "", 0) < 0) {
63+
if (PyBytesWriter_WriteBytes(&writer, &str, "", 0) < 0) {
6464
goto error;
6565
}
6666

@@ -89,7 +89,7 @@ test_byteswriter_discard(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
8989
if (writer == NULL) {
9090
return NULL;
9191
}
92-
assert(PyBytesWriter_WriteBytes(writer, &str, "abc", 3) == 0);
92+
assert(PyBytesWriter_WriteBytes(&writer, &str, "abc", 3) == 0);
9393

9494
PyBytesWriter_Discard(writer);
9595
Py_RETURN_NONE;
@@ -106,10 +106,10 @@ test_byteswriter_writebytes(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(arg
106106
return NULL;
107107
}
108108

109-
if (PyBytesWriter_WriteBytes(writer, &str, "abc", 3) < 0) {
109+
if (PyBytesWriter_WriteBytes(&writer, &str, "abc", 3) < 0) {
110110
goto error;
111111
}
112-
if (PyBytesWriter_WriteBytes(writer, &str, "def", 3) < 0) {
112+
if (PyBytesWriter_WriteBytes(&writer, &str, "def", 3) < 0) {
113113
goto error;
114114
}
115115

@@ -140,11 +140,11 @@ test_byteswriter_prepare(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
140140
}
141141

142142
// test error on purpose (negative size)
143-
assert(PyBytesWriter_Prepare(writer, &str, -3) < 0);
143+
assert(PyBytesWriter_Prepare(&writer, &str, -3) < 0);
144144
assert(PyErr_ExceptionMatches(PyExc_ValueError));
145145
PyErr_Clear();
146146

147-
if (PyBytesWriter_Prepare(writer, &str, 3) < 0) {
147+
if (PyBytesWriter_Prepare(&writer, &str, 3) < 0) {
148148
PyBytesWriter_Discard(writer);
149149
return NULL;
150150
}

‎Objects/bytesobject.c‎

Lines changed: 85 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3383,7 +3383,7 @@ _PyBytesWriter_Init(_PyBytesWriter *writer)
33833383
#endif
33843384
}
33853385

3386-
3386+
/*
33873387
PyBytesWriter* PyBytesWriter_Create(Py_ssize_t size, char **pstr)
33883388
{
33893389
_PyBytesWriter *writer = PyMem_Malloc(sizeof(_PyBytesWriter));
@@ -3405,22 +3405,22 @@ PyBytesWriter* PyBytesWriter_Create(Py_ssize_t size, char **pstr)
34053405
*pstr = str;
34063406
return (PyBytesWriter*)writer;
34073407
}
3408-
3408+
*/
34093409

34103410
void
34113411
_PyBytesWriter_Dealloc(_PyBytesWriter *writer)
34123412
{
34133413
Py_CLEAR(writer->buffer);
34143414
}
34153415

3416-
3416+
/*
34173417
void
34183418
PyBytesWriter_Discard(PyBytesWriter *writer)
34193419
{
34203420
_PyBytesWriter_Dealloc((_PyBytesWriter*)writer);
34213421
PyMem_Free(writer);
34223422
}
3423-
3423+
*/
34243424

34253425
Py_LOCAL_INLINE(char*)
34263426
_PyBytesWriter_AsString(_PyBytesWriter *writer)
@@ -3593,7 +3593,7 @@ _PyBytesWriter_Prepare(_PyBytesWriter *writer, void *str, Py_ssize_t size)
35933593
return str;
35943594
}
35953595

3596-
3596+
/*
35973597
int
35983598
PyBytesWriter_Prepare(PyBytesWriter *writer, char **str, Py_ssize_t size)
35993599
{
@@ -3609,6 +3609,7 @@ PyBytesWriter_Prepare(PyBytesWriter *writer, char **str, Py_ssize_t size)
36093609
*str = str2;
36103610
return 0;
36113611
}
3612+
*/
36123613

36133614

36143615
/* Allocate the buffer to write size bytes.
@@ -3688,14 +3689,15 @@ _PyBytesWriter_Finish(_PyBytesWriter *writer, void *str)
36883689
return result;
36893690
}
36903691

3691-
3692+
/*
36923693
PyObject *
36933694
PyBytesWriter_Finish(PyBytesWriter *writer, char *str)
36943695
{
36953696
PyObject *res = _PyBytesWriter_Finish((_PyBytesWriter*)writer, str);
36963697
PyMem_Free(writer);
36973698
return res;
36983699
}
3700+
*/
36993701

37003702

37013703
void*
@@ -3714,7 +3716,7 @@ _PyBytesWriter_WriteBytes(_PyBytesWriter *writer, void *ptr,
37143716
return str;
37153717
}
37163718

3717-
3719+
/*
37183720
int
37193721
PyBytesWriter_WriteBytes(PyBytesWriter *writer, char **str,
37203722
const void *bytes, Py_ssize_t size)
@@ -3728,7 +3730,7 @@ PyBytesWriter_WriteBytes(PyBytesWriter *writer, char **str,
37283730
*str = str2;
37293731
return 0;
37303732
}
3731-
3733+
*/
37323734

37333735
void
37343736
_PyBytes_Repeat(char* dest, Py_ssize_t len_dest,
@@ -3753,3 +3755,78 @@ _PyBytes_Repeat(char* dest, Py_ssize_t len_dest,
37533755
}
37543756
}
37553757

3758+
static inline char* _PyBytesWriter_AS_STRING(PyObject *op)
3759+
{
3760+
return _Py_CAST(PyBytesObject*, op)->ob_sval;
3761+
}
3762+
3763+
// Create a bytes writer instance.
3764+
PyBytesWriter* PyBytesWriter_Create(Py_ssize_t size, char **str) {
3765+
if (size < 0) {
3766+
size = 0;
3767+
}
3768+
PyObject *res = PyBytes_FromStringAndSize(NULL, size);
3769+
if (!res) return NULL;
3770+
Py_SET_TYPE(res, NULL);
3771+
*str = _PyBytesWriter_AS_STRING(res);
3772+
return (PyBytesWriter*)res;
3773+
}
3774+
3775+
// Return the final Python bytes object and destroy the writer instance.
3776+
PyObject* PyBytesWriter_Finish(PyBytesWriter *writer, char *str) {
3777+
PyObject *o = (PyObject *)writer;
3778+
assert(Py_TYPE(o) == NULL);
3779+
Py_ssize_t size_written = str - _PyBytesWriter_AS_STRING(o);
3780+
Py_SET_TYPE(o, &PyBytes_Type);
3781+
if (size_written == Py_SIZE(o)) {
3782+
return o;
3783+
}
3784+
PyObject *res = PyBytes_FromStringAndSize(_PyBytesWriter_AS_STRING(o), size_written);
3785+
Py_DECREF(o);
3786+
return res;
3787+
}
3788+
3789+
// Discard the internal bytes buffer and destroy the writer instance.
3790+
void PyBytesWriter_Discard(PyBytesWriter *writer) {
3791+
PyObject *o = (PyObject *)writer;
3792+
Py_SET_TYPE(o, &PyBytes_Type);
3793+
Py_DECREF(o);
3794+
}
3795+
3796+
// Allocate *size* bytes to prepare writing *size* bytes into *writer*.
3797+
int PyBytesWriter_Prepare(PyBytesWriter **writer, char **str, Py_ssize_t size) {
3798+
if (size < 0) {
3799+
PyErr_SetString(PyExc_ValueError, "PyBytesWriter: size must not be negative");
3800+
return -1;
3801+
}
3802+
PyObject *o = (PyObject *)*writer;
3803+
Py_ssize_t size_written = *str - _PyBytesWriter_AS_STRING(o);
3804+
assert(size_written >= 0);
3805+
Py_ssize_t size_left = Py_SIZE(o) - size_written;
3806+
assert(size_left >= 0);
3807+
if (size_left >= size) {
3808+
return 0;
3809+
}
3810+
char *new_str;
3811+
PyBytesWriter *new_writer = PyBytesWriter_Create(size_written + size, &new_str);
3812+
if (!new_writer) {
3813+
return -1;
3814+
}
3815+
memcpy(new_str, _PyBytesWriter_AS_STRING(o), size_written);
3816+
*str = new_str + size_written;
3817+
Py_SET_TYPE(new_writer, NULL);
3818+
Py_SET_TYPE(o, &PyBytes_Type);
3819+
Py_SETREF(*writer, new_writer);
3820+
return 0;
3821+
}
3822+
3823+
// Write a the bytes string *bytes* of *size* bytes into *writer*.
3824+
int PyBytesWriter_WriteBytes(PyBytesWriter **writer, char **str, const void *bytes, Py_ssize_t size) {
3825+
if (PyBytesWriter_Prepare(writer, str, size) == -1) {
3826+
return -1;
3827+
}
3828+
memcpy(*str, bytes, size);
3829+
*str += size;
3830+
return 0;
3831+
}
3832+

0 commit comments

Comments
 (0)