Skip to content

Commit eff71b5

Browse files
committed
Add PyBytesWriter_WriteBytes()
1 parent 92e1294 commit eff71b5

File tree

4 files changed

+62
-15
lines changed

4 files changed

+62
-15
lines changed

Include/cpython/bytesobject.h

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,11 @@ PyAPI_FUNC(Py_ssize_t) PyBytesWriter_GetSize(
6666
PyAPI_FUNC(Py_ssize_t) PyBytesWriter_GetAllocated(
6767
PyBytesWriter *writer);
6868

69+
PyAPI_FUNC(int) PyBytesWriter_WriteBytes(
70+
PyBytesWriter *writer,
71+
const void *bytes,
72+
Py_ssize_t size);
73+
6974
PyAPI_FUNC(int) PyBytesWriter_Resize(
7075
PyBytesWriter *writer,
7176
Py_ssize_t size);

Lib/test/test_capi/test_bytes.py

Lines changed: 11 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -297,15 +297,13 @@ class PyBytesWriterTest(unittest.TestCase):
297297
def create_writer(self, alloc=0, string=b''):
298298
return _testcapi.PyBytesWriter(alloc, string)
299299

300-
def test_empty(self):
300+
def test_create(self):
301301
# Test PyBytesWriter_Create()
302302
writer = self.create_writer()
303303
self.assertEqual(writer.get_size(), 0)
304304
self.assertEqual(writer.get_allocated(), self.SMALL_BUFFER)
305305
self.assertEqual(writer.finish(), b'')
306306

307-
def test_abc(self):
308-
# Test PyBytesWriter_Create()
309307
writer = self.create_writer(3, b'abc')
310308
self.assertEqual(writer.get_size(), 3)
311309
self.assertEqual(writer.get_allocated(), self.SMALL_BUFFER)
@@ -316,21 +314,19 @@ def test_abc(self):
316314
self.assertEqual(writer.get_allocated(), self.SMALL_BUFFER)
317315
self.assertEqual(writer.finish_with_size(3), b'abc')
318316

319-
# def test_write_bytes(self):
320-
# # Test PyBytesWriter_WriteBytes()
317+
def test_write_bytes(self):
318+
# Test PyBytesWriter_WriteBytes()
319+
writer = self.create_writer()
320+
writer.write_bytes(b'Hello World!', -1)
321+
self.assertEqual(writer.finish(), b'Hello World!')
321322

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!')
323+
writer = self.create_writer()
324+
writer.write_bytes(b'Hello ', -1)
325+
writer.write_bytes(b'World! <truncated>', 6)
326+
self.assertEqual(writer.finish(), b'Hello World!')
330327

331328
def test_resize(self):
332-
# Test PyBytesWriter_Extend()
333-
329+
# Test PyBytesWriter_Resize()
334330
writer = self.create_writer()
335331
writer.resize(len(b'number=123456'), b'number=123456')
336332
writer.resize(len(b'number=123456'), b'')

Modules/_testcapi/bytes.c

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,27 @@ writer_check(WriterObject *self)
130130
}
131131

132132

133+
static PyObject*
134+
writer_write_bytes(PyObject *self_raw, PyObject *args)
135+
{
136+
WriterObject *self = (WriterObject *)self_raw;
137+
if (writer_check(self) < 0) {
138+
return NULL;
139+
}
140+
141+
char *bytes;
142+
Py_ssize_t size;
143+
if (!PyArg_ParseTuple(args, "yn", &bytes, &size)) {
144+
return NULL;
145+
}
146+
147+
if (PyBytesWriter_WriteBytes(self->writer, bytes, size) < 0) {
148+
return NULL;
149+
}
150+
Py_RETURN_NONE;
151+
}
152+
153+
133154
static PyObject*
134155
writer_resize(PyObject *self_raw, PyObject *args)
135156
{
@@ -220,6 +241,7 @@ writer_finish_with_size(PyObject *self_raw, PyObject *args)
220241

221242

222243
static PyMethodDef writer_methods[] = {
244+
{"write_bytes", _PyCFunction_CAST(writer_write_bytes), METH_VARARGS},
223245
{"resize", _PyCFunction_CAST(writer_resize), METH_VARARGS},
224246
{"get_size", _PyCFunction_CAST(writer_get_size), METH_NOARGS},
225247
{"get_allocated", _PyCFunction_CAST(writer_get_allocated), METH_NOARGS},

Objects/bytesobject.c

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3925,3 +3925,27 @@ PyBytesWriter_Resize(PyBytesWriter *writer, Py_ssize_t size)
39253925
writer->size = size;
39263926
return 0;
39273927
}
3928+
3929+
3930+
int
3931+
PyBytesWriter_WriteBytes(PyBytesWriter *writer,
3932+
const void *bytes, Py_ssize_t size)
3933+
{
3934+
if (size < 0) {
3935+
size = strlen(bytes);
3936+
}
3937+
3938+
Py_ssize_t pos = writer->size;
3939+
if (size > PY_SSIZE_T_MAX - pos) {
3940+
PyErr_NoMemory();
3941+
return -1;
3942+
}
3943+
Py_ssize_t total = pos + size;
3944+
3945+
if (PyBytesWriter_Resize(writer, total) < 0) {
3946+
return -1;
3947+
}
3948+
char *buf = byteswriter_data(writer);
3949+
memcpy(buf + pos, bytes, size);
3950+
return 0;
3951+
}

0 commit comments

Comments
 (0)