gh-125196: Use PyUnicodeWriter for JSON encoder (#125249)

Replace the private _PyUnicodeWriter with the public PyUnicodeWriter.
This commit is contained in:
Victor Stinner 2024-10-10 15:33:00 +02:00 committed by GitHub
parent 120b891e4d
commit c914212474
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 49 additions and 38 deletions

View File

@ -86,11 +86,11 @@ encoder_dealloc(PyObject *self);
static int static int
encoder_clear(PyEncoderObject *self); encoder_clear(PyEncoderObject *self);
static int static int
encoder_listencode_list(PyEncoderObject *s, _PyUnicodeWriter *writer, PyObject *seq, PyObject *newline_indent); encoder_listencode_list(PyEncoderObject *s, PyUnicodeWriter *writer, PyObject *seq, PyObject *newline_indent);
static int static int
encoder_listencode_obj(PyEncoderObject *s, _PyUnicodeWriter *writer, PyObject *obj, PyObject *newline_indent); encoder_listencode_obj(PyEncoderObject *s, PyUnicodeWriter *writer, PyObject *obj, PyObject *newline_indent);
static int static int
encoder_listencode_dict(PyEncoderObject *s, _PyUnicodeWriter *writer, PyObject *dct, PyObject *newline_indent); encoder_listencode_dict(PyEncoderObject *s, PyUnicodeWriter *writer, PyObject *dct, PyObject *newline_indent);
static PyObject * static PyObject *
_encoded_const(PyObject *obj); _encoded_const(PyObject *obj);
static void static void
@ -1268,38 +1268,39 @@ encoder_call(PyEncoderObject *self, PyObject *args, PyObject *kwds)
{ {
/* Python callable interface to encode_listencode_obj */ /* Python callable interface to encode_listencode_obj */
static char *kwlist[] = {"obj", "_current_indent_level", NULL}; static char *kwlist[] = {"obj", "_current_indent_level", NULL};
PyObject *obj, *result; PyObject *obj;
Py_ssize_t indent_level; Py_ssize_t indent_level;
_PyUnicodeWriter writer;
if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:_iterencode", kwlist, if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:_iterencode", kwlist,
&obj, &indent_level)) &obj, &indent_level))
return NULL; return NULL;
_PyUnicodeWriter_Init(&writer); PyUnicodeWriter *writer = PyUnicodeWriter_Create(0);
writer.overallocate = 1; if (writer == NULL) {
return NULL;
}
PyObject *newline_indent = NULL; PyObject *newline_indent = NULL;
if (self->indent != Py_None) { if (self->indent != Py_None) {
newline_indent = _create_newline_indent(self->indent, indent_level); newline_indent = _create_newline_indent(self->indent, indent_level);
if (newline_indent == NULL) { if (newline_indent == NULL) {
_PyUnicodeWriter_Dealloc(&writer); PyUnicodeWriter_Discard(writer);
return NULL; return NULL;
} }
} }
if (encoder_listencode_obj(self, &writer, obj, newline_indent)) { if (encoder_listencode_obj(self, writer, obj, newline_indent)) {
_PyUnicodeWriter_Dealloc(&writer); PyUnicodeWriter_Discard(writer);
Py_XDECREF(newline_indent); Py_XDECREF(newline_indent);
return NULL; return NULL;
} }
Py_XDECREF(newline_indent); Py_XDECREF(newline_indent);
result = PyTuple_New(1); PyObject *str = PyUnicodeWriter_Finish(writer);
if (result == NULL || if (str == NULL) {
PyTuple_SetItem(result, 0, _PyUnicodeWriter_Finish(&writer)) < 0) {
Py_XDECREF(result);
return NULL; return NULL;
} }
PyObject *result = PyTuple_Pack(1, str);
Py_DECREF(str);
return result; return result;
} }
@ -1370,16 +1371,16 @@ encoder_encode_string(PyEncoderObject *s, PyObject *obj)
} }
static int static int
_steal_accumulate(_PyUnicodeWriter *writer, PyObject *stolen) _steal_accumulate(PyUnicodeWriter *writer, PyObject *stolen)
{ {
/* Append stolen and then decrement its reference count */ /* Append stolen and then decrement its reference count */
int rval = _PyUnicodeWriter_WriteStr(writer, stolen); int rval = PyUnicodeWriter_WriteStr(writer, stolen);
Py_DECREF(stolen); Py_DECREF(stolen);
return rval; return rval;
} }
static int static int
encoder_listencode_obj(PyEncoderObject *s, _PyUnicodeWriter *writer, encoder_listencode_obj(PyEncoderObject *s, PyUnicodeWriter *writer,
PyObject *obj, PyObject *newline_indent) PyObject *obj, PyObject *newline_indent)
{ {
/* Encode Python object obj to a JSON term */ /* Encode Python object obj to a JSON term */
@ -1387,13 +1388,13 @@ encoder_listencode_obj(PyEncoderObject *s, _PyUnicodeWriter *writer,
int rv; int rv;
if (obj == Py_None) { if (obj == Py_None) {
return _PyUnicodeWriter_WriteASCIIString(writer, "null", 4); return PyUnicodeWriter_WriteUTF8(writer, "null", 4);
} }
else if (obj == Py_True) { else if (obj == Py_True) {
return _PyUnicodeWriter_WriteASCIIString(writer, "true", 4); return PyUnicodeWriter_WriteUTF8(writer, "true", 4);
} }
else if (obj == Py_False) { else if (obj == Py_False) {
return _PyUnicodeWriter_WriteASCIIString(writer, "false", 5); return PyUnicodeWriter_WriteUTF8(writer, "false", 5);
} }
else if (PyUnicode_Check(obj)) { else if (PyUnicode_Check(obj)) {
PyObject *encoded = encoder_encode_string(s, obj); PyObject *encoded = encoder_encode_string(s, obj);
@ -1402,6 +1403,10 @@ encoder_listencode_obj(PyEncoderObject *s, _PyUnicodeWriter *writer,
return _steal_accumulate(writer, encoded); return _steal_accumulate(writer, encoded);
} }
else if (PyLong_Check(obj)) { else if (PyLong_Check(obj)) {
if (PyLong_CheckExact(obj)) {
// Fast-path for exact integers
return PyUnicodeWriter_WriteRepr(writer, obj);
}
PyObject *encoded = PyLong_Type.tp_repr(obj); PyObject *encoded = PyLong_Type.tp_repr(obj);
if (encoded == NULL) if (encoded == NULL)
return -1; return -1;
@ -1478,7 +1483,7 @@ encoder_listencode_obj(PyEncoderObject *s, _PyUnicodeWriter *writer,
} }
static int static int
encoder_encode_key_value(PyEncoderObject *s, _PyUnicodeWriter *writer, bool *first, encoder_encode_key_value(PyEncoderObject *s, PyUnicodeWriter *writer, bool *first,
PyObject *dct, PyObject *key, PyObject *value, PyObject *dct, PyObject *key, PyObject *value,
PyObject *newline_indent, PyObject *newline_indent,
PyObject *item_separator) PyObject *item_separator)
@ -1518,7 +1523,7 @@ encoder_encode_key_value(PyEncoderObject *s, _PyUnicodeWriter *writer, bool *fir
*first = false; *first = false;
} }
else { else {
if (_PyUnicodeWriter_WriteStr(writer, item_separator) < 0) { if (PyUnicodeWriter_WriteStr(writer, item_separator) < 0) {
Py_DECREF(keystr); Py_DECREF(keystr);
return -1; return -1;
} }
@ -1533,7 +1538,7 @@ encoder_encode_key_value(PyEncoderObject *s, _PyUnicodeWriter *writer, bool *fir
if (_steal_accumulate(writer, encoded) < 0) { if (_steal_accumulate(writer, encoded) < 0) {
return -1; return -1;
} }
if (_PyUnicodeWriter_WriteStr(writer, s->key_separator) < 0) { if (PyUnicodeWriter_WriteStr(writer, s->key_separator) < 0) {
return -1; return -1;
} }
if (encoder_listencode_obj(s, writer, value, newline_indent) < 0) { if (encoder_listencode_obj(s, writer, value, newline_indent) < 0) {
@ -1544,7 +1549,7 @@ encoder_encode_key_value(PyEncoderObject *s, _PyUnicodeWriter *writer, bool *fir
} }
static int static int
encoder_listencode_dict(PyEncoderObject *s, _PyUnicodeWriter *writer, encoder_listencode_dict(PyEncoderObject *s, PyUnicodeWriter *writer,
PyObject *dct, PyObject *newline_indent) PyObject *dct, PyObject *newline_indent)
{ {
/* Encode Python dict dct a JSON term */ /* Encode Python dict dct a JSON term */
@ -1555,8 +1560,10 @@ encoder_listencode_dict(PyEncoderObject *s, _PyUnicodeWriter *writer,
PyObject *new_newline_indent = NULL; PyObject *new_newline_indent = NULL;
PyObject *separator_indent = NULL; PyObject *separator_indent = NULL;
if (PyDict_GET_SIZE(dct) == 0) /* Fast path */ if (PyDict_GET_SIZE(dct) == 0) {
return _PyUnicodeWriter_WriteASCIIString(writer, "{}", 2); /* Fast path */
return PyUnicodeWriter_WriteUTF8(writer, "{}", 2);
}
if (s->markers != Py_None) { if (s->markers != Py_None) {
int has_key; int has_key;
@ -1574,8 +1581,9 @@ encoder_listencode_dict(PyEncoderObject *s, _PyUnicodeWriter *writer,
} }
} }
if (_PyUnicodeWriter_WriteChar(writer, '{')) if (PyUnicodeWriter_WriteChar(writer, '{')) {
goto bail; goto bail;
}
PyObject *current_item_separator = s->item_separator; // borrowed reference PyObject *current_item_separator = s->item_separator; // borrowed reference
if (s->indent != Py_None) { if (s->indent != Py_None) {
@ -1589,7 +1597,7 @@ encoder_listencode_dict(PyEncoderObject *s, _PyUnicodeWriter *writer,
} }
// update item separator with a borrowed reference // update item separator with a borrowed reference
current_item_separator = separator_indent; current_item_separator = separator_indent;
if (_PyUnicodeWriter_WriteStr(writer, new_newline_indent) < 0) { if (PyUnicodeWriter_WriteStr(writer, new_newline_indent) < 0) {
goto bail; goto bail;
} }
} }
@ -1635,13 +1643,14 @@ encoder_listencode_dict(PyEncoderObject *s, _PyUnicodeWriter *writer,
Py_CLEAR(new_newline_indent); Py_CLEAR(new_newline_indent);
Py_CLEAR(separator_indent); Py_CLEAR(separator_indent);
if (_PyUnicodeWriter_WriteStr(writer, newline_indent) < 0) { if (PyUnicodeWriter_WriteStr(writer, newline_indent) < 0) {
goto bail; goto bail;
} }
} }
if (_PyUnicodeWriter_WriteChar(writer, '}')) if (PyUnicodeWriter_WriteChar(writer, '}')) {
goto bail; goto bail;
}
return 0; return 0;
bail: bail:
@ -1653,7 +1662,7 @@ bail:
} }
static int static int
encoder_listencode_list(PyEncoderObject *s, _PyUnicodeWriter *writer, encoder_listencode_list(PyEncoderObject *s, PyUnicodeWriter *writer,
PyObject *seq, PyObject *newline_indent) PyObject *seq, PyObject *newline_indent)
{ {
PyObject *ident = NULL; PyObject *ident = NULL;
@ -1668,7 +1677,7 @@ encoder_listencode_list(PyEncoderObject *s, _PyUnicodeWriter *writer,
return -1; return -1;
if (PySequence_Fast_GET_SIZE(s_fast) == 0) { if (PySequence_Fast_GET_SIZE(s_fast) == 0) {
Py_DECREF(s_fast); Py_DECREF(s_fast);
return _PyUnicodeWriter_WriteASCIIString(writer, "[]", 2); return PyUnicodeWriter_WriteUTF8(writer, "[]", 2);
} }
if (s->markers != Py_None) { if (s->markers != Py_None) {
@ -1687,8 +1696,9 @@ encoder_listencode_list(PyEncoderObject *s, _PyUnicodeWriter *writer,
} }
} }
if (_PyUnicodeWriter_WriteChar(writer, '[')) if (PyUnicodeWriter_WriteChar(writer, '[')) {
goto bail; goto bail;
}
PyObject *separator = s->item_separator; // borrowed reference PyObject *separator = s->item_separator; // borrowed reference
if (s->indent != Py_None) { if (s->indent != Py_None) {
@ -1697,7 +1707,7 @@ encoder_listencode_list(PyEncoderObject *s, _PyUnicodeWriter *writer,
goto bail; goto bail;
} }
if (_PyUnicodeWriter_WriteStr(writer, new_newline_indent) < 0) { if (PyUnicodeWriter_WriteStr(writer, new_newline_indent) < 0) {
goto bail; goto bail;
} }
@ -1710,7 +1720,7 @@ encoder_listencode_list(PyEncoderObject *s, _PyUnicodeWriter *writer,
for (i = 0; i < PySequence_Fast_GET_SIZE(s_fast); i++) { for (i = 0; i < PySequence_Fast_GET_SIZE(s_fast); i++) {
PyObject *obj = PySequence_Fast_GET_ITEM(s_fast, i); PyObject *obj = PySequence_Fast_GET_ITEM(s_fast, i);
if (i) { if (i) {
if (_PyUnicodeWriter_WriteStr(writer, separator) < 0) if (PyUnicodeWriter_WriteStr(writer, separator) < 0)
goto bail; goto bail;
} }
if (encoder_listencode_obj(s, writer, obj, new_newline_indent)) { if (encoder_listencode_obj(s, writer, obj, new_newline_indent)) {
@ -1727,13 +1737,14 @@ encoder_listencode_list(PyEncoderObject *s, _PyUnicodeWriter *writer,
if (s->indent != Py_None) { if (s->indent != Py_None) {
Py_CLEAR(new_newline_indent); Py_CLEAR(new_newline_indent);
Py_CLEAR(separator_indent); Py_CLEAR(separator_indent);
if (_PyUnicodeWriter_WriteStr(writer, newline_indent) < 0) { if (PyUnicodeWriter_WriteStr(writer, newline_indent) < 0) {
goto bail; goto bail;
} }
} }
if (_PyUnicodeWriter_WriteChar(writer, ']')) if (PyUnicodeWriter_WriteChar(writer, ']')) {
goto bail; goto bail;
}
Py_DECREF(s_fast); Py_DECREF(s_fast);
return 0; return 0;