FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

gh-157710: Defer allocation in PyUnicodeWriter_Create() (#157969) · python/cpython@71f6c00 · GitHub

  1. python
  2. cpython

Repository navigation

Commit 71f6c00

Browse files
vstinner
authored
gh-157710: Defer allocation in PyUnicodeWriter_Create() (#157969)
PyUnicodeWriter_Create(length) no longer allocates 'length' characters immediately. The allocation of the buffer is now done lazily at the first write, except if the read-only optimization is used. So the read-only optimization can also be used even if length is greater than 0. No longer overallocate the first buffer (at the first write). Only overallocate when the buffer is resized (at the second write). It avoids the need to truncate in PyUnicodeWriter_Finish() when PyUnicodeWriter_Create(length) used the exact output length. Add get_buffer() method to writer tests.
1 parent 7cfb655 commit 71f6c00

3 files changed

Lines changed: 144 additions & 63 deletions

File tree

‎Lib/test/test_capi/test_unicode.py‎

Lines changed: 91 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1880,6 +1880,36 @@ def test_basic(self):
18801880
self.assertEqual(writer.finish(),
18811881
"var=long value 'repr'")
18821882

1883+
def test_create(self):
1884+
# Test PyUnicodeWriter_Create() with non-zero size
1885+
s = 'Monty Python'
1886+
1887+
# Preallocate the exact length. Use 2 writes to force the creation
1888+
# of a buffer:
1889+
# 1. Use the read-only optimization.
1890+
# 2. Allocate a buffer of length character.
1891+
# No resize needed in finish().
1892+
writer = self.create_writer(len(s))
1893+
writer.write_str(s[:5])
1894+
self.assertEqual(writer.get_buffer(), (5, 127, True))
1895+
writer.write_str(s[5:])
1896+
self.assertEqual(writer.get_buffer(), (len(s), 127, False))
1897+
self.assertEqual(writer.finish(), s)
1898+
1899+
# Preallocate len(s)-1 characters. Use 3 writes:
1900+
# 1. Use read-only optimization.
1901+
# 2. Allocate a buffer of len-1 characters.
1902+
# 3. Resize the buffer with overallocation.
1903+
# finish() has to truncate the buffer.
1904+
writer = self.create_writer(len(s) - 1)
1905+
writer.write_str(s[:2])
1906+
self.assertEqual(writer.get_buffer(), (2, 127, True))
1907+
writer.write_str(s[2:5])
1908+
self.assertEqual(writer.get_buffer(), (len(s) - 1, 127, False))
1909+
writer.write_str(s[5:])
1910+
self.assertGreater(writer.get_buffer()[0], len(s))
1911+
self.assertEqual(writer.finish(), s)
1912+
18831913
def test_repr_null(self):
18841914
writer = self.create_writer(0)
18851915
writer.write_utf8(b'var=', -1)
@@ -2087,32 +2117,39 @@ def test_substring_empty(self):
20872117
def test_singletons(self):
20882118
for size in (0, 123):
20892119
with self.subTest(size=size):
2120+
# PyUnicodeWriter_Finish() returns the empty string singleton
2121+
# if no character has been written.
20902122
writer = self.create_writer(size)
20912123
writer.write_utf8(b'utf8', 0)
20922124
writer.write_ascii(b'ascii', 0)
20932125
writer.write_widechar(b'wstr', 0)
20942126
writer.write_ucs4(b'ucs4', 0)
2095-
writer.write_substring('text', 0, 0)
2127+
writer.write_substring('text', 2, 2)
2128+
self.assertEqual(writer.get_buffer(), (None, 127, False))
20962129
self.assertIs(writer.finish(), '')
20972130

20982131
for size in (0, 123):
20992132
for ch in range(256):
21002133
with self.subTest(size=size, ch=ch):
21012134
ch = chr(ch)
2135+
maxchar = (255 if ord(ch) >= 128 else 127)
21022136

2103-
# If the first write is a Latin1 character and no buffer
2104-
# was allocated yet, use the singleton as the read-only
2105-
# buffer
2137+
# PyUnicodeWriter_WriteChar(ch) uses the read-only
2138+
# optimization with the character singleton if ch is a
2139+
# Latin1 character and no buffer was allocated yet.
21062140
writer = self.create_writer(size)
21072141
writer.write_char(ord(ch))
2142+
self.assertEqual(writer.get_buffer(),
2143+
(1, maxchar, True))
21082144
self.assertIs(writer.finish(), ch)
21092145

2110-
# PyUnicodeWriter_Finish() replaces the buffer
2111-
# with the singleton
2146+
# PyUnicodeWriter_Finish() replaces the buffer with the
2147+
# singleton. Use PyUnicodeWriter_WriteSubstring() to avoid
2148+
# the read-only buffer optimization.
21122149
writer = self.create_writer(size)
2113-
# Use PyUnicodeWriter_WriteSubstring() to avoid
2114-
# the read-only buffer optimization
2115-
writer.write_substring(ch + 'xxx', 0, 1)
2150+
writer.write_substring('xxx' + ch + 'y', 3, 4)
2151+
self.assertEqual(writer.get_buffer(),
2152+
(size or 1, maxchar, False))
21162153
self.assertIs(writer.finish(), ch)
21172154

21182155
@unittest.skipUnless(support.Py_DEBUG, 'need debug build (Py_DEBUG)')
@@ -2135,60 +2172,83 @@ def test_detect_overflow(self):
21352172
def test_memory_error(self):
21362173
# Inject MemoryError in PyUnicodeWriter_WriteStr()
21372174
writer = self.create_writer(0)
2138-
writer.write_str("start")
2175+
writer.write_utf8(b"start", -1)
2176+
self.assertEqual(writer.get_buffer(), (5, 127, False))
21392177
with self.assertRaises(MemoryError):
21402178
with support.inject_memory_error_cm():
21412179
# Resize the internal str object
21422180
writer.write_str("s" * 1024)
21432181
writer.write_str(" end")
21442182
self.assertEqual(writer.finish(), "start end")
21452183

2146-
# Inject MemoryError in PyUnicodeWriter_Finish()
2184+
# Inject MemoryError in PyUnicodeWriter_Finish(). Use write_utf8() to
2185+
# allocate a buffer of 1024 character. finish() needs to truncate the
2186+
# buffer to 3 characters.
21472187
writer = self.create_writer(1024)
2148-
writer.write_str("abc")
2188+
writer.write_utf8(b"abc", -1)
2189+
self.assertEqual(writer.get_buffer(), (1024, 127, False))
21492190
with self.assertRaises(MemoryError):
21502191
with support.inject_memory_error_cm():
2151-
# Need to truncate the internal str object
21522192
writer.finish()
21532193

21542194
def test_change_kind(self):
21552195
writer = self.create_writer(0)
2196+
21562197
# Create an ASCII buffer
21572198
writer.write_str('ascii ')
2199+
self.assertEqual(writer.get_buffer()[1], 127)
2200+
21582201
# Change the buffer to UCS1
21592202
writer.write_str('latin1:\xe9 ')
2203+
self.assertEqual(writer.get_buffer()[1], 255)
2204+
21602205
# Change the buffer to UCS2
21612206
writer.write_str('ucs2:\u20ac ')
2207+
self.assertEqual(writer.get_buffer()[1], 0xffff)
2208+
21622209
# Change the buffer to UCS4
21632210
writer.write_str('ucs4:\U0010ffff')
2211+
self.assertEqual(writer.get_buffer()[1], 0x10_ffff)
2212+
21642213
self.assertEqual(writer.finish(),
21652214
'ascii latin1:\xe9 ucs2:\u20ac ucs4:\U0010ffff')
21662215

21672216
def test_readonly_optim(self):
21682217
# Read-only optimization: if the first and only write is a Python str
21692218
# object and no buffer was allocated yet, return the object unchanged
21702219
unique_string = 'unique string'
2171-
writer = self.create_writer(0)
2172-
writer.write_str(unique_string)
2173-
self.assertIs(writer.finish(), unique_string)
2220+
expected = (len(unique_string), 127, True)
2221+
for size in (0, 123):
2222+
with self.subTest(size=size):
2223+
# PyUnicodeWriter_WriteStr() optimization
2224+
writer = self.create_writer(size)
2225+
writer.write_str(unique_string)
2226+
self.assertEqual(writer.get_buffer(), expected)
2227+
self.assertIs(writer.finish(), unique_string)
21742228

2175-
writer = self.create_writer(0)
2176-
writer.write_substring(unique_string, 0, len(unique_string))
2177-
self.assertIs(writer.finish(), unique_string)
2229+
# PyUnicodeWriter_WriteSubstring() optimization
2230+
writer = self.create_writer(size)
2231+
writer.write_substring(unique_string, 0, len(unique_string))
2232+
self.assertEqual(writer.get_buffer(), expected)
2233+
self.assertIs(writer.finish(), unique_string)
21782234

2179-
class MyStr:
2180-
def __str__(self):
2181-
return unique_string
2182-
writer = self.create_writer(0)
2183-
writer.write_str(MyStr())
2184-
self.assertIs(writer.finish(), unique_string)
2235+
# PyUnicodeWriter_WriteStr() optimization
2236+
class MyStr:
2237+
def __str__(self):
2238+
return unique_string
2239+
writer = self.create_writer(size)
2240+
writer.write_str(MyStr())
2241+
self.assertEqual(writer.get_buffer(), expected)
2242+
self.assertIs(writer.finish(), unique_string)
21852243

2186-
class MyRepr:
2187-
def __repr__(self):
2188-
return unique_string
2189-
writer = self.create_writer(0)
2190-
writer.write_repr(MyRepr())
2191-
self.assertIs(writer.finish(), unique_string)
2244+
# PyUnicodeWriter_WriteRepr() optimization
2245+
class MyRepr:
2246+
def __repr__(self):
2247+
return unique_string
2248+
writer = self.create_writer(size)
2249+
writer.write_repr(MyRepr())
2250+
self.assertEqual(writer.get_buffer(), expected)
2251+
self.assertIs(writer.finish(), unique_string)
21922252

21932253

21942254
# Test PyUnicodeWriter_Format()

‎Modules/_testcapi/unicode.c‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -730,6 +730,31 @@ writer_get_pointer(PyObject *self_raw, PyObject *args)
730730
}
731731

732732

733+
static PyObject*
734+
writer_get_buffer(PyObject *self_raw, PyObject *args)
735+
{
736+
WriterObject *self = (WriterObject *)self_raw;
737+
if (writer_check(self) < 0) {
738+
return NULL;
739+
}
740+
741+
_PyUnicodeWriter *writer = (_PyUnicodeWriter*)self->writer;
742+
PyObject *allocated;
743+
Py_UCS4 maxchar;
744+
if (writer->buffer) {
745+
allocated = PyLong_FromSsize_t(PyUnicode_GET_LENGTH(writer->buffer));
746+
maxchar = PyUnicode_MAX_CHAR_VALUE(writer->buffer);
747+
}
748+
else {
749+
allocated = Py_None;
750+
maxchar = writer->min_char;
751+
}
752+
return Py_BuildValue("(NkN)",
753+
allocated, (unsigned long)maxchar,
754+
PyBool_FromLong(writer->readonly));
755+
}
756+
757+
733758
static PyObject*
734759
writer_finish(PyObject *self_raw, PyObject *Py_UNUSED(args))
735760
{
@@ -755,6 +780,7 @@ static PyMethodDef writer_methods[] = {
755780
{"write_substring", _PyCFunction_CAST(writer_write_substring), METH_VARARGS},
756781
{"decodeutf8stateful", _PyCFunction_CAST(writer_decodeutf8stateful), METH_VARARGS},
757782
{"get_pointer", _PyCFunction_CAST(writer_get_pointer), METH_VARARGS},
783+
{"get_buffer", _PyCFunction_CAST(writer_get_buffer), METH_VARARGS},
758784
{"finish", _PyCFunction_CAST(writer_finish), METH_NOARGS},
759785
{NULL, NULL} /* sentinel */
760786
};

‎Objects/unicode_writer.c‎

Lines changed: 27 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -151,10 +151,9 @@ PyUnicodeWriter_Create(Py_ssize_t length)
151151
_PyUnicodeWriter *writer = (_PyUnicodeWriter *)pub_writer;
152152

153153
_PyUnicodeWriter_Init(writer);
154-
if (_PyUnicodeWriter_Prepare(writer, length, 127) < 0) {
155-
PyUnicodeWriter_Discard(pub_writer);
156-
return NULL;
157-
}
154+
// The buffer is created lazily at the first write, except if
155+
// the read-only optimization is used.
156+
writer->min_length = length;
158157
writer->overallocate = 1;
159158

160159
return pub_writer;
@@ -189,9 +188,6 @@ int
189188
_PyUnicodeWriter_PrepareInternal(_PyUnicodeWriter *writer,
190189
Py_ssize_t length, Py_UCS4 maxchar)
191190
{
192-
Py_ssize_t newlen;
193-
PyObject *newbuffer;
194-
195191
assert(length >= 0);
196192
assert(maxchar <= _Py_MAX_UNICODE);
197193

@@ -203,50 +199,50 @@ _PyUnicodeWriter_PrepareInternal(_PyUnicodeWriter *writer,
203199
PyErr_NoMemory();
204200
return -1;
205201
}
206-
newlen = writer->pos + length;
202+
Py_ssize_t alloc = writer->pos + length;
207203

208204
maxchar = Py_MAX(maxchar, writer->min_char);
209205

206+
PyObject *newbuffer;
210207
if (writer->buffer == NULL) {
211208
assert(!writer->readonly);
212-
if (writer->overallocate
213-
&& newlen <= (PY_SSIZE_T_MAX - newlen / OVERALLOCATE_FACTOR)) {
214-
/* overallocate to limit the number of realloc() */
215-
newlen += newlen / OVERALLOCATE_FACTOR;
216-
}
217-
if (newlen < writer->min_length)
218-
newlen = writer->min_length;
209+
// Do not overallocate at the first allocation, but use min_length
210+
if (alloc < writer->min_length)
211+
alloc = writer->min_length;
219212

220-
writer->buffer = PyUnicode_New(newlen, maxchar);
213+
writer->buffer = PyUnicode_New(alloc, maxchar);
221214
if (writer->buffer == NULL)
222215
return -1;
223216
}
224-
else if (newlen > writer->size) {
225-
if (writer->overallocate
226-
&& newlen <= (PY_SSIZE_T_MAX - newlen / OVERALLOCATE_FACTOR)) {
217+
else if (alloc > writer->size) {
218+
// Do not overallocate at the first allocation, but use min_length
219+
int overallocate = (writer->overallocate && !writer->readonly);
220+
if (overallocate
221+
&& alloc <= (PY_SSIZE_T_MAX - alloc / OVERALLOCATE_FACTOR)) {
227222
/* overallocate to limit the number of realloc() */
228-
newlen += newlen / OVERALLOCATE_FACTOR;
223+
alloc += alloc / OVERALLOCATE_FACTOR;
229224
}
230-
if (newlen < writer->min_length)
231-
newlen = writer->min_length;
225+
if (alloc < writer->min_length)
226+
alloc = writer->min_length;
232227

233228
if (maxchar > writer->maxchar || writer->readonly) {
234229
/* resize + widen */
235230
maxchar = Py_MAX(maxchar, writer->maxchar);
236-
newbuffer = PyUnicode_New(newlen, maxchar);
231+
newbuffer = PyUnicode_New(alloc, maxchar);
237232
if (newbuffer == NULL)
238233
return -1;
239234
_PyUnicode_FastCopyCharacters(newbuffer, 0,
240235
writer->buffer, 0, writer->pos);
241-
Py_DECREF(writer->buffer);
242236
writer->readonly = 0;
237+
Py_DECREF(writer->buffer);
238+
writer->buffer = newbuffer;
243239
}
244240
else {
245-
newbuffer = _PyUnicode_ResizeCompact(writer->buffer, newlen);
241+
newbuffer = _PyUnicode_ResizeCompact(writer->buffer, alloc);
246242
if (newbuffer == NULL)
247243
return -1;
244+
writer->buffer = newbuffer;
248245
}
249-
writer->buffer = newbuffer;
250246
}
251247
else if (maxchar > writer->maxchar) {
252248
assert(!writer->readonly);
@@ -310,13 +306,12 @@ _PyUnicodeWriter_WriteStr(_PyUnicodeWriter *writer, PyObject *str)
310306
{
311307
assert(PyUnicode_Check(str));
312308

313-
Py_UCS4 maxchar;
314-
Py_ssize_t len;
315-
316-
len = PyUnicode_GET_LENGTH(str);
317-
if (len == 0)
309+
Py_ssize_t len = PyUnicode_GET_LENGTH(str);
310+
if (len == 0) {
318311
return 0;
319-
maxchar = PyUnicode_MAX_CHAR_VALUE(str);
312+
}
313+
Py_UCS4 maxchar = PyUnicode_MAX_CHAR_VALUE(str);
314+
320315
if (maxchar > writer->maxchar || len > writer->size - writer->pos) {
321316
if (writer->buffer == NULL) {
322317
assert(_PyUnicode_CheckConsistency(str, 1));

0 commit comments

Comments
 (0)
Back | FazBrowse Home | New Git URL