Skip to content

Commit 8542958

Browse files
authored
gh-157710: Add _PyUnicodeWriter_CanWrite() function (#157712)
Make sure that a writer can be modified before writing into it. * Document that PyUnicodeWriter is not thread safe. * Add singletons tests to test_capi.test_unicode. * Check PyUnicode_CheckExact() earlier in _PyUnicode_IsModifiable().
1 parent 5539c2a commit 8542958

6 files changed

Lines changed: 61 additions & 7 deletions

File tree

Doc/c-api/unicode.rst

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1797,6 +1797,9 @@ object.
17971797
The instance must be destroyed by :c:func:`PyUnicodeWriter_Finish` on
17981798
success, or :c:func:`PyUnicodeWriter_Discard` on error.
17991799
1800+
The API is **not thread safe**. To share a writer with multiple threads, a
1801+
critical section or a lock is needed.
1802+
18001803
.. c:function:: PyUnicodeWriter* PyUnicodeWriter_Create(Py_ssize_t length)
18011804
18021805
Create a Unicode writer instance.

Include/internal/pycore_unicodeobject.h

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,9 @@ extern "C" {
1616
#define _Py_MAX_UNICODE 0x10ffff
1717

1818

19-
extern int _PyUnicode_IsModifiable(PyObject *unicode);
19+
// Export for '_multibytecodec' shared extension. _PyUnicodeWriter_CanWrite()
20+
// calls this function when assertions are enabled.
21+
PyAPI_FUNC(int) _PyUnicode_IsModifiable(PyObject *unicode);
2022
extern void _PyUnicodeWriter_InitWithBuffer(
2123
_PyUnicodeWriter *writer,
2224
PyObject *buffer);
@@ -105,12 +107,31 @@ _PyUnicode_EnsureUnicode(PyObject *obj)
105107
return 0;
106108
}
107109

110+
#ifndef NDEBUG
111+
static inline int
112+
_PyUnicodeWriter_CanWrite(_PyUnicodeWriter *writer)
113+
{
114+
// Code adapted from _PyUnicode_IsModifiable()
115+
assert(!writer->readonly);
116+
PyObject *buffer = writer->buffer;
117+
assert(buffer != NULL);
118+
// Do not use _PyObject_IsUniquelyReferenced(): the caller can have its own
119+
// lock to prevent a writer being used by two theads at the same time.
120+
assert(Py_REFCNT(buffer) == 1);
121+
assert(PyUnstable_Unicode_GET_CACHED_HASH(buffer) == -1);
122+
assert(!PyUnicode_CHECK_INTERNED(buffer));
123+
assert(!_Py_IsImmortal(buffer));
124+
return 1;
125+
}
126+
#endif
127+
108128
static inline int
109129
_PyUnicodeWriter_WriteCharInline(_PyUnicodeWriter *writer, Py_UCS4 ch)
110130
{
111131
assert(ch <= _Py_MAX_UNICODE);
112132
if (_PyUnicodeWriter_Prepare(writer, 1, ch) < 0)
113133
return -1;
134+
assert(_PyUnicodeWriter_CanWrite(writer));
114135
PyUnicode_WRITE(writer->kind, writer->data, writer->pos, ch);
115136
writer->pos++;
116137
return 0;

Lib/test/test_capi/test_unicode.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1981,6 +1981,17 @@ def test_substring_empty(self):
19811981
writer.write_substring("abc", 1, 1)
19821982
self.assertEqual(writer.finish(), '')
19831983

1984+
def test_singletons(self):
1985+
writer = self.create_writer(5)
1986+
self.assertIs(writer.finish(), '')
1987+
1988+
for ch in range(256):
1989+
with self.subTest(ch=ch):
1990+
ch = chr(ch)
1991+
writer = self.create_writer(0)
1992+
writer.write_substring(ch + 'xxx', 0, 1)
1993+
self.assertIs(writer.finish(), ch)
1994+
19841995

19851996
@unittest.skipIf(ctypes is None, 'need ctypes')
19861997
class PyUnicodeWriterFormatTest(unittest.TestCase):

Objects/longobject.c

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2220,6 +2220,7 @@ long_to_decimal_string_internal(PyObject *aa,
22202220
Py_DECREF(scratch);
22212221
return -1;
22222222
}
2223+
assert(_PyUnicodeWriter_CanWrite(writer));
22232224
}
22242225
else if (bytes_writer) {
22252226
*bytes_str = PyBytesWriter_GrowAndUpdatePointer(bytes_writer, strlen,
@@ -2390,8 +2391,10 @@ long_format_binary(PyObject *aa, int base, int alternate,
23902391
}
23912392

23922393
if (writer) {
2393-
if (_PyUnicodeWriter_Prepare(writer, sz, 'x') == -1)
2394+
if (_PyUnicodeWriter_Prepare(writer, sz, 'x') == -1) {
23942395
return -1;
2396+
}
2397+
assert(_PyUnicodeWriter_CanWrite(writer));
23952398
}
23962399
else if (bytes_writer) {
23972400
*bytes_str = PyBytesWriter_GrowAndUpdatePointer(bytes_writer, sz,

Objects/unicode_writer.c

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -350,6 +350,8 @@ _PyUnicodeWriter_WriteStr(_PyUnicodeWriter *writer, PyObject *str)
350350
if (_PyUnicodeWriter_PrepareInternal(writer, len, maxchar) == -1)
351351
return -1;
352352
}
353+
354+
assert(_PyUnicodeWriter_CanWrite(writer));
353355
_PyUnicode_FastCopyCharacters(writer->buffer, writer->pos,
354356
str, 0, len);
355357
writer->pos += len;
@@ -428,6 +430,7 @@ _PyUnicodeWriter_WriteSubstring(_PyUnicodeWriter *writer, PyObject *str,
428430
if (_PyUnicodeWriter_Prepare(writer, len, maxchar) < 0) {
429431
return -1;
430432
}
433+
assert(_PyUnicodeWriter_CanWrite(writer));
431434

432435
_PyUnicode_FastCopyCharacters(writer->buffer, writer->pos,
433436
str, start, len);
@@ -485,8 +488,10 @@ _PyUnicodeWriter_WriteASCIIString(_PyUnicodeWriter *writer,
485488
return 0;
486489
}
487490

488-
if (_PyUnicodeWriter_Prepare(writer, len, 127) == -1)
491+
if (_PyUnicodeWriter_Prepare(writer, len, 127) == -1) {
489492
return -1;
493+
}
494+
assert(_PyUnicodeWriter_CanWrite(writer));
490495

491496
switch (writer->kind)
492497
{
@@ -591,6 +596,7 @@ _PyUnicodeWriter_WriteLatin1String(_PyUnicodeWriter *writer,
591596
maxchar = ucs1lib_find_max_char((const Py_UCS1*)str, (const Py_UCS1*)str + len);
592597
if (_PyUnicodeWriter_Prepare(writer, len, maxchar) == -1)
593598
return -1;
599+
assert(_PyUnicodeWriter_CanWrite(writer));
594600
unicode_write_cstr(writer->buffer, writer->pos, str, len);
595601
writer->pos += len;
596602
return 0;

Objects/unicodeobject.c

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1743,18 +1743,21 @@ unicode_is_singleton(PyObject *unicode)
17431743
}
17441744
#endif
17451745

1746+
// If this function is updated, update also _PyUnicodeWriter_CanWrite().
17461747
int
17471748
_PyUnicode_IsModifiable(PyObject *unicode)
17481749
{
17491750
assert(_PyUnicode_CHECK(unicode));
1751+
if (!PyUnicode_CheckExact(unicode))
1752+
return 0;
1753+
// On Free Threading, this test fails if called from a thread other
1754+
// than the one which created the str object.
17501755
if (!_PyObject_IsUniquelyReferenced(unicode))
17511756
return 0;
17521757
if (PyUnicode_HASH(unicode) != -1)
17531758
return 0;
17541759
if (PyUnicode_CHECK_INTERNED(unicode))
17551760
return 0;
1756-
if (!PyUnicode_CheckExact(unicode))
1757-
return 0;
17581761
#ifdef Py_DEBUG
17591762
/* singleton refcount is greater than 1 */
17601763
assert(!unicode_is_singleton(unicode));
@@ -2008,6 +2011,7 @@ PyUnicodeWriter_WriteWideChar(PyUnicodeWriter *pub_writer,
20082011
if (_PyUnicodeWriter_Prepare(writer, size - num_surrogates, maxchar) < 0) {
20092012
return -1;
20102013
}
2014+
assert(_PyUnicodeWriter_CanWrite(writer));
20112015

20122016
int kind = writer->kind;
20132017
void *data = (Py_UCS1*)writer->data + writer->pos * kind;
@@ -2266,6 +2270,7 @@ PyUnicodeWriter_WriteUCS4(PyUnicodeWriter *pub_writer,
22662270
if (_PyUnicodeWriter_Prepare(writer, size, max_char) < 0) {
22672271
return -1;
22682272
}
2273+
assert(_PyUnicodeWriter_CanWrite(writer));
22692274

22702275
int kind = writer->kind;
22712276
void *data = (Py_UCS1*)writer->data + writer->pos * kind;
@@ -2552,8 +2557,10 @@ unicode_fromformat_write_str(_PyUnicodeWriter *writer, PyObject *str,
25522557
else
25532558
maxchar = writer->maxchar;
25542559

2555-
if (_PyUnicodeWriter_Prepare(writer, arglen, maxchar) == -1)
2560+
if (_PyUnicodeWriter_Prepare(writer, arglen, maxchar) == -1) {
25562561
return -1;
2562+
}
2563+
assert(_PyUnicodeWriter_CanWrite(writer));
25572564

25582565
fill = Py_MAX(width - length, 0);
25592566
if (fill && !(flags & F_LJUST)) {
@@ -2843,8 +2850,10 @@ unicode_fromformat_arg(_PyUnicodeWriter *writer,
28432850
Py_ssize_t spacepad = Py_MAX(width - precision - sign, 0);
28442851
Py_ssize_t zeropad = Py_MAX(precision - len, 0);
28452852

2846-
if (_PyUnicodeWriter_Prepare(writer, width, 127) == -1)
2853+
if (_PyUnicodeWriter_Prepare(writer, width, 127) == -1) {
28472854
return NULL;
2855+
}
2856+
assert(_PyUnicodeWriter_CanWrite(writer));
28482857

28492858
if (spacepad && !(flags & F_LJUST)) {
28502859
if (PyUnicode_Fill(writer->buffer, writer->pos, spacepad, ' ') == -1)
@@ -5371,6 +5380,7 @@ _PyUnicode_DecodeUTF8Writer(_PyUnicodeWriter *writer,
53715380
if (_PyUnicodeWriter_Prepare(writer, size, 127) < 0) {
53725381
return -1;
53735382
}
5383+
assert(_PyUnicodeWriter_CanWrite(writer));
53745384

53755385
const char *starts = s;
53765386
const char *end = s + size;

0 commit comments

Comments
 (0)