Add PyFrozenSetWriter_UpdateArray() function

Add PyFrozenSetWriter_UpdateArray() and
PyFrozenSetWriter_UpdateArraySteal() functions.

Add tests on PyFrozenSetWriter_AddSteal().

Typoe: rename frozsetwriter_check_consistency() to
frozensetwriter_check_consistency().

Author: vstinner

Doc/c-api/set.rst

@@ -289,3 +289,22 @@ See also the :c:func:`PyFrozenSet_New` function.
 
    On success, return ``0``.
    On error, set an exception and return ``-1``.
+
+   On error, elements already added are kept.
+
+.. c:function:: int PyFrozenSetWriter_UpdateArray(PyFrozenSetWriter *writer, PyObject *const *items, Py_ssize_t nitem)
+
+   Update the frozenset writer, adding *nitem* elements from *items* array:
+   copy references from *items* to the writer.
+
+   On success, return ``0``.
+   On error, set an exception and return ``-1``.
+
+   On error, elements already added are kept.
+
+   *nitem* must be positive or zero.
+
+.. c:function:: int PyFrozenSetWriter_UpdateArraySteal(PyFrozenSetWriter *writer, PyObject *const *items, Py_ssize_t nitem)
+
+   Similar to :c:func:`PyFrozenSetWriter_UpdateArray`, but :term:`steal`
+   references instead of copying references.

Doc/whatsnew/3.16.rst

@@ -1304,6 +1304,8 @@ New features
   * :c:func:`PyFrozenSetWriter_Add`
   * :c:func:`PyFrozenSetWriter_AddSteal`
   * :c:func:`PyFrozenSetWriter_Update`
+  * :c:func:`PyFrozenSetWriter_UpdateArray`
+  * :c:func:`PyFrozenSetWriter_UpdateArraySteal`
 
   (Contributed by Victor Stinner in :gh:`160147`.)
 

Include/cpython/setobject.h

@@ -87,3 +87,11 @@ PyAPI_FUNC(int) PyFrozenSetWriter_AddSteal(
 PyAPI_FUNC(int) PyFrozenSetWriter_Update(
     PyFrozenSetWriter *writer,
     PyObject *iterable);
+PyAPI_FUNC(int) PyFrozenSetWriter_UpdateArray(
+    PyFrozenSetWriter *writer,
+    PyObject *const *items,
+    Py_ssize_t nitem);
+PyAPI_FUNC(int) PyFrozenSetWriter_UpdateArraySteal(
+    PyFrozenSetWriter *writer,
+    PyObject *const *items,
+    Py_ssize_t nitem);

Lib/test/test_capi/test_set.py

@@ -279,6 +279,7 @@ def test_set_next_entry(self):
 
 class FrozenSetWriterTests(unittest.TestCase):
     PySet_MINSIZE = _testcapi.PySet_MINSIZE
+    UNHASHABLE = [1, 2, 3]
 
     def create_writer(self, size=0):
         # Test PyFrozenSetWriter_Create()
@@ -296,6 +297,10 @@ def frozenset_equal(self, result, expected):
         self.assertEqual(result, expected)
         self.assertEqual(type(result), frozenset)
 
+    def assert_unhashable(self):
+        err_msg = f"unhashable type: '{type(self.UNHASHABLE).__name__}'"
+        return self.assertRaisesRegex(TypeError, err_msg)
+
     def test_get_table_set(self):
         # Test size parameter of PyFrozenSetWriter_Create()
         PySet_MINSIZE = self.PySet_MINSIZE
@@ -313,40 +318,78 @@ def create_items(self, nitem):
         # Use mortal (and unique) objects to be able to detect reference leaks
         return tuple(object() for _ in range(nitem))
 
-    def test_add(self):
-        # Test PyFrozenSetWriter_Append()
+    def check_add_method(self, add_method):
         items = self.create_items(self.PySet_MINSIZE * 2 - 1)
         table_size = self.compute_table_size(len(items))
         for prealloc in (0, len(items) // 2, len(items), len(items) * 10):
             with self.subTest(size=prealloc):
                 writer = self.create_writer()
+                add = getattr(writer, add_method)
                 for item in items:
-                    writer.add(item)
+                    add(item)
                 self.frozenset_equal(writer.finish(), items)
 
         # Test error (unhashable type)
         writer = self.create_writer(0)
         writer.update('abc')
-        with self.assertRaisesRegex(TypeError, "unhashable type: 'list'"):
-            writer.add([])
+        add = getattr(writer, add_method)
+        with self.assert_unhashable():
+            add(self.UNHASHABLE)
         self.frozenset_equal(writer.finish(), frozenset('abc'))
 
-    def test_update(self):
-        # Test PyFrozenSetWriter_Update()
-        items1 = self.create_items(self.PySet_MINSIZE - 1)
+    def test_add(self):
+        # Test PyFrozenSetWriter_Add()
+        self.check_add_method('add')
+
+    def test_add_steal(self):
+        # Test PyFrozenSetWriter_AddSteal()
+        self.check_add_method('add_steal')
+
+    def check_update_method(self, update_method):
+        items = self.create_items(self.PySet_MINSIZE - 1)
         items2 = self.create_items(self.PySet_MINSIZE - 1)
         common = self.create_items(self.PySet_MINSIZE - 1)
-        expected = set(items1) | set(items2) | set(common)
+        expected = set(items) | set(items2) | set(common)
         total = len(expected)
-        self.assertEqual(total, len(items1) + len(items2) + len(common))
+        self.assertEqual(total, len(items) + len(items2) + len(common))
+        empty = ()
 
         for prealloc in (0, total // 2, total, total * 10):
             with self.subTest(size=prealloc):
                 writer = self.create_writer(prealloc)
-                writer.update(items1 + common)
-                writer.update(common + items2)
+                update = getattr(writer, update_method)
+                update(items + common)
+                update(common + items2)
+                update(empty)  # test nitem=0
                 self.frozenset_equal(writer.finish(), expected)
 
+        # Test error path: first or last item is not hashable
+        for first in (True, False):
+            with self.subTest(first=first):
+                if first:
+                    items_error = (self.UNHASHABLE, *items)
+                    expected = frozenset()
+                else:
+                    items_error = (*items, self.UNHASHABLE)
+                    expected = frozenset(items)
+                writer = self.create_writer(0)
+                update = getattr(writer, update_method)
+                with self.assert_unhashable():
+                    update(items_error)
+                self.frozenset_equal(writer.finish(), expected)
+
+    def test_update(self):
+        # Test PyFrozenSetWriter_Update()
+        self.check_update_method('update')
+
+    def test_update_array(self):
+        # Test PyFrozenSetWriter_UpdateArray()
+        self.check_update_method('update_array')
+
+    def test_update_array_steal(self):
+        # Test PyFrozenSetWriter_UpdateArraySteal()
+        self.check_update_method('update_array_steal')
+
     def test_finish(self):
         # Test PyFrozenSetWriter_Finish(): track the frozenset in the GC if at
         # least one item is tracked by the GC

Misc/NEWS.d/next/C_API/2026-10-21-18-20-57.gh-issue-160147.7uo4pF.rst

@@ -7,5 +7,7 @@ Add a new :ref:`PyFrozenSetWriter <pyfrozensetwriter>` C API to create a
 * :c:func:`PyFrozenSetWriter_Add`
 * :c:func:`PyFrozenSetWriter_AddSteal`
 * :c:func:`PyFrozenSetWriter_Update`
+* :c:func:`PyFrozenSetWriter_UpdateArray`
+* :c:func:`PyFrozenSetWriter_UpdateArraySteal`
 
 Patch by Victor Stinner.

Modules/_testcapi/set.c

@@ -150,6 +150,58 @@ writer_update(PyObject *self_raw, PyObject *iterable)
 }
 
 
+static PyObject*
+_writer_update_array(PyObject *self_raw, PyObject *tuple, int steal)
+{
+    WriterObject *self = (WriterObject *)self_raw;
+    if (writer_check(self) < 0) {
+        return NULL;
+    }
+
+    if (!PyTuple_Check(tuple)) {
+        PyErr_Format(PyExc_TypeError,
+                     "expect tuple, got %T", tuple);
+        return NULL;
+    }
+    PyObject **items = &PyTuple_GET_ITEM(tuple, 0);
+    Py_ssize_t len = PyTuple_GET_SIZE(tuple);
+
+    if (steal) {
+        for (Py_ssize_t i = 0; i < len; i++) {
+            Py_INCREF(items[i]);
+        }
+    }
+
+    int res;
+    if (steal) {
+        res = PyFrozenSetWriter_UpdateArraySteal(self->writer, items, len);
+    }
+    else {
+        res = PyFrozenSetWriter_UpdateArray(self->writer, items, len);
+    }
+    if (res < 0) {
+        return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
+
+// Test PyFrozenSetWriter_UpdateArray()
+static PyObject*
+writer_update_array(PyObject *self, PyObject *tuple)
+{
+    return _writer_update_array(self, tuple, 0);
+}
+
+
+// Test PyFrozenSetWriter_UpdateArraySteal()
+static PyObject*
+writer_update_array_steal(PyObject *self, PyObject *tuple)
+{
+    return _writer_update_array(self, tuple, 1);
+}
+
+
 static PyObject*
 writer_get_table_size(PyObject *self_raw, PyObject *Py_UNUSED(args))
 {
@@ -198,6 +250,8 @@ static PyMethodDef writer_methods[] = {
     {"add", _PyCFunction_CAST(writer_add), METH_O},
     {"add_steal", _PyCFunction_CAST(writer_add_steal), METH_O},
     {"update", _PyCFunction_CAST(writer_update), METH_O},
+    {"update_array", _PyCFunction_CAST(writer_update_array), METH_O},
+    {"update_array_steal", _PyCFunction_CAST(writer_update_array_steal), METH_O},
     {"get_table_size", _PyCFunction_CAST(writer_get_table_size), METH_NOARGS},
     {"finish", _PyCFunction_CAST(writer_finish), METH_NOARGS},
     {"discard", _PyCFunction_CAST(writer_discard), METH_NOARGS},

Objects/setobject.c

@@ -3229,10 +3229,11 @@ static PyObject _dummy_struct = _PyObject_HEAD_INIT(&_PySetDummy_Type);
 
 #ifndef NDEBUG
 static int
-frozsetwriter_check_consistency(PyFrozenSetWriter *writer)
+frozensetwriter_check_consistency(PyFrozenSetWriter *writer)
 {
     assert(writer != NULL);
     PyObject *frozenset = _PyObject_CAST(writer);
+    assert(PyFrozenSet_CheckExact(frozenset));
     // Don't test _PyObject_IsUniquelyReferenced() since a frozenset writer
     // can be shared between multiple threads with external locking.
     assert(Py_REFCNT(frozenset) == 1);
@@ -3263,7 +3264,7 @@ PyFrozenSetWriter_Create(Py_ssize_t size)
     }
 
     PyFrozenSetWriter *writer = (PyFrozenSetWriter*)frozenset;
-    assert(frozsetwriter_check_consistency(writer));
+    assert(frozensetwriter_check_consistency(writer));
     return writer;
 }
 
@@ -3275,7 +3276,7 @@ PyFrozenSetWriter_Discard(PyFrozenSetWriter *writer)
         return;
     }
 
-    assert(frozsetwriter_check_consistency(writer));
+    assert(frozensetwriter_check_consistency(writer));
     PyObject *frozenset = _PyObject_CAST(writer);
     Py_DECREF(frozenset);
 }
@@ -3284,7 +3285,7 @@ PyFrozenSetWriter_Discard(PyFrozenSetWriter *writer)
 PyObject*
 PyFrozenSetWriter_Finish(PyFrozenSetWriter *writer)
 {
-    assert(frozsetwriter_check_consistency(writer));
+    assert(frozensetwriter_check_consistency(writer));
     PySetObject *frozenset = _PySet_CAST(writer);
 
     // Now that the frozenset is fully initialized, it's safe to track it in
@@ -3306,7 +3307,7 @@ PyFrozenSetWriter_Finish(PyFrozenSetWriter *writer)
 int
 PyFrozenSetWriter_AddSteal(PyFrozenSetWriter *writer, PyObject *item)
 {
-    assert(frozsetwriter_check_consistency(writer));
+    assert(frozensetwriter_check_consistency(writer));
 
     Py_hash_t hash = PyObject_Hash(item);
     if (hash == -1) {
@@ -3330,7 +3331,78 @@ PyFrozenSetWriter_Add(PyFrozenSetWriter *writer, PyObject *item)
 int
 PyFrozenSetWriter_Update(PyFrozenSetWriter *writer, PyObject *iterable)
 {
-    assert(frozsetwriter_check_consistency(writer));
+    assert(frozensetwriter_check_consistency(writer));
     PySetObject *frozenset = _PySet_CAST(writer);
     return set_update_local(frozenset, iterable);
 }
+
+
+static inline int
+frozensetwriter_update_array(PyFrozenSetWriter *writer,
+                             PyObject *const *items, Py_ssize_t nitem,
+                             int steal)
+{
+    if (nitem < 0) {
+        PyErr_SetString(PyExc_ValueError, "nitem must be >= 0");
+        return -1;
+    }
+
+#define CLEANUP(START) \
+        if (steal) { \
+            for (Py_ssize_t j = (START); j < nitem; j++) { \
+                Py_DECREF(items[j]); \
+            } \
+        }
+
+    assert(frozensetwriter_check_consistency(writer));
+    PySetObject *frozenset = _PySet_CAST(writer);
+
+    /* Do one big resize at the start, rather than
+     * incrementally resizing as we insert new keys.
+     */
+    if ((frozenset->fill + nitem) * 5 >= frozenset->mask * 3) {
+        if (set_table_resize(frozenset, (frozenset->used + nitem) * 2) != 0) {
+            CLEANUP(0);
+            return -1;
+        }
+    }
+
+    for (Py_ssize_t i = 0; i < nitem; i++) {
+        PyObject *item = items[i];
+
+        Py_hash_t hash = PyObject_Hash(item);
+        if (hash == -1) {
+            set_unhashable_type(item);
+            // CLEANUP() calls Py_DECREF(item)
+            CLEANUP(i);
+            return -1;
+        }
+
+        if (!steal) {
+            item = Py_NewRef(item);
+        }
+        if (set_add_entry_takeref(frozenset, item, hash) < 0) {
+            CLEANUP(i);
+            return -1;
+        }
+    }
+    return 0;
+
+#undef CLEANUP
+}
+
+
+int
+PyFrozenSetWriter_UpdateArray(PyFrozenSetWriter *writer,
+                              PyObject *const *items, Py_ssize_t nitem)
+{
+    return frozensetwriter_update_array(writer, items, nitem, 0);
+}
+
+
+int
+PyFrozenSetWriter_UpdateArraySteal(PyFrozenSetWriter *writer,
+                                   PyObject *const *items, Py_ssize_t nitem)
+{
+    return frozensetwriter_update_array(writer, items, nitem, 1);
+}