[C API] Add an efficient public PyUnicodeWriter API

[vstinner]

Feature or enhancement

Creating a Python string object in an efficient way is complicated. Python has private _PyUnicodeWriter API. It's being used by these projects:

Affected projects (5):

• Cython (3.0.9)
• asyncpg (0.29.0)
• catboost (1.2.3)
• frozendict (2.4.0)
• immutables (0.20)

I propose making the API public to promote it and help C extensions maintainers to write more efficient code to create Python string objects.

API:

typedef struct PyUnicodeWriter PyUnicodeWriter;

PyAPI_FUNC(PyUnicodeWriter*) PyUnicodeWriter_Create(void);
PyAPI_FUNC(void) PyUnicodeWriter_Discard(PyUnicodeWriter *writer);
PyAPI_FUNC(PyObject*) PyUnicodeWriter_Finish(PyUnicodeWriter *writer);

PyAPI_FUNC(void) PyUnicodeWriter_SetOverallocate(
    PyUnicodeWriter *writer,
    int overallocate);

PyAPI_FUNC(int) PyUnicodeWriter_WriteChar(
    PyUnicodeWriter *writer,
    Py_UCS4 ch);
PyAPI_FUNC(int) PyUnicodeWriter_WriteUTF8(
    PyUnicodeWriter *writer,
    const char *str,  // decoded from UTF-8
    Py_ssize_t len);  // use strlen() if len < 0
PyAPI_FUNC(int) PyUnicodeWriter_Format(
    PyUnicodeWriter *writer,
    const char *format,
    ...);

// Write str(obj)
PyAPI_FUNC(int) PyUnicodeWriter_WriteStr(
    PyUnicodeWriter *writer,
    PyObject *obj);

// Write repr(obj)
PyAPI_FUNC(int) PyUnicodeWriter_WriteRepr(
    PyUnicodeWriter *writer,
    PyObject *obj);

// Write str[start:end]
PyAPI_FUNC(int) PyUnicodeWriter_WriteSubstring(
    PyUnicodeWriter *writer,
    PyObject *str,
    Py_ssize_t start,
    Py_ssize_t end);

The internal writer buffer is overallocated by default. PyUnicodeWriter_Finish() truncates the buffer to the exact size if the buffer was overallocated.

Overallocation reduces the cost of exponential complexity when adding short strings in a loop. Use PyUnicodeWriter_SetOverallocate(writer, 0) to disable overallocation just before the last write.

The writer takes care of the internal buffer kind: Py_UCS1 (latin1), Py_UCS2 (BMP) or Py_UCS4 (full Unicode Character Set). It also implements an optimization if a single write is made using PyUnicodeWriter_WriteStr(): it returns the string unchanged without any copy.

---

Example of usage (simplified code from Python/unionobject.c):

static PyObject *
union_repr(PyObject *self)
{
    unionobject *alias = (unionobject *)self;
    Py_ssize_t len = PyTuple_GET_SIZE(alias->args);

    PyUnicodeWriter *writer = PyUnicodeWriter_Create();
    if (writer == NULL) {
        return NULL;
    }

    for (Py_ssize_t i = 0; i < len; i++) {
        if (i > 0 && PyUnicodeWriter_WriteUTF8(writer, " | ", 3) < 0) {
            goto error;
        }
        PyObject *p = PyTuple_GET_ITEM(alias->args, i);
        if (PyUnicodeWriter_WriteRepr(writer, p) < 0) {
            goto error;
        }
    }
    return PyUnicodeWriter_Finish(writer);

error:
    PyUnicodeWriter_Discard(writer);
    return NULL;
}

Linked PRs

• gh-119182: Add PyUnicodeWriter C API #119184
• gh-119396: Optimize PyUnicode_FromFormat() UTF-8 decoder #119398
• gh-119182: Decode PyUnicode_FromFormat() format string from UTF-8 #120248
• gh-119182: Use strict error handler in PyUnicode_FromFormat() #120307
• gh-119182: Add PyUnicodeWriter_DecodeUTF8Stateful() #120639
• gh-119182: Optimize PyUnicode_FromFormat() #120796
• gh-119182: Use public PyUnicodeWriter API in union_repr() #120797
• gh-119182: Use public PyUnicodeWriter API in ga_repr() #120799
• gh-119182: Use public PyUnicodeWriter in contextvar_tp_repr() #120809
• gh-119182: Rewrite PyUnicodeWriter tests in Python #120845
• gh-119182: Add PyUnicodeWriter_WriteUCS4() function #120849
• gh-119182: Use PyUnicodeWriter_WriteWideChar() #120851
• gh-119182: Add checks to PyUnicodeWriter APIs #120870
• gh-119182: Complete PyUnicodeWriter documentation #127607
• gh-119182: Use public PyUnicodeWriter in wrap_strftime() #129206
• gh-119182: Use public PyUnicodeWriter in time_strftime() #129207
• gh-119182: Use public PyUnicodeWriter in ast_unparse.c #129208
• gh-119182: Use public PyUnicodeWriter in Python-ast.c #129209
• gh-119182: Use public PyUnicodeWriter in stringio.c #129243
• gh-119182: Use public PyUnicodeWriter in _json.c #129249

