pythongh-119182: Use public PyUnicodeWriter API in ga_repr() (python#120799)

vstinner · mrahtz · commit c0cd6da1a77c · 2024-06-30T12:20:49.000+01:00
The public PyUnicodeWriter API enables overallocation by default and
so is more efficient.

Benchmark:

python -m pyperf timeit \
    -s 't = list[int, float, complex, str, bytes, bytearray, ' \
                 'memoryview, list, dict]' \
    'str(t)'

Result:

1.49 us +- 0.03 us -&gt; 1.10 us +- 0.02 us: 1.35x faster
diff --git a/Objects/genericaliasobject.c b/Objects/genericaliasobject.c
@@ -51,16 +51,15 @@ ga_traverse(PyObject *self, visitproc visit, void *arg)
 }
 
 static int
-ga_repr_item(_PyUnicodeWriter *writer, PyObject *p)
+ga_repr_item(PyUnicodeWriter *writer, PyObject *p)
 {
     PyObject *qualname = NULL;
     PyObject *module = NULL;
-    PyObject *r = NULL;
     int rc;
 
     if (p == Py_Ellipsis) {
         // The Ellipsis object
-        r = PyUnicode_FromString("...");
+        rc = PyUnicodeWriter_WriteUTF8(writer, "...", 3);
         goto done;
     }
 
@@ -71,17 +70,17 @@ ga_repr_item(_PyUnicodeWriter *writer, PyObject *p)
         goto use_repr;
     }
     if (rc < 0) {
-        goto done;
+        goto error;
     }
 
     if (PyObject_GetOptionalAttr(p, &_Py_ID(__qualname__), &qualname) < 0) {
-        goto done;
+        goto error;
     }
     if (qualname == NULL) {
         goto use_repr;
     }
     if (PyObject_GetOptionalAttr(p, &_Py_ID(__module__), &module) < 0) {
-        goto done;
+        goto error;
     }
     if (module == NULL || module == Py_None) {
         goto use_repr;
@@ -92,45 +91,42 @@ ga_repr_item(_PyUnicodeWriter *writer, PyObject *p)
         _PyUnicode_EqualToASCIIString(module, "builtins"))
     {
         // builtins don't need a module name
-        r = PyObject_Str(qualname);
+        rc = PyUnicodeWriter_WriteStr(writer, qualname);
         goto done;
     }
     else {
-        r = PyUnicode_FromFormat("%S.%S", module, qualname);
+        rc = PyUnicodeWriter_Format(writer, "%S.%S", module, qualname);
         goto done;
     }
 
+error:
+    rc = -1;
+    goto done;
+
 use_repr:
-    r = PyObject_Repr(p);
+    rc = PyUnicodeWriter_WriteRepr(writer, p);
+    goto done;
 
 done:
     Py_XDECREF(qualname);
     Py_XDECREF(module);
-    if (r == NULL) {
-        // error if any of the above PyObject_Repr/PyUnicode_From* fail
-        rc = -1;
-    }
-    else {
-        rc = _PyUnicodeWriter_WriteStr(writer, r);
-        Py_DECREF(r);
-    }
     return rc;
 }
 
 static int
-ga_repr_items_list(_PyUnicodeWriter *writer, PyObject *p)
+ga_repr_items_list(PyUnicodeWriter *writer, PyObject *p)
 {
     assert(PyList_CheckExact(p));
 
     Py_ssize_t len = PyList_GET_SIZE(p);
 
-    if (_PyUnicodeWriter_WriteASCIIString(writer, "[", 1) < 0) {
+    if (PyUnicodeWriter_WriteChar(writer, '[') < 0) {
         return -1;
     }
 
     for (Py_ssize_t i = 0; i < len; i++) {
         if (i > 0) {
-            if (_PyUnicodeWriter_WriteASCIIString(writer, ", ", 2) < 0) {
+            if (PyUnicodeWriter_WriteUTF8(writer, ", ", 2) < 0) {
                 return -1;
             }
         }
@@ -140,7 +136,7 @@ ga_repr_items_list(_PyUnicodeWriter *writer, PyObject *p)
         }
     }
 
-    if (_PyUnicodeWriter_WriteASCIIString(writer, "]", 1) < 0) {
+    if (PyUnicodeWriter_WriteChar(writer, ']') < 0) {
         return -1;
     }
 
@@ -153,49 +149,55 @@ ga_repr(PyObject *self)
     gaobject *alias = (gaobject *)self;
     Py_ssize_t len = PyTuple_GET_SIZE(alias->args);
 
-    _PyUnicodeWriter writer;
-    _PyUnicodeWriter_Init(&writer);
+    // Estimation based on the shortest format: "int[int, int, int]"
+    Py_ssize_t estimate = (len <= PY_SSIZE_T_MAX / 5) ? len * 5 : len;
+    estimate = 3 + 1 + estimate + 1;
+    PyUnicodeWriter *writer = PyUnicodeWriter_Create(estimate);
+    if (writer == NULL) {
+        return NULL;
+    }
 
     if (alias->starred) {
-        if (_PyUnicodeWriter_WriteASCIIString(&writer, "*", 1) < 0) {
+        if (PyUnicodeWriter_WriteChar(writer, '*') < 0) {
             goto error;
         }
     }
-    if (ga_repr_item(&writer, alias->origin) < 0) {
+    if (ga_repr_item(writer, alias->origin) < 0) {
         goto error;
     }
-    if (_PyUnicodeWriter_WriteASCIIString(&writer, "[", 1) < 0) {
+    if (PyUnicodeWriter_WriteChar(writer, '[') < 0) {
         goto error;
     }
     for (Py_ssize_t i = 0; i < len; i++) {
         if (i > 0) {
-            if (_PyUnicodeWriter_WriteASCIIString(&writer, ", ", 2) < 0) {
+            if (PyUnicodeWriter_WriteUTF8(writer, ", ", 2) < 0) {
                 goto error;
             }
         }
         PyObject *p = PyTuple_GET_ITEM(alias->args, i);
         if (PyList_CheckExact(p)) {
             // Looks like we are working with ParamSpec's list of type args:
-            if (ga_repr_items_list(&writer, p) < 0) {
+            if (ga_repr_items_list(writer, p) < 0) {
                 goto error;
             }
         }
-        else if (ga_repr_item(&writer, p) < 0) {
+        else if (ga_repr_item(writer, p) < 0) {
             goto error;
         }
     }
     if (len == 0) {
         // for something like tuple[()] we should print a "()"
-        if (_PyUnicodeWriter_WriteASCIIString(&writer, "()", 2) < 0) {
+        if (PyUnicodeWriter_WriteUTF8(writer, "()", 2) < 0) {
             goto error;
         }
     }
-    if (_PyUnicodeWriter_WriteASCIIString(&writer, "]", 1) < 0) {
+    if (PyUnicodeWriter_WriteChar(writer, ']') < 0) {
         goto error;
     }
-    return _PyUnicodeWriter_Finish(&writer);
+    return PyUnicodeWriter_Finish(writer);
+
 error:
-    _PyUnicodeWriter_Dealloc(&writer);
+    PyUnicodeWriter_Discard(writer);
     return NULL;
 }
 

Original file line number	Diff line number	Diff line change
`@@ -51,16 +51,15 @@ ga_traverse(PyObject self, visitproc visit, void arg)`
`51`	`51`	`}`
`52`	`52`
`53`	`53`	`static int`
`54`		`-ga_repr_item(_PyUnicodeWriter writer, PyObject p)`
	`54`	`+ga_repr_item(PyUnicodeWriter writer, PyObject p)`
`55`	`55`	`{`
`56`	`56`	`PyObject *qualname = NULL;`
`57`	`57`	`PyObject *module = NULL;`
`58`		`- PyObject *r = NULL;`
`59`	`58`	`int rc;`
`60`	`59`
`61`	`60`	`if (p == Py_Ellipsis) {`
`62`	`61`	`// The Ellipsis object`
`63`		`- r = PyUnicode_FromString("...");`
	`62`	`+ rc = PyUnicodeWriter_WriteUTF8(writer, "...", 3);`
`64`	`63`	`goto done;`
`65`	`64`	`}`
`66`	`65`
`@@ -71,17 +70,17 @@ ga_repr_item(_PyUnicodeWriter writer, PyObject p)`
`71`	`70`	`goto use_repr;`
`72`	`71`	`}`
`73`	`72`	`if (rc < 0) {`
`74`		`- goto done;`
	`73`	`+ goto error;`
`75`	`74`	`}`
`76`	`75`
`77`	`76`	`if (PyObject_GetOptionalAttr(p, &_Py_ID(__qualname__), &qualname) < 0) {`
`78`		`- goto done;`
	`77`	`+ goto error;`
`79`	`78`	`}`
`80`	`79`	`if (qualname == NULL) {`
`81`	`80`	`goto use_repr;`
`82`	`81`	`}`
`83`	`82`	`if (PyObject_GetOptionalAttr(p, &_Py_ID(__module__), &module) < 0) {`
`84`		`- goto done;`
	`83`	`+ goto error;`
`85`	`84`	`}`
`86`	`85`	`if (module == NULL \|\| module == Py_None) {`
`87`	`86`	`goto use_repr;`
`@@ -92,45 +91,42 @@ ga_repr_item(_PyUnicodeWriter writer, PyObject p)`
`92`	`91`	`_PyUnicode_EqualToASCIIString(module, "builtins"))`
`93`	`92`	`{`
`94`	`93`	`// builtins don't need a module name`
`95`		`- r = PyObject_Str(qualname);`
	`94`	`+ rc = PyUnicodeWriter_WriteStr(writer, qualname);`
`96`	`95`	`goto done;`
`97`	`96`	`}`
`98`	`97`	`else {`
`99`		`- r = PyUnicode_FromFormat("%S.%S", module, qualname);`
	`98`	`+ rc = PyUnicodeWriter_Format(writer, "%S.%S", module, qualname);`
`100`	`99`	`goto done;`
`101`	`100`	`}`
`102`	`101`
	`102`	`+error:`
	`103`	`+ rc = -1;`
	`104`	`+ goto done;`
	`105`	`+`
`103`	`106`	`use_repr:`
`104`		`- r = PyObject_Repr(p);`
	`107`	`+ rc = PyUnicodeWriter_WriteRepr(writer, p);`
	`108`	`+ goto done;`
`105`	`109`
`106`	`110`	`done:`
`107`	`111`	`Py_XDECREF(qualname);`
`108`	`112`	`Py_XDECREF(module);`
`109`		`- if (r == NULL) {`
`110`		`- // error if any of the above PyObject_Repr/PyUnicode_From* fail`
`111`		`- rc = -1;`
`112`		`- }`
`113`		`- else {`
`114`		`- rc = _PyUnicodeWriter_WriteStr(writer, r);`
`115`		`- Py_DECREF(r);`
`116`		`- }`
`117`	`113`	`return rc;`
`118`	`114`	`}`
`119`	`115`
`120`	`116`	`static int`
`121`		`-ga_repr_items_list(_PyUnicodeWriter writer, PyObject p)`
	`117`	`+ga_repr_items_list(PyUnicodeWriter writer, PyObject p)`
`122`	`118`	`{`
`123`	`119`	`assert(PyList_CheckExact(p));`
`124`	`120`
`125`	`121`	`Py_ssize_t len = PyList_GET_SIZE(p);`
`126`	`122`
`127`		`- if (_PyUnicodeWriter_WriteASCIIString(writer, "[", 1) < 0) {`
	`123`	`+ if (PyUnicodeWriter_WriteChar(writer, '[') < 0) {`
`128`	`124`	`return -1;`
`129`	`125`	`}`
`130`	`126`
`131`	`127`	`for (Py_ssize_t i = 0; i < len; i++) {`
`132`	`128`	`if (i > 0) {`
`133`		`- if (_PyUnicodeWriter_WriteASCIIString(writer, ", ", 2) < 0) {`
	`129`	`+ if (PyUnicodeWriter_WriteUTF8(writer, ", ", 2) < 0) {`
`134`	`130`	`return -1;`
`135`	`131`	`}`
`136`	`132`	`}`
`@@ -140,7 +136,7 @@ ga_repr_items_list(_PyUnicodeWriter writer, PyObject p)`
`140`	`136`	`}`
`141`	`137`	`}`
`142`	`138`
`143`		`- if (_PyUnicodeWriter_WriteASCIIString(writer, "]", 1) < 0) {`
	`139`	`+ if (PyUnicodeWriter_WriteChar(writer, ']') < 0) {`
`144`	`140`	`return -1;`
`145`	`141`	`}`
`146`	`142`
`@@ -153,49 +149,55 @@ ga_repr(PyObject *self)`
`153`	`149`	`gaobject alias = (gaobject )self;`
`154`	`150`	`Py_ssize_t len = PyTuple_GET_SIZE(alias->args);`
`155`	`151`
`156`		`- _PyUnicodeWriter writer;`
`157`		`- _PyUnicodeWriter_Init(&writer);`
	`152`	`+ // Estimation based on the shortest format: "int[int, int, int]"`
	`153`	`+ Py_ssize_t estimate = (len <= PY_SSIZE_T_MAX / 5) ? len * 5 : len;`
	`154`	`+ estimate = 3 + 1 + estimate + 1;`
	`155`	`+ PyUnicodeWriter *writer = PyUnicodeWriter_Create(estimate);`
	`156`	`+ if (writer == NULL) {`
	`157`	`+ return NULL;`
	`158`	`+ }`
`158`	`159`
`159`	`160`	`if (alias->starred) {`
`160`		`- if (_PyUnicodeWriter_WriteASCIIString(&writer, "*", 1) < 0) {`
	`161`	`+ if (PyUnicodeWriter_WriteChar(writer, '*') < 0) {`
`161`	`162`	`goto error;`
`162`	`163`	`}`
`163`	`164`	`}`
`164`		`- if (ga_repr_item(&writer, alias->origin) < 0) {`
	`165`	`+ if (ga_repr_item(writer, alias->origin) < 0) {`
`165`	`166`	`goto error;`
`166`	`167`	`}`
`167`		`- if (_PyUnicodeWriter_WriteASCIIString(&writer, "[", 1) < 0) {`
	`168`	`+ if (PyUnicodeWriter_WriteChar(writer, '[') < 0) {`
`168`	`169`	`goto error;`
`169`	`170`	`}`
`170`	`171`	`for (Py_ssize_t i = 0; i < len; i++) {`
`171`	`172`	`if (i > 0) {`
`172`		`- if (_PyUnicodeWriter_WriteASCIIString(&writer, ", ", 2) < 0) {`
	`173`	`+ if (PyUnicodeWriter_WriteUTF8(writer, ", ", 2) < 0) {`
`173`	`174`	`goto error;`
`174`	`175`	`}`
`175`	`176`	`}`
`176`	`177`	`PyObject *p = PyTuple_GET_ITEM(alias->args, i);`
`177`	`178`	`if (PyList_CheckExact(p)) {`
`178`	`179`	`// Looks like we are working with ParamSpec's list of type args:`
`179`		`- if (ga_repr_items_list(&writer, p) < 0) {`
	`180`	`+ if (ga_repr_items_list(writer, p) < 0) {`
`180`	`181`	`goto error;`
`181`	`182`	`}`
`182`	`183`	`}`
`183`		`- else if (ga_repr_item(&writer, p) < 0) {`
	`184`	`+ else if (ga_repr_item(writer, p) < 0) {`
`184`	`185`	`goto error;`
`185`	`186`	`}`
`186`	`187`	`}`
`187`	`188`	`if (len == 0) {`
`188`	`189`	`// for something like tuple[()] we should print a "()"`
`189`		`- if (_PyUnicodeWriter_WriteASCIIString(&writer, "()", 2) < 0) {`
	`190`	`+ if (PyUnicodeWriter_WriteUTF8(writer, "()", 2) < 0) {`
`190`	`191`	`goto error;`
`191`	`192`	`}`
`192`	`193`	`}`
`193`		`- if (_PyUnicodeWriter_WriteASCIIString(&writer, "]", 1) < 0) {`
	`194`	`+ if (PyUnicodeWriter_WriteChar(writer, ']') < 0) {`
`194`	`195`	`goto error;`
`195`	`196`	`}`
`196`		`- return _PyUnicodeWriter_Finish(&writer);`
	`197`	`+ return PyUnicodeWriter_Finish(writer);`
	`198`	`+`
`197`	`199`	`error:`
`198`		`- _PyUnicodeWriter_Dealloc(&writer);`
	`200`	`+ PyUnicodeWriter_Discard(writer);`
`199`	`201`	`return NULL;`
`200`	`202`	`}`
`201`	`203`