Use PyTupleWriter in multiple functions

Author: vstinner

Modules/_tkinter.c

@@ -3080,60 +3080,33 @@ Tkapp_Dealloc(PyObject *op)
 
 /**** Tkinter Module ****/
 
-typedef struct {
-    PyObject* tuple;
-    Py_ssize_t size; /* current size */
-    Py_ssize_t maxsize; /* allocated size */
-} FlattenContext;
-
 static int
-_bump(FlattenContext* context, Py_ssize_t size)
-{
-    /* expand tuple to hold (at least) size new items.
-       return true if successful, false if an exception was raised */
-
-    Py_ssize_t maxsize = context->maxsize * 2;  /* never overflows */
-
-    if (maxsize < context->size + size)
-        maxsize = context->size + size;  /* never overflows */
-
-    context->maxsize = maxsize;
-
-    return _PyTuple_Resize(&context->tuple, maxsize) >= 0;
-}
-
-static int
-_flatten1(FlattenContext* context, PyObject* item, int depth)
+_flatten1(PyTupleWriter *writer, PyObject *item, int depth)
 {
     /* add tuple or list to argument tuple (recursively) */
 
-    Py_ssize_t i, size;
-
     if (depth > 1000) {
         PyErr_SetString(PyExc_ValueError,
                         "nesting too deep in _flatten");
         return 0;
-    } else if (PyTuple_Check(item) || PyList_Check(item)) {
-        size = PySequence_Fast_GET_SIZE(item);
-        /* preallocate (assume no nesting) */
-        if (context->size + size > context->maxsize &&
-            !_bump(context, size))
-            return 0;
+    }
+
+    if (PyTuple_Check(item) || PyList_Check(item)) {
+        Py_ssize_t size = PySequence_Fast_GET_SIZE(item);
         /* copy items to output tuple */
-        for (i = 0; i < size; i++) {
+        for (Py_ssize_t i = 0; i < size; i++) {
             PyObject *o = PySequence_Fast_GET_ITEM(item, i);
             if (PyList_Check(o) || PyTuple_Check(o)) {
-                if (!_flatten1(context, o, depth + 1))
+                if (!_flatten1(writer, o, depth + 1))
                     return 0;
             } else if (o != Py_None) {
-                if (context->size + 1 > context->maxsize &&
-                    !_bump(context, 1))
+                if (PyTupleWriter_Add(writer, o) < 0) {
                     return 0;
-                PyTuple_SET_ITEM(context->tuple,
-                                 context->size++, Py_NewRef(o));
+                }
             }
         }
-    } else {
+    }
+    else {
         PyErr_SetString(PyExc_TypeError, "argument must be sequence");
         return 0;
     }
@@ -3152,29 +3125,25 @@ static PyObject *
 _tkinter__flatten(PyObject *module, PyObject *item)
 /*[clinic end generated code: output=cad02a3f97f29862 input=6b9c12260aa1157f]*/
 {
-    FlattenContext context;
-
-    context.maxsize = PySequence_Size(item);
-    if (context.maxsize < 0)
+    Py_ssize_t maxsize = PySequence_Size(item);
+    if (maxsize < 0) {
         return NULL;
-    if (context.maxsize == 0)
+    }
+    if (maxsize == 0) {
         return PyTuple_New(0);
+    }
 
-    context.tuple = PyTuple_New(context.maxsize);
-    if (!context.tuple)
-        return NULL;
-
-    context.size = 0;
-
-    if (!_flatten1(&context, item, 0)) {
-        Py_XDECREF(context.tuple);
+    PyTupleWriter *writer = PyTupleWriter_Create(maxsize);
+    if (writer == NULL) {
         return NULL;
     }
 
-    if (_PyTuple_Resize(&context.tuple, context.size))
+    if (!_flatten1(writer, item, 0)) {
+        PyTupleWriter_Discard(writer);
         return NULL;
+    }
 
-    return context.tuple;
+    return PyTupleWriter_Finish(writer);
 }
 
 /*[clinic input]

Objects/genericaliasobject.c

@@ -4,6 +4,7 @@
 #include "pycore_ceval.h"         // _PyEval_GetBuiltin()
 #include "pycore_modsupport.h"    // _PyArg_NoKeywords()
 #include "pycore_object.h"
+#include "pycore_tuple.h"         // _PyTuple_ITEMS()
 #include "pycore_typevarobject.h" // _Py_typing_type_repr
 #include "pycore_unicodeobject.h" // _PyUnicode_EqualToASCIIString()
 #include "pycore_unionobject.h"   // _Py_union_type_or, _PyGenericAlias_Check
@@ -161,22 +162,6 @@ tuple_add(PyObject *self, Py_ssize_t len, PyObject *item)
     return 0;
 }
 
-static Py_ssize_t
-tuple_extend(PyObject **dst, Py_ssize_t dstindex,
-             PyObject **src, Py_ssize_t count)
-{
-    assert(count >= 0);
-    if (_PyTuple_Resize(dst, PyTuple_GET_SIZE(*dst) + count - 1) != 0) {
-        return -1;
-    }
-    assert(dstindex + count <= PyTuple_GET_SIZE(*dst));
-    for (Py_ssize_t i = 0; i < count; ++i) {
-        PyObject *item = src[i];
-        PyTuple_SET_ITEM(*dst, dstindex + i, Py_NewRef(item));
-    }
-    return dstindex + count;
-}
-
 PyObject *
 _Py_make_parameters(PyObject *args)
 {
@@ -274,43 +259,54 @@ subs_tvars(PyObject *obj, PyObject *params,
     if (PyObject_GetOptionalAttr(obj, &_Py_ID(__parameters__), &subparams) < 0) {
         return NULL;
     }
+
     if (subparams && PyTuple_Check(subparams) && PyTuple_GET_SIZE(subparams)) {
         Py_ssize_t nparams = PyTuple_GET_SIZE(params);
         Py_ssize_t nsubargs = PyTuple_GET_SIZE(subparams);
-        PyObject *subargs = PyTuple_New(nsubargs);
-        if (subargs == NULL) {
-            Py_DECREF(subparams);
-            return NULL;
+        PyTupleWriter *writer = PyTupleWriter_Create(nsubargs);
+        if (writer == NULL) {
+            goto error;
         }
-        Py_ssize_t j = 0;
         for (Py_ssize_t i = 0; i < nsubargs; ++i) {
             PyObject *arg = PyTuple_GET_ITEM(subparams, i);
             Py_ssize_t iparam = tuple_index(params, nparams, arg);
             if (iparam >= 0) {
                 PyObject *param = PyTuple_GET_ITEM(params, iparam);
                 arg = argitems[iparam];
                 if (Py_TYPE(param)->tp_iter && PyTuple_Check(arg)) {  // TypeVarTuple
-                    j = tuple_extend(&subargs, j,
-                                    &PyTuple_GET_ITEM(arg, 0),
-                                    PyTuple_GET_SIZE(arg));
-                    if (j < 0) {
-                        return NULL;
+                    if (PyTupleWriter_AddArray(writer,
+                                               _PyTuple_ITEMS(arg),
+                                               PyTuple_GET_SIZE(arg)) < 0) {
+                        goto error;
                     }
                     continue;
                 }
             }
-            PyTuple_SET_ITEM(subargs, j, Py_NewRef(arg));
-            j++;
+            if (PyTupleWriter_Add(writer, arg) < 0) {
+                goto error;
+            }
         }
-        assert(j == PyTuple_GET_SIZE(subargs));
 
-        obj = PyObject_GetItem(obj, subargs);
+        PyObject *subargs = PyTupleWriter_Finish(writer);
+        if (subargs != NULL) {
+            obj = PyObject_GetItem(obj, subargs);
+            Py_DECREF(subargs);
+        }
+        else {
+            obj = NULL;
+        }
+        goto done;
 
-        Py_DECREF(subargs);
+error:
+        PyTupleWriter_Discard(writer);
+        Py_DECREF(subparams);
+        return NULL;
     }
     else {
         Py_INCREF(obj);
     }
+
+done:
     Py_XDECREF(subparams);
     return obj;
 }
@@ -454,60 +450,50 @@ _Py_subs_parameters(PyObject *self, PyObject *args, PyObject *parameters, PyObje
         }
     }
     Py_ssize_t nargs = PyTuple_GET_SIZE(args);
-    PyObject *newargs = PyTuple_New(nargs);
+    PyTupleWriter *newargs = PyTupleWriter_Create(nargs);
     if (newargs == NULL) {
-        Py_DECREF(item);
-        Py_XDECREF(tuple_args);
-        return NULL;
+        goto error;
     }
-    for (Py_ssize_t iarg = 0, jarg = 0; iarg < nargs; iarg++) {
+
+    for (Py_ssize_t iarg = 0; iarg < nargs; iarg++) {
         PyObject *arg = PyTuple_GET_ITEM(args, iarg);
         if (PyType_Check(arg)) {
-            PyTuple_SET_ITEM(newargs, jarg, Py_NewRef(arg));
-            jarg++;
+            if (PyTupleWriter_Add(newargs, arg) < 0) {
+                goto error;
+            }
             continue;
         }
+
         // Recursively substitute params in lists/tuples.
         if (PyTuple_Check(arg) || PyList_Check(arg)) {
             PyObject *subargs = _Py_subs_parameters(self, arg, parameters, item);
             if (subargs == NULL) {
-                Py_DECREF(newargs);
-                Py_DECREF(item);
-                Py_XDECREF(tuple_args);
-                return NULL;
+                goto error;
             }
             if (PyTuple_Check(arg)) {
-                PyTuple_SET_ITEM(newargs, jarg, subargs);
+                if (PyTupleWriter_AddSteal(newargs, subargs) < 0) {
+                    goto error;
+                }
             }
             else {
                 // _Py_subs_parameters returns a tuple. If the original arg was a list,
                 // convert subargs to a list as well.
                 PyObject *subargs_list = PySequence_List(subargs);
                 Py_DECREF(subargs);
-                if (subargs_list == NULL) {
-                    Py_DECREF(newargs);
-                    Py_DECREF(item);
-                    Py_XDECREF(tuple_args);
-                    return NULL;
+                if (PyTupleWriter_AddSteal(newargs, subargs_list) < 0) {
+                    goto error;
                 }
-                PyTuple_SET_ITEM(newargs, jarg, subargs_list);
             }
-            jarg++;
             continue;
         }
+
         int unpack = _is_unpacked_typevartuple(arg);
         if (unpack < 0) {
-            Py_DECREF(newargs);
-            Py_DECREF(item);
-            Py_XDECREF(tuple_args);
-            return NULL;
+            goto error;
         }
         PyObject *subst;
         if (PyObject_GetOptionalAttr(arg, &_Py_ID(__typing_subst__), &subst) < 0) {
-            Py_DECREF(newargs);
-            Py_DECREF(item);
-            Py_XDECREF(tuple_args);
-            return NULL;
+            goto error;
         }
         if (subst) {
             Py_ssize_t iparam = tuple_index(parameters, nparams, arg);
@@ -518,43 +504,42 @@ _Py_subs_parameters(PyObject *self, PyObject *args, PyObject *parameters, PyObje
         else {
             arg = subs_tvars(arg, parameters, argitems, nitems);
         }
-        if (arg == NULL) {
-            Py_DECREF(newargs);
-            Py_DECREF(item);
-            Py_XDECREF(tuple_args);
-            return NULL;
-        }
         if (unpack) {
+            if (arg == NULL) {
+                goto error;
+            }
             if (!PyTuple_Check(arg)) {
-                Py_DECREF(newargs);
-                Py_DECREF(item);
-                Py_XDECREF(tuple_args);
                 PyObject *original = PyTuple_GET_ITEM(args, iarg);
                 PyErr_Format(PyExc_TypeError,
                              "expected __typing_subst__ of %T objects to return a tuple, not %T",
                              original, arg);
                 Py_DECREF(arg);
-                return NULL;
+                goto error;
             }
-            jarg = tuple_extend(&newargs, jarg,
-                    &PyTuple_GET_ITEM(arg, 0), PyTuple_GET_SIZE(arg));
-            Py_DECREF(arg);
-            if (jarg < 0) {
-                Py_DECREF(item);
-                Py_XDECREF(tuple_args);
-                assert(newargs == NULL);
-                return NULL;
+            if (PyTupleWriter_AddArray(newargs,
+                                       _PyTuple_ITEMS(arg),
+                                       PyTuple_GET_SIZE(arg)) < 0) {
+                Py_DECREF(arg);
+                goto error;
             }
+            Py_DECREF(arg);
         }
         else {
-            PyTuple_SET_ITEM(newargs, jarg, arg);
-            jarg++;
+            if (PyTupleWriter_AddSteal(newargs, arg) < 0) {
+                goto error;
+            }
         }
     }
 
     Py_DECREF(item);
     Py_XDECREF(tuple_args);
-    return newargs;
+    return PyTupleWriter_Finish(newargs);
+
+error:
+    PyTupleWriter_Discard(newargs);
+    Py_DECREF(item);
+    Py_XDECREF(tuple_args);
+    return NULL;
 }
 
 PyDoc_STRVAR(genericalias__doc__,