Skip to content

Commit 92e1294

Browse files
vstinnervstinner
committed
Add tests
* Add PyBytesWriter_GetSize() * Rename: * PyBytesWriter_Data() => PyBytesWriter_GetData() * PyBytesWriter_Allocated() => PyBytesWriter_GetAllocated()
1 parent 8761a9b commit 92e1294

File tree

7 files changed

+318
-29
lines changed

7 files changed

+318
-29
lines changed

‎Include/cpython/bytesobject.h‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ _PyBytes_Join(PyObject *sep, PyObject *iterable)
4747
typedef struct PyBytesWriter PyBytesWriter;
4848

4949
PyAPI_FUNC(PyBytesWriter *) PyBytesWriter_Create(
50-
Py_ssize_t alloc);
50+
Py_ssize_t size);
5151
PyAPI_FUNC(void) PyBytesWriter_Discard(
5252
PyBytesWriter *writer);
5353
PyAPI_FUNC(PyObject*) PyBytesWriter_Finish(
@@ -59,9 +59,11 @@ PyAPI_FUNC(PyObject*) PyBytesWriter_FinishWithEndPointer(
5959
PyBytesWriter *writer,
6060
void *data);
6161

62-
PyAPI_FUNC(void*) PyBytesWriter_Data(
62+
PyAPI_FUNC(void*) PyBytesWriter_GetData(
63+
PyBytesWriter *writer);
64+
PyAPI_FUNC(Py_ssize_t) PyBytesWriter_GetSize(
6365
PyBytesWriter *writer);
64-
PyAPI_FUNC(Py_ssize_t) PyBytesWriter_Allocated(
66+
PyAPI_FUNC(Py_ssize_t) PyBytesWriter_GetAllocated(
6567
PyBytesWriter *writer);
6668

6769
PyAPI_FUNC(int) PyBytesWriter_Resize(

‎Lib/test/test_capi/test_bytes.py‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -291,5 +291,81 @@ def test_join(self):
291291
bytes_join(b'', NULL)
292292

293293

294+
class PyBytesWriterTest(unittest.TestCase):
295+
SMALL_BUFFER = 256 # bytes
296+
297+
def create_writer(self, alloc=0, string=b''):
298+
return _testcapi.PyBytesWriter(alloc, string)
299+
300+
def test_empty(self):
301+
# Test PyBytesWriter_Create()
302+
writer = self.create_writer()
303+
self.assertEqual(writer.get_size(), 0)
304+
self.assertEqual(writer.get_allocated(), self.SMALL_BUFFER)
305+
self.assertEqual(writer.finish(), b'')
306+
307+
def test_abc(self):
308+
# Test PyBytesWriter_Create()
309+
writer = self.create_writer(3, b'abc')
310+
self.assertEqual(writer.get_size(), 3)
311+
self.assertEqual(writer.get_allocated(), self.SMALL_BUFFER)
312+
self.assertEqual(writer.finish(), b'abc')
313+
314+
writer = self.create_writer(10, b'abc')
315+
self.assertEqual(writer.get_size(), 10)
316+
self.assertEqual(writer.get_allocated(), self.SMALL_BUFFER)
317+
self.assertEqual(writer.finish_with_size(3), b'abc')
318+
319+
# def test_write_bytes(self):
320+
# # Test PyBytesWriter_WriteBytes()
321+
322+
# writer = self.create_writer()
323+
# writer.write_bytes(b'Hello World!', -1)
324+
# self.assertEqual(writer.finish(), b'Hello World!')
325+
326+
# writer = self.create_writer()
327+
# writer.write_bytes(b'Hello ', -1)
328+
# writer.write_bytes(b'World! <truncated>', 6)
329+
# self.assertEqual(writer.finish(), b'Hello World!')
330+
331+
def test_resize(self):
332+
# Test PyBytesWriter_Extend()
333+
334+
writer = self.create_writer()
335+
writer.resize(len(b'number=123456'), b'number=123456')
336+
writer.resize(len(b'number=123456'), b'')
337+
self.assertEqual(writer.get_size(), len(b'number=123456'))
338+
self.assertEqual(writer.finish(), b'number=123456')
339+
340+
writer = self.create_writer()
341+
writer.resize(0, b'')
342+
writer.resize(len(b'number=123456'), b'number=123456')
343+
self.assertEqual(writer.finish(), b'number=123456')
344+
345+
writer = self.create_writer()
346+
writer.resize(len(b'number='), b'number=')
347+
writer.resize(len(b'number=123456'), b'123456')
348+
self.assertEqual(writer.finish(), b'number=123456')
349+
350+
writer = self.create_writer()
351+
writer.resize(len(b'number='), b'number=')
352+
writer.resize(len(b'number='), b'')
353+
writer.resize(len(b'number=123456'), b'123456')
354+
self.assertEqual(writer.finish(), b'number=123456')
355+
356+
writer = self.create_writer()
357+
writer.resize(len(b'number'), b'number')
358+
writer.resize(len(b'number='), b'=')
359+
writer.resize(len(b'number=123'), b'123')
360+
writer.resize(len(b'number=123456'), b'456')
361+
self.assertEqual(writer.finish(), b'number=123456')
362+
363+
# def test_format(self):
364+
# # Test PyBytesWriter_Format()
365+
# writer = self.create_writer()
366+
# writer.format_i(123456)
367+
# self.assertEqual(writer.finish(), b'123456')
368+
369+
294370
if __name__ == "__main__":
295371
unittest.main()

‎Modules/_pickle.c‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2624,19 +2624,19 @@ raw_unicode_escape(PyObject *obj)
26242624
if (writer == NULL) {
26252625
return NULL;
26262626
}
2627-
char *p = PyBytesWriter_Data(writer);
2627+
char *p = PyBytesWriter_GetData(writer);
26282628

26292629
for (Py_ssize_t i=0; i < size; i++) {
26302630
Py_UCS4 ch = PyUnicode_READ(kind, data, i);
26312631
/* Map 32-bit characters to '\Uxxxxxxxx' */
26322632
if (ch >= 0x10000) {
26332633
/* -1: subtract 1 preallocated byte */
26342634
alloc += 10-1;
2635-
Py_ssize_t pos = p - (char*)PyBytesWriter_Data(writer);
2635+
Py_ssize_t pos = p - (char*)PyBytesWriter_GetData(writer);
26362636
if (PyBytesWriter_Resize(writer, alloc) < 0) {
26372637
goto error;
26382638
}
2639-
p = (char*)PyBytesWriter_Data(writer) + pos;
2639+
p = (char*)PyBytesWriter_GetData(writer) + pos;
26402640

26412641
*p++ = '\\';
26422642
*p++ = 'U';
@@ -2656,11 +2656,11 @@ raw_unicode_escape(PyObject *obj)
26562656
{
26572657
/* -1: subtract 1 preallocated byte */
26582658
alloc += 6-1;
2659-
Py_ssize_t pos = p - (char*)PyBytesWriter_Data(writer);
2659+
Py_ssize_t pos = p - (char*)PyBytesWriter_GetData(writer);
26602660
if (PyBytesWriter_Resize(writer, alloc) < 0) {
26612661
goto error;
26622662
}
2663-
p = (char*)PyBytesWriter_Data(writer) + pos;
2663+
p = (char*)PyBytesWriter_GetData(writer) + pos;
26642664

26652665
*p++ = '\\';
26662666
*p++ = 'u';

‎Modules/_struct.c‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2291,7 +2291,7 @@ s_pack(PyObject *self, PyObject *const *args, Py_ssize_t nargs)
22912291
if (writer == NULL) {
22922292
return NULL;
22932293
}
2294-
char *buf = PyBytesWriter_Data(writer);
2294+
char *buf = PyBytesWriter_GetData(writer);
22952295

22962296
/* Call the guts */
22972297
if ( s_pack_internal(soself, args, 0, buf, state) != 0 ) {

‎Modules/_testcapi/bytes.c‎

Lines changed: 203 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,199 @@ bytes_join(PyObject *Py_UNUSED(module), PyObject *args)
5151
}
5252

5353

54+
// --- PyBytesWriter type ---------------------------------------------------
55+
56+
typedef struct {
57+
PyObject_HEAD
58+
PyBytesWriter *writer;
59+
} WriterObject;
60+
61+
62+
static PyObject *
63+
writer_new(PyTypeObject *type, PyObject *args, PyObject *kwargs)
64+
{
65+
WriterObject *self = (WriterObject *)type->tp_alloc(type, 0);
66+
if (!self) {
67+
return NULL;
68+
}
69+
self->writer = NULL;
70+
return (PyObject*)self;
71+
}
72+
73+
74+
static int
75+
writer_init(PyObject *self_raw, PyObject *args, PyObject *kwargs)
76+
{
77+
WriterObject *self = (WriterObject *)self_raw;
78+
if (self->writer) {
79+
PyBytesWriter_Discard(self->writer);
80+
}
81+
82+
if (kwargs && PyDict_GET_SIZE(kwargs)) {
83+
PyErr_Format(PyExc_TypeError,
84+
"PyBytesWriter() takes exactly no keyword arguments");
85+
return -1;
86+
}
87+
88+
Py_ssize_t alloc;
89+
char *str;
90+
Py_ssize_t str_size;
91+
if (!PyArg_ParseTuple(args, "ny#", &alloc, &str, &str_size)) {
92+
return -1;
93+
}
94+
95+
self->writer = PyBytesWriter_Create(alloc);
96+
if (self->writer == NULL) {
97+
return -1;
98+
}
99+
100+
if (str_size) {
101+
char *buf = PyBytesWriter_GetData(self->writer);
102+
memcpy(buf, str, str_size);
103+
}
104+
105+
return 0;
106+
}
107+
108+
109+
static void
110+
writer_dealloc(PyObject *self_raw)
111+
{
112+
WriterObject *self = (WriterObject *)self_raw;
113+
PyTypeObject *tp = Py_TYPE(self);
114+
if (self->writer) {
115+
PyBytesWriter_Discard(self->writer);
116+
}
117+
tp->tp_free(self);
118+
Py_DECREF(tp);
119+
}
120+
121+
122+
static inline int
123+
writer_check(WriterObject *self)
124+
{
125+
if (self->writer == NULL) {
126+
PyErr_SetString(PyExc_ValueError, "operation on finished writer");
127+
return -1;
128+
}
129+
return 0;
130+
}
131+
132+
133+
static PyObject*
134+
writer_resize(PyObject *self_raw, PyObject *args)
135+
{
136+
WriterObject *self = (WriterObject *)self_raw;
137+
if (writer_check(self) < 0) {
138+
return NULL;
139+
}
140+
141+
Py_ssize_t size;
142+
char *str;
143+
Py_ssize_t str_size;
144+
if (!PyArg_ParseTuple(args,
145+
"ny#",
146+
&size, &str, &str_size)) {
147+
return NULL;
148+
}
149+
assert(size >= str_size);
150+
151+
Py_ssize_t pos = PyBytesWriter_GetSize(self->writer);
152+
if (PyBytesWriter_Resize(self->writer, size) < 0) {
153+
return NULL;
154+
}
155+
156+
char *buf = PyBytesWriter_GetData(self->writer);
157+
memcpy(buf + pos, str, str_size);
158+
159+
Py_RETURN_NONE;
160+
}
161+
162+
163+
static PyObject*
164+
writer_get_size(PyObject *self_raw, PyObject *Py_UNUSED(args))
165+
{
166+
WriterObject *self = (WriterObject *)self_raw;
167+
if (writer_check(self) < 0) {
168+
return NULL;
169+
}
170+
171+
Py_ssize_t alloc = PyBytesWriter_GetSize(self->writer);
172+
return PyLong_FromSsize_t(alloc);
173+
}
174+
175+
176+
static PyObject*
177+
writer_get_allocated(PyObject *self_raw, PyObject *Py_UNUSED(args))
178+
{
179+
WriterObject *self = (WriterObject *)self_raw;
180+
if (writer_check(self) < 0) {
181+
return NULL;
182+
}
183+
184+
Py_ssize_t alloc = PyBytesWriter_GetAllocated(self->writer);
185+
return PyLong_FromSsize_t(alloc);
186+
}
187+
188+
189+
static PyObject*
190+
writer_finish(PyObject *self_raw, PyObject *Py_UNUSED(args))
191+
{
192+
WriterObject *self = (WriterObject *)self_raw;
193+
if (writer_check(self) < 0) {
194+
return NULL;
195+
}
196+
197+
PyObject *str = PyBytesWriter_Finish(self->writer);
198+
self->writer = NULL;
199+
return str;
200+
}
201+
202+
203+
static PyObject*
204+
writer_finish_with_size(PyObject *self_raw, PyObject *args)
205+
{
206+
WriterObject *self = (WriterObject *)self_raw;
207+
if (writer_check(self) < 0) {
208+
return NULL;
209+
}
210+
211+
Py_ssize_t size;
212+
if (!PyArg_ParseTuple(args, "n", &size)) {
213+
return NULL;
214+
}
215+
216+
PyObject *str = PyBytesWriter_FinishWithSize(self->writer, size);
217+
self->writer = NULL;
218+
return str;
219+
}
220+
221+
222+
static PyMethodDef writer_methods[] = {
223+
{"resize", _PyCFunction_CAST(writer_resize), METH_VARARGS},
224+
{"get_size", _PyCFunction_CAST(writer_get_size), METH_NOARGS},
225+
{"get_allocated", _PyCFunction_CAST(writer_get_allocated), METH_NOARGS},
226+
{"finish", _PyCFunction_CAST(writer_finish), METH_NOARGS},
227+
{"finish_with_size", _PyCFunction_CAST(writer_finish_with_size), METH_VARARGS},
228+
{NULL, NULL} /* sentinel */
229+
};
230+
231+
static PyType_Slot Writer_Type_slots[] = {
232+
{Py_tp_new, writer_new},
233+
{Py_tp_init, writer_init},
234+
{Py_tp_dealloc, writer_dealloc},
235+
{Py_tp_methods, writer_methods},
236+
{0, 0}, /* sentinel */
237+
};
238+
239+
static PyType_Spec Writer_spec = {
240+
.name = "_testcapi.PyBytesWriter",
241+
.basicsize = sizeof(WriterObject),
242+
.flags = Py_TPFLAGS_DEFAULT,
243+
.slots = Writer_Type_slots,
244+
};
245+
246+
54247
static PyMethodDef test_methods[] = {
55248
{"bytes_resize", bytes_resize, METH_VARARGS},
56249
{"bytes_join", bytes_join, METH_VARARGS},
@@ -64,5 +257,15 @@ _PyTestCapi_Init_Bytes(PyObject *m)
64257
return -1;
65258
}
66259

260+
PyTypeObject *writer_type = (PyTypeObject *)PyType_FromSpec(&Writer_spec);
261+
if (writer_type == NULL) {
262+
return -1;
263+
}
264+
if (PyModule_AddType(m, writer_type) < 0) {
265+
Py_DECREF(writer_type);
266+
return -1;
267+
}
268+
Py_DECREF(writer_type);
269+
67270
return 0;
68271
}

0 commit comments

Comments
 (0)