Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 42 additions & 4 deletions Lib/test/test_capi/test_bytes.py
Original file line number Diff line number Diff line change
Expand Up @@ -299,11 +299,11 @@ def test_join(self):
bytes_join(b'', NULL)


class BytesWriterTest(unittest.TestCase):
result_type = bytes
class BaseWriterTest:
result_type = NotImplementedError

def create_writer(self, alloc=0, string=b''):
return _testcapi.PyBytesWriter(alloc, string, 0)
raise NotImplementedError

def test_create(self):
# Test PyBytesWriter_Create()
Expand Down Expand Up @@ -388,10 +388,48 @@ def test_example_highlevel(self):
self.assertEqual(_testcapi.byteswriter_highlevel(), b'Hello World!')


class ByteArrayWriterTest(BytesWriterTest):
class BytesWriterTest(BaseWriterTest, unittest.TestCase):
result_type = bytes

def create_writer(self, alloc=0, string=b''):
# Test PyBytesWriter_Create()
return _testcapi.PyBytesWriter(alloc, string, 0)

# Only PyBytesWriter_Create() returns singletons
def test_singletons(self):
empty = b''
singletons = {ch: bytes((ch,)) for ch in range(256)}
small_buffer = _testcapi.PyBytesWriter_small_buffer

writer = self.create_writer()
self.assertIs(writer.finish(), empty)

# Test writer larger than small_buffer
writer = self.create_writer()
unused_text = b'x' * (small_buffer * 2)
writer.write_bytes(unused_text, len(unused_text))
self.assertIs(writer.finish_with_size(0), empty)

for ch in range(256):
text = bytes((ch,))

writer = self.create_writer()
writer.write_bytes(text, 1)
self.assertIs(writer.finish(), singletons[ch])

# Test writer larger than small_buffer
writer = self.create_writer()
writer.write_bytes(text, 1)
unused_text = b'x' * (small_buffer * 2)
writer.write_bytes(unused_text, len(unused_text))
self.assertIs(writer.finish_with_size(1), singletons[ch])


class ByteArrayWriterTest(BaseWriterTest, unittest.TestCase):
result_type = bytearray

def create_writer(self, alloc=0, string=b''):
# Test private _PyBytesWriter_CreateByteArray()
return _testcapi.PyBytesWriter(alloc, string, 1)


Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
:c:func:`PyBytesWriter_FinishWithSize` now returns single byte singletons if
*size* equals to ``1``. Patch by Victor Stinner.
13 changes: 11 additions & 2 deletions Modules/_testcapi/bytes.c
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
#include "parts.h"
#include "util.h"

#include <stddef.h> // offsetof()

#include "pycore_bytesobject.h" // _PyBytesWriter_CreateByteArray()


Expand Down Expand Up @@ -150,8 +152,8 @@ writer_write_bytes(PyObject *self_raw, PyObject *args)
}

char *bytes;
Py_ssize_t size;
if (!PyArg_ParseTuple(args, "yn", &bytes, &size)) {
Py_ssize_t unused_size, size;
if (!PyArg_ParseTuple(args, "y#n", &bytes, &unused_size, &size)) {
return NULL;
}

Expand Down Expand Up @@ -377,5 +379,12 @@ _PyTestCapi_Init_Bytes(PyObject *m)
}
Py_DECREF(writer_type);

// PyBytesWriter.obj is the second member, small_buffer is the first member
long size = (long)offsetof(PyBytesWriter, obj);
if (PyModule_AddIntConstant(m, "PyBytesWriter_small_buffer", size) < 0) {
Py_DECREF(writer_type);
return -1;
}

return 0;
}
10 changes: 10 additions & 0 deletions Objects/bytesobject.c
Original file line number Diff line number Diff line change
Expand Up @@ -3766,13 +3766,23 @@ PyBytesWriter_FinishWithSize(PyBytesWriter *writer, Py_ssize_t size)
}
}
}

result = writer->obj;
writer->obj = NULL;

if (size == 1 && !writer->use_bytearray) {
// Get the single byte singleton
unsigned char ch = PyBytes_AS_STRING(result)[0];
PyObject *op = (PyObject*)CHARACTER(ch);
assert(_Py_IsImmortal(op));
Py_SETREF(result, op);
}
}
else if (writer->use_bytearray) {
result = PyByteArray_FromStringAndSize(writer->small_buffer, size);
}
else {
// The function returns single byte singleton if size equals 1
result = PyBytes_FromStringAndSize(writer->small_buffer, size);
}
PyBytesWriter_Discard(writer);
Expand Down
Loading