gh-160147: Add PyFrozenSetWriter C API

Doc/c-api/set.rst

@@ -102,6 +102,8 @@ the constructor functions work with any iterable Python object.
    set on success or ``NULL`` on failure.  Raise :exc:`TypeError` if *iterable* is
    not actually iterable.
 
+   See also the :ref:`PyFrozenSetWriter API <pyfrozensetwriter>`.
+
    .. note::
 
       The operation is atomic on :term:`free threading <free-threaded build>`
@@ -213,3 +215,97 @@ Deprecated API
    If looking for the size of a set, use :c:func:`PySet_Size` instead.
 
    .. soft-deprecated:: 3.14
+
+
+.. _pyfrozensetwriter:
+
+PyFrozenSetWriter
+-----------------
+
+The :c:type:`PyFrozenSetWriter` API can be used to create a Python :class:`frozenset`
+object.
+
+See also the :c:func:`PyFrozenSet_New` function.
+
+.. versionadded:: next
+
+
+.. c:type:: PyFrozenSetWriter
+
+   A frozenset writer object.
+
+   The API is **not thread safe**. To share a writer with multiple threads, a
+   critical section or a lock is needed.
+
+   The instance must be destroyed by :c:func:`PyFrozenSetWriter_Finish` on
+   success, or :c:func:`PyFrozenSetWriter_Discard` on error.
+
+
+.. c:function:: PyFrozenSetWriter* PyFrozenSetWriter_Create(Py_ssize_t size)
+
+   Create a :c:type:`PyFrozenSetWriter`.
+
+   If *size* is greater than zero, allocate an internal table big enough to
+   store *size* items.
+
+   On error, set an exception and return ``NULL``.
+
+   *size* must be positive or zero.
+
+.. c:function:: PyObject* PyFrozenSetWriter_Finish(PyFrozenSetWriter *writer)
+
+   Finish a :c:type:`PyFrozenSetWriter` created by
+   :c:func:`PyFrozenSetWriter_Create`.
+
+   On success, return a Python :class:`frozenset` object.
+   On error, set an exception and return ``NULL``.
+
+   The writer instance is invalid after the call in any case.
+   No API can be called on the writer after :c:func:`PyFrozenSetWriter_Finish`.
+
+.. c:function:: void PyFrozenSetWriter_Discard(PyFrozenSetWriter *writer)
+
+   Discard a :c:type:`PyFrozenSetWriter` created by :c:func:`PyFrozenSetWriter_Create`.
+
+   Do nothing if *writer* is ``NULL``.
+
+   The writer instance is invalid after the call.
+   No API can be called on the writer after :c:func:`PyFrozenSetWriter_Discard`.
+
+.. c:function:: int PyFrozenSetWriter_Add(PyFrozenSetWriter *writer, PyObject *item)
+
+   Add *item* to a frozenset writer.
+
+   On success, return ``0``.
+   On error, set an exception and return ``-1``.
+
+.. c:function:: int PyFrozenSetWriter_AddSteal(PyFrozenSetWriter *writer, PyObject *item)
+
+   Similar to :c:func:`PyFrozenSetWriter_Add`, but :term:`steal` the *item*
+   reference.
+
+.. c:function:: int PyFrozenSetWriter_Update(PyFrozenSetWriter *writer, PyObject *iterable)
+
+   Update the frozenset writer, adding elements from *iterable*.
+
+   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

@@ -1295,6 +1295,20 @@ New features
 * Add :c:func:`Py_HashBuffer` to the limited C API.
   (Contributed by Victor Stinner in :gh:`148233`.)
 
+* Add a new :ref:`PyFrozenSetWriter <pyfrozensetwriter>` C API to create
+  a :class:`frozenset` object. Add functions:
+
+  * :c:func:`PyFrozenSetWriter_Create`
+  * :c:func:`PyFrozenSetWriter_Finish`
+  * :c:func:`PyFrozenSetWriter_Discard`
+  * :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`.)
+
 * Add :c:macro:`Py_SETREF` and :c:macro:`Py_XSETREF` macros to the limited C
   API.
   (Contributed by Nathan Goldbaum and Victor Stinner in :gh:`158842`.)

Include/cpython/setobject.h

@@ -69,3 +69,29 @@ static inline Py_ssize_t PySet_GET_SIZE(PyObject *so) {
 #endif
 }
 #define PySet_GET_SIZE(so) PySet_GET_SIZE(_PyObject_CAST(so))
+
+
+// --- PyFrozenSetWriter API ------------------------------------------------
+
+typedef struct PyFrozenSetWriter PyFrozenSetWriter;
+
+PyAPI_FUNC(PyFrozenSetWriter*) PyFrozenSetWriter_Create(Py_ssize_t size);
+PyAPI_FUNC(PyObject*) PyFrozenSetWriter_Finish(PyFrozenSetWriter *writer);
+PyAPI_FUNC(void) PyFrozenSetWriter_Discard(PyFrozenSetWriter *writer);
+PyAPI_FUNC(int) PyFrozenSetWriter_Add(
+    PyFrozenSetWriter *writer,
+    PyObject *item);
+PyAPI_FUNC(int) PyFrozenSetWriter_AddSteal(
+    PyFrozenSetWriter *writer,
+    PyObject *item);
+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

@@ -1,13 +1,17 @@
+import gc
+import threading
 import unittest
 
 from test.support import import_helper
+from test.support import threading_helper
 
 # Skip this test if the _testcapi, _testlimitedcapi or _testinternalcapi
 # modules aren't available.
 _testcapi = import_helper.import_module('_testcapi')
 _testlimitedcapi = import_helper.import_module('_testlimitedcapi')
 _testinternalcapi = import_helper.import_module('_testinternalcapi')
 
+
 class set_subclass(set):
     pass
 
@@ -273,5 +277,182 @@ def test_set_next_entry(self):
         # CRASHES: set_next(NULL, 0)
 
 
+class FrozenSetWriterTests(unittest.TestCase):
+    PySet_MINSIZE = _testcapi.PySet_MINSIZE
+    UNHASHABLE = [1, 2, 3]
+
+    def create_writer(self, size=0):
+        # Test PyFrozenSetWriter_Create()
+        return _testcapi.PyFrozenSetWriter(size)
+
+    def compute_table_size(self, size):
+        table_size = self.PySet_MINSIZE
+        while table_size <= size:
+            table_size <<= 1
+        return table_size
+
+    def frozenset_equal(self, result, expected):
+        if not isinstance(expected, frozenset):
+            expected = frozenset(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
+
+        def check_table_size(size, table_size):
+            writer = self.create_writer(size)
+            self.assertEqual(writer.get_table_size(), table_size)
+
+        check_table_size(0, PySet_MINSIZE)
+        check_table_size(PySet_MINSIZE // 2, PySet_MINSIZE)
+        check_table_size(PySet_MINSIZE - 1, PySet_MINSIZE * 2)
+        check_table_size(PySet_MINSIZE, PySet_MINSIZE * 4)
+        check_table_size(2 ** 15 - 1, 2 ** 16)
+
+    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 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(prealloc=prealloc):
+                writer = self.create_writer()
+                add = getattr(writer, add_method)
+                for item in items:
+                    add(item)
+                self.frozenset_equal(writer.finish(), items)
+
+        # Test error (unhashable type)
+        writer = self.create_writer(0)
+        writer.update('abc')
+        add = getattr(writer, add_method)
+        with self.assert_unhashable():
+            add(self.UNHASHABLE)
+        self.frozenset_equal(writer.finish(), frozenset('abc'))
+
+    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(items) | set(items2) | set(common)
+        total = len(expected)
+        self.assertEqual(total, len(items) + len(items2) + len(common))
+        empty = ()
+
+        for prealloc in (0, total // 2, total, total * 10):
+            with self.subTest(prealloc=prealloc):
+                writer = self.create_writer(prealloc)
+                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(): only track the frozenset in the GC
+        # if at least one item is tracked by the GC
+
+        # Integers are not tracked by the GC
+        items = [2, 5, 7]
+        self.assertTrue(all(not gc.is_tracked(item) for item in items))
+        writer = self.create_writer()
+        writer.update(items)
+        self.assertFalse(gc.is_tracked(writer.finish()))
+
+        # Create an object tracked by the GC
+        class MyObject:
+            pass
+        obj = MyObject()
+        self.assertTrue(gc.is_tracked(obj))
+
+        writer = self.create_writer()
+        writer.add(obj)
+        self.assertTrue(gc.is_tracked(writer.finish()))
+
+        # test again, now using PyFrozenSetWriter_Update()
+        writer = self.create_writer()
+        writer.update([obj])
+        self.assertTrue(gc.is_tracked(writer.finish()))
+
+    def test_discard(self):
+        # Test PyFrozenSetWriter_Discard(): must not leak memory
+
+        writer = self.create_writer(0)
+        writer.discard()
+
+        for prealloc in (0, 3, 100):
+            for items in (
+                (),          # empty
+                [2, 5, 7],   # use the small table
+                range(100),  # allocate a new table
+            ):
+                with self.subTest(prealloc=prealloc, items=items):
+                    writer = self.create_writer(prealloc)
+                    writer.update(items)
+                    writer.discard()
+
+    @threading_helper.requires_working_threading()
+    def test_thread(self):
+        # PyFrozenSetWriter can be used by multiple threads: it's up to the caller
+        # to implement a lock to prevent concurrent accesses.
+        writer = self.create_writer()
+
+        def thread_func(writer):
+            writer.add('a')
+            writer.add_steal('b')
+            writer.update(('c',))
+            writer.update_array(('d',))
+            writer.update_array_steal(('e',))
+
+        thread = threading.Thread(target=thread_func, args=(writer,))
+        thread.start()
+        threading_helper.join_thread(thread)
+
+        self.frozenset_equal(writer.finish(), frozenset('abcde'))
+
+
 if __name__ == "__main__":
     unittest.main()

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

@@ -0,0 +1,13 @@
+Add a new :ref:`PyFrozenSetWriter <pyfrozensetwriter>` C API to create a
+:class:`frozenset` object. Add functions:
+
+* :c:func:`PyFrozenSetWriter_Create`
+* :c:func:`PyFrozenSetWriter_Finish`
+* :c:func:`PyFrozenSetWriter_Discard`
+* :c:func:`PyFrozenSetWriter_Add`
+* :c:func:`PyFrozenSetWriter_AddSteal`
+* :c:func:`PyFrozenSetWriter_Update`
+* :c:func:`PyFrozenSetWriter_UpdateArray`
+* :c:func:`PyFrozenSetWriter_UpdateArraySteal`
+
+Patch by Victor Stinner.