[vstinner]

Benchmark using:

• bench.patch
• bench.py

bench_concat: Mean +- std dev: 2.07 us +- 0.03 us
bench_writer: Mean +- std dev: 894 ns +- 13 ns

PyUnicodeWriter is 2.3x faster than PyUnicode_Concat()+PyUnicode_Append().

The difference comes from overallocation: if I add PyUnicodeWriter_SetOverallocate(writer, 0); after PyUnicodeWriter_Create(), PyUnicodeWriter has the same performance than PyUnicode_Concat()+PyUnicode_Append(). Overallocation avoids str += str quadratic complexity (well, at least, it reduces the complexity).

The PyUnicodeWriter API makes overallocation easy to use.

cc @serhiy-storchaka

[vstinner]

By the way, PyPy provides __pypy__.builders.StringBuilder for "Fast String Concatenation": https://doc.pypy.org/en/latest/__pypy__-module.html#fast-string-concatenation to work around the str += str quadratic complexity.

$ pypy3.9 
>>>> import __pypy__
>>>> b=__pypy__.builders.StringBuilder()
>>>> b.append('x')
>>>> b.append('=')
>>>> b.append('value')
>>>> b.build()
'x=value'

[vstinner]

Article about this performance problem in Python: https://lwn.net/Articles/816415/

[gvanrossum]

Curious if this warrants a further API PyUnicodeWriter_WriteStr(writer, obj) which appends repr(obj) (just as WriteStr(writer, obj) can be seen to append str(obj)), and eventually the development of a new type slot that writes the repr or str of an object to a writer rather than returning a string object. (And maybe even an "WriteAscii" to write ascii(obj)) and WriteFormat to do something with formats. :-)

I know, I know, hyper-generalization, yet this is what the Union example is screaming for... I suppose we can add those later.

How long has the internal writer API existed?

Would these be in the Stable ABI / Limited API from the start? (API-wise these look stable.)

[vstinner]

Curious if this warrants a further API PyUnicodeWriter_WriteStr(writer, obj) which appends repr(obj)

I suppose that you mean PyUnicodeWriter_WriteRepr().

Curious if this warrants a further API PyUnicodeWriter_WriteStr(writer, obj) which appends repr(obj) (just as WriteStr(writer, obj) can be seen to append str(obj)), and eventually the development of a new type slot that writes the repr or str of an object to a writer rather than returning a string object. (And maybe even an "WriteAscii" to write ascii(obj)) and WriteFormat to do something with formats. :-)

There is already a collection of helper function accepting a writer and I find this really cool. It's not "slot-based", since each function has many formatting options.

extern int _PyLong_FormatWriter(
    _PyUnicodeWriter *writer,
    PyObject *obj,
    int base,
    int alternate);

extern int _PyLong_FormatAdvancedWriter(
    _PyUnicodeWriter *writer,
    PyObject *obj,
    PyObject *format_spec,
    Py_ssize_t start,
    Py_ssize_t end);

extern int _PyFloat_FormatAdvancedWriter(
    _PyUnicodeWriter *writer,
    PyObject *obj,
    PyObject *format_spec,
    Py_ssize_t start,
    Py_ssize_t end);

extern int _PyComplex_FormatAdvancedWriter(
    _PyUnicodeWriter *writer,
    PyObject *obj,
    PyObject *format_spec,
    Py_ssize_t start,
    Py_ssize_t end);

extern int _PyUnicode_FormatAdvancedWriter(
    _PyUnicodeWriter *writer,
    PyObject *obj,
    PyObject *format_spec,
    Py_ssize_t start,
    Py_ssize_t end);

extern Py_ssize_t _PyUnicode_InsertThousandsGrouping(
    _PyUnicodeWriter *writer,
    Py_ssize_t n_buffer,
    PyObject *digits,
    Py_ssize_t d_pos,
    Py_ssize_t n_digits,
    Py_ssize_t min_width,
    const char *grouping,
    PyObject *thousands_sep,
    Py_UCS4 *maxchar);

These functions avoid memory copies. For example, _PyLong_FormatWriter() writes directly digits in the writter buffer, without the need of a temporary buffer.

How long has the internal writer API existed?

12 years: I added it in 2012.

commit 202fdca133ce8f5b0c37cca1353070e0721c688d
Author: Victor Stinner <victor.stinner@gmail.com>
Date:   Mon May 7 12:47:02 2012 +0200

    Close #14716: str.format() now uses the new "unicode writer" API instead of the
    PyAccu API. For example, it makes str.format() from 25% to 30% faster on Linux.

I wrote this API to fix the major performance regression after PEP 393 – Flexible String Representation was implemented. After my optimization work, many string operations on Unicode objects became faster than Python 2 operations on bytes! Especially when treating only ASCII characters which is the most common case. I mostly optimized str.format() and str % args where are powerful but complex.

In 2016, I wrote an article about the two "writer" APIs that I wrote to optimize: https://vstinner.github.io/pybyteswriter.html

Would these be in the Stable ABI / Limited API from the start? (API-wise these look stable.)

I would prefer to not add it to the limited C API directly, but wait one Python version to see how it goes.

[gvanrossum]

(Yes, I meant WriteRepr.) I like these other helpers -- can we just add them all to the public API? Or are there issues with any of them?

[vstinner]

(Yes, I meant WriteRepr.) I like these other helpers -- can we just add them all to the public API? Or are there issues with any of them?

I added the following function which should fit most of these use cases:

PyAPI_FUNC(int) PyUnicodeWriter_FromFormat(
    PyUnicodeWriter *writer,
    const char *format,
    ...);

Example to write repr(obj):

PyUnicodeWriter_FromFormat(writer, "%R", obj);

Example to write str(obj):

PyUnicodeWriter_FromFormat(writer, "%S", obj);

It's the same format than PyUnicode_FromFormat(). Example:

PyUnicodeWriter_FromFormat(writer, "Hello %s, %i.", "Python", 123);

[encukou]

Thank you, this looks very useful!

I see that PyUnicodeWriter_Finish frees the writer. That's great; it allows optimizations we can also use in other writers/builders in the future. (Those should have a consistent API.)
One thing to note is that PyUnicodeWriter_Finish should free the writer even when an error occurs.
Maybe PyUnicodeWriter_Free should be named e.g. PyUnicodeWriter_Discard to emphasize that you should only call it if you didn't Finish.

The va_arg function is problematic for non-C languages, but it's possible to get the functionality with other functions – especially if we add a number-writing helper, so I'm OK with adding it.

The proposed API is nice and minimal. My bet about what users will ask for next goes to PyUnicodeWriter_WriteUTF8String (for IO) & PyUnicodeWriter_WriteUTF16String (for Windows or Java interop).

Name bikeshedding:

• PyUnicodeWriter_WriteUCS4Char rather than PyUnicodeWriter_WriteChar -- character is an overloaded term, let's be specific.
• PyUnicodeWriter_WriteFormat (or WriteFromFormat?) rather than PyUnicodeWriter_FromFormat -- it's writing, not creating a writer.

---

I see the PR hides underscored API that some existing projects use. I thought we weren't doing that any more.

[vstinner]

PyUnicodeWriter_WriteUCS4Char rather than PyUnicodeWriter_WriteChar -- character is an overloaded term, let's be specific.

"WriteChar" name comes from PyUnicode_ReadChar() and PyUnicode_WriteChar() names. I don't think that mentioning UCS4 is useful.

PyUnicodeWriter_WriteFormat (or WriteFromFormat?) rather than PyUnicodeWriter_FromFormat -- it's writing, not creating a writer.

I would prefer just "PyUnicodeWriter_Format()". I prefer to not support str.format() which is more a "Python API" than a C API. It's less convenient to use in C. If we don't support str.format(), "PyUnicodeWriter_Format()" is fine for the "PyUnicode_FormFormat()" variant.

[encukou]

Yeah, PyUnicodeWriter_Format sounds good. It avoids the PyX_FromY scheme we use for constructing new objects.

I think that using unqualified Char for a UCS4 codepoint was a mistake we shouldn't continue, but I'm happy to be outvoted on that.