[C API] Add PyTupleWriter API

Include/cpython/tupleobject.h

@@ -30,7 +30,12 @@ static inline Py_ssize_t PyTuple_GET_SIZE(PyObject *op) {
 
 /* Function *only* to be used to fill in brand new tuples */
 static inline void
-PyTuple_SET_ITEM(PyObject *op, Py_ssize_t index, PyObject *value) {
+PyTuple_SET_ITEM(PyObject *op, Py_ssize_t index, PyObject *value)
+{
+    // For backward compatibility, accept mutating a tuple even if
+    // _PyObject_IsUniquelyReferenced(op) is false and/or tuple->ob_hash != -1.
+    // See for example _PyTuple_Recycle(). Accept also setting an item which
+    // was already set: the caller must Py_DECREF() the old value in this case.
     PyTupleObject *tuple = _PyTuple_CAST(op);
     assert(0 <= index);
     assert(index < Py_SIZE(tuple));
@@ -42,3 +47,29 @@ PyTuple_SET_ITEM(PyObject *op, Py_ssize_t index, PyObject *value) {
 PyAPI_FUNC(PyObject*) PyTuple_FromArray(
     PyObject *const *array,
     Py_ssize_t size);
+
+
+/* --- PyTupleWriter API -------------------------------------------------- */
+
+struct PyTupleWriter;
+typedef struct PyTupleWriter PyTupleWriter;
+
+PyAPI_FUNC(PyTupleWriter*) PyTupleWriter_Create(Py_ssize_t size);
+PyAPI_FUNC(PyObject*) PyTupleWriter_Finish(PyTupleWriter *writer);
+PyAPI_FUNC(PyObject*) PyTupleWriter_FinishWithSize(
+    PyTupleWriter *writer,
+    Py_ssize_t final_size);
+PyAPI_FUNC(void) PyTupleWriter_Discard(PyTupleWriter *writer);
+PyAPI_FUNC(Py_ssize_t) PyTupleWriter_GetSize(PyTupleWriter *writer);
+PyAPI_FUNC(int) PyTupleWriter_SetItem(
+    PyTupleWriter *writer,
+    Py_ssize_t index,
+    PyObject *item);
+PyAPI_FUNC(int) PyTupleWriter_SetItems(
+    PyTupleWriter *writer,
+    Py_ssize_t index,
+    PyObject *const *items,
+    Py_ssize_t nitem);
+PyAPI_FUNC(int) PyTupleWriter_Resize(
+    PyTupleWriter **writer,
+    Py_ssize_t new_size);

Include/internal/pycore_tuple.h

@@ -76,6 +76,77 @@
 #endif
 #define _PyTuple_HASH_EMPTY (_PyTuple_HASH_XXPRIME_5 + (_PyTuple_HASH_XXPRIME_5 ^ 3527539UL))
 
+
+#ifndef NDEBUG
+// Check that a tuple can be mutated in-place.
+// Usage: assert(_PyTuple_CanMutate(v));
+static int
+_PyTuple_CanMutate(PyTupleObject *tuple)
+{
+    assert(PyTuple_Check(tuple));
+
+    // Do not check _PyObject_IsUniquelyReferenced(): the function is called
+    // by _PyTupleWriter_CheckConsistency() and a tuple writer can be used by
+    // multiple threads using external locking.
+    assert(Py_REFCNT(tuple) == 1);
+
+    // The empty tuple must not be modified
+    assert(1 <= Py_SIZE(tuple));
+    assert(tuple != &_Py_SINGLETON(tuple_empty));
+
+    // The hash value must not be cached yet. Otherwise, it means that the
+    // tuple object was already used in Python somehow (ex: as a dictionary
+    // key).
+    assert(tuple->ob_hash == -1);
+
+    return 1;
+}
+#endif
+
+
+/* --- PyTupleWriter API -------------------------------------------------- */
+
+#ifndef NDEBUG
+static inline int
+_PyTupleWriter_CheckConsistency(PyTupleWriter *writer)
+{
+    assert(writer != NULL);
+    assert(PyTuple_CheckExact(_PyObject_CAST(writer)));
+    PyTupleObject *tuple = (PyTupleObject*)writer;
+    // Do not use _PyObject_IsUniquelyReferenced(): the caller can have
+    // its own lock to prevent a writer from being used by two threads
+    // at the same time.
+    assert(Py_REFCNT(tuple) == 1);
+    assert(0 <= Py_SIZE(tuple));
+    return 1;
+}
+#endif
+
+
+static inline void
+_PyTupleWriter_SET_ITEM_STEAL(PyTupleWriter *writer,
+                              Py_ssize_t index, PyObject *item)
+{
+    assert(_PyTupleWriter_CheckConsistency(writer));
+    PyTupleObject *tuple = _PyTuple_CAST(writer);
+
+    assert(_PyTuple_CanMutate(tuple));
+    assert(0 <= index && index < Py_SIZE(tuple));
+    assert(tuple->ob_item[index] == NULL);
+    assert(item != NULL);
+
+    tuple->ob_item[index] = item;
+}
+
+
+static inline void
+_PyTupleWriter_SET_ITEM(PyTupleWriter *writer,
+                        Py_ssize_t index, PyObject *item)
+{
+    assert(item != NULL);
+    return _PyTupleWriter_SET_ITEM_STEAL(writer, index, Py_NewRef(item));
+}
+
 #ifdef __cplusplus
 }
 #endif

Modules/_tkinter.c

@@ -31,6 +31,7 @@ Copyright (C) 1994 Steen Lumholt.
 #endif
 
 #include "pycore_long.h"          // _PyLong_IsNegative()
+#include "pycore_tuple.h"         // _PyTupleWriter_SET_ITEM()
 #include "pycore_unicodeobject.h" // _PyUnicode_AsUTF8String
 
 #include <signal.h>               // SIGINT
@@ -3401,7 +3402,7 @@ Tkapp_Dealloc(PyObject *op)
 /**** Tkinter Module ****/
 
 typedef struct {
-    PyObject* tuple;
+    PyTupleWriter *writer;
     Py_ssize_t size; /* current size */
     Py_ssize_t maxsize; /* allocated size */
 } FlattenContext;
@@ -3419,7 +3420,7 @@ _bump(FlattenContext* context, Py_ssize_t size)
 
     context->maxsize = maxsize;
 
-    return _PyTuple_Resize(&context->tuple, maxsize) >= 0;
+    return PyTupleWriter_Resize(&context->writer, maxsize) >= 0;
 }
 
 static int
@@ -3449,8 +3450,7 @@ _flatten1(FlattenContext* context, PyObject* item, int depth)
                 if (context->size + 1 > context->maxsize &&
                     !_bump(context, 1))
                     return 0;
-                PyTuple_SET_ITEM(context->tuple,
-                                 context->size++, Py_NewRef(o));
+                _PyTupleWriter_SET_ITEM(context->writer, context->size++, o);
             }
         }
     } else {
@@ -3480,21 +3480,18 @@ _tkinter__flatten(PyObject *module, PyObject *item)
     if (context.maxsize == 0)
         return PyTuple_New(0);
 
-    context.tuple = PyTuple_New(context.maxsize);
-    if (!context.tuple)
+    context.writer = PyTupleWriter_Create(context.maxsize);
+    if (!context.writer) {
         return NULL;
-
+    }
     context.size = 0;
 
     if (!_flatten1(&context, item, 0)) {
-        Py_XDECREF(context.tuple);
+        PyTupleWriter_Discard(context.writer);
         return NULL;
     }
 
-    if (_PyTuple_Resize(&context.tuple, context.size))
-        return NULL;
-
-    return context.tuple;
+    return PyTupleWriter_FinishWithSize(context.writer, context.size);
 }
 
 /*[clinic input]

Modules/itertoolsmodule.c

@@ -207,25 +207,24 @@ batched_next(PyObject *op)
     Py_ssize_t n = FT_ATOMIC_LOAD_SSIZE_RELAXED(bo->batch_size);
     PyObject *it = bo->it;
     PyObject *item;
-    PyObject *result;
 
     if (n < 0) {
         return NULL;
     }
-    result = PyTuple_New(n);
-    if (result == NULL) {
+
+    PyTupleWriter *writer = PyTupleWriter_Create(n);
+    if (writer == NULL) {
         return NULL;
     }
     iternextfunc iternext = *Py_TYPE(it)->tp_iternext;
-    PyObject **items = _PyTuple_ITEMS(result);
     for (i=0 ; i < n ; i++) {
         item = iternext(it);
         if (item == NULL) {
             goto null_item;
         }
-        items[i] = item;
+        _PyTupleWriter_SET_ITEM_STEAL(writer, i, item);
     }
-    return result;
+    return PyTupleWriter_Finish(writer);
 
  null_item:
     if (PyErr_Occurred()) {
@@ -235,7 +234,7 @@ batched_next(PyObject *op)
 #ifndef Py_GIL_DISABLED
             Py_CLEAR(bo->it);
 #endif
-            Py_DECREF(result);
+            PyTupleWriter_Discard(writer);
             return NULL;
         }
         PyErr_Clear();
@@ -245,20 +244,19 @@ batched_next(PyObject *op)
 #ifndef Py_GIL_DISABLED
         Py_CLEAR(bo->it);
 #endif
-        Py_DECREF(result);
+        PyTupleWriter_Discard(writer);
         return NULL;
     }
     if (bo->strict) {
         FT_ATOMIC_STORE_SSIZE_RELAXED(bo->batch_size, -1);
 #ifndef Py_GIL_DISABLED
         Py_CLEAR(bo->it);
 #endif
-        Py_DECREF(result);
+        PyTupleWriter_Discard(writer);
         PyErr_SetString(PyExc_ValueError, "batched(): incomplete batch");
         return NULL;
     }
-    _PyTuple_Resize(&result, i);
-    return result;
+    return PyTupleWriter_FinishWithSize(writer, i);
 }
 
 static PyType_Slot batched_slots[] = {