diff --git a/Include/internal/pycore_unicodeobject.h b/Include/internal/pycore_unicodeobject.h index 3b930e857e3704..112337d90968b3 100644 --- a/Include/internal/pycore_unicodeobject.h +++ b/Include/internal/pycore_unicodeobject.h @@ -189,7 +189,8 @@ _PyUnicodeWriter_WriteCharInline(_PyUnicodeWriter *writer, Py_UCS4 ch) return 0; } -extern PyObject* _PyUnicodeWriter_FinishWithSize( +// Export for '_testcapi' shared extension +PyAPI_FUNC(PyObject*) _PyUnicodeWriter_FinishWithSize( _PyUnicodeWriter *writer, Py_ssize_t size); diff --git a/Lib/test/test_capi/test_unicode.py b/Lib/test/test_capi/test_unicode.py index 6965b9dc11b57d..9c91d6eb512d2a 100644 --- a/Lib/test/test_capi/test_unicode.py +++ b/Lib/test/test_capi/test_unicode.py @@ -2322,6 +2322,19 @@ def test_readonly_optim_large_int(self): self.assertEqual(writer.get_buffer(), expected) self.assertEqual(writer.finish(), large_int_str) + def test_finish_with_size(self): + # Test _PyUnicodeWriter_FinishWithSize(). Truncate text requires + # to change the buffer kind. + text = 'a\xff\u20ac\U0010ffff' + expected = (0x10ffff, False) # do not test the buffer size + for size in range(len(text) + 1): + writer = self.create_writer(0) + writer.write_utf8(text.encode(), -1) + self.assertEqual(writer.get_buffer()[1:], expected) + self.assertEqual(writer.finish_with_size(size), text[:size]) + + # CRASHES writer.finish_with_size(len(text) + 1) + # Test PyUnicodeWriter_Format() @unittest.skipIf(ctypes is None, 'need ctypes') diff --git a/Modules/_testcapi/unicode.c b/Modules/_testcapi/unicode.c index ba9b205e07ef8b..5138fbab1f4ed6 100644 --- a/Modules/_testcapi/unicode.c +++ b/Modules/_testcapi/unicode.c @@ -795,6 +795,31 @@ writer_finish(PyObject *self_raw, PyObject *Py_UNUSED(args)) } +static PyObject* +writer_finish_with_size(PyObject *self_raw, PyObject *args) +{ + PyAPI_FUNC(PyObject*) _PyUnicodeWriter_FinishWithSize( + _PyUnicodeWriter *writer, + Py_ssize_t size); + + WriterObject *self = (WriterObject *)self_raw; + if (writer_check(self) < 0) { + return NULL; + } + + Py_ssize_t size; + if (!PyArg_ParseTuple(args, "n", &size)) { + return NULL; + } + + _PyUnicodeWriter *writer = (_PyUnicodeWriter*)self->writer; + PyObject *str = _PyUnicodeWriter_FinishWithSize(writer, size); + PyUnicodeWriter_Discard(self->writer); + self->writer = NULL; + return str; +} + + static PyMethodDef writer_methods[] = { {"write_char", _PyCFunction_CAST(writer_write_char), METH_VARARGS}, {"write_utf8", _PyCFunction_CAST(writer_write_utf8), METH_VARARGS}, @@ -809,6 +834,7 @@ static PyMethodDef writer_methods[] = { {"get_pointer", _PyCFunction_CAST(writer_get_pointer), METH_VARARGS}, {"get_buffer", _PyCFunction_CAST(writer_get_buffer), METH_VARARGS}, {"finish", _PyCFunction_CAST(writer_finish), METH_NOARGS}, + {"finish_with_size", _PyCFunction_CAST(writer_finish_with_size), METH_VARARGS}, {NULL, NULL} /* sentinel */ };