diff --git a/Doc/c-api/set.rst b/Doc/c-api/set.rst index db537aff2e6ce5a..7d3acd06412bebc 100644 --- a/Doc/c-api/set.rst +++ b/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 `. + .. note:: The operation is atomic on :term:`free threading ` @@ -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. diff --git a/Doc/whatsnew/3.16.rst b/Doc/whatsnew/3.16.rst index bd92a0a48f1f70c..7ebd9d7d4addf6e 100644 --- a/Doc/whatsnew/3.16.rst +++ b/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 ` 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`.) diff --git a/Include/cpython/setobject.h b/Include/cpython/setobject.h index 89565cb29212fc1..9c6ac18066b791c 100644 --- a/Include/cpython/setobject.h +++ b/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); diff --git a/Lib/test/test_capi/test_set.py b/Lib/test/test_capi/test_set.py index a41c2fcaf12270f..64447b57a418397 100644 --- a/Lib/test/test_capi/test_set.py +++ b/Lib/test/test_capi/test_set.py @@ -1,6 +1,9 @@ +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. @@ -8,6 +11,7 @@ _testlimitedcapi = import_helper.import_module('_testlimitedcapi') _testinternalcapi = import_helper.import_module('_testinternalcapi') + class set_subclass(set): pass @@ -273,5 +277,175 @@ 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(): track the frozenset in the GC + writer = self.create_writer() + writer.update([2, 5, 7]) + self.assertTrue(gc.is_tracked(writer.finish())) + + class MyObject: + pass + + obj = MyObject() + 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() diff --git a/Misc/NEWS.d/next/C_API/2026-10-21-18-20-57.gh-issue-160147.7uo4pF.rst b/Misc/NEWS.d/next/C_API/2026-10-21-18-20-57.gh-issue-160147.7uo4pF.rst new file mode 100644 index 000000000000000..d2d0277ee8da281 --- /dev/null +++ b/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 ` 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. diff --git a/Modules/_abc.c b/Modules/_abc.c index bca9066d340ee43..65be6d653127bcd 100644 --- a/Modules/_abc.c +++ b/Modules/_abc.c @@ -362,7 +362,7 @@ static int compute_abstract_methods(PyObject *self) { int ret = -1; - PyObject *abstracts = PyFrozenSet_New(NULL); + PyFrozenSetWriter *abstracts = PyFrozenSetWriter_Create(0); if (abstracts == NULL) { return -1; } @@ -404,7 +404,7 @@ compute_abstract_methods(PyObject *self) Py_INCREF(key); int is_abstract = _PyObject_IsAbstract(value); if (is_abstract < 0 || - (is_abstract && PySet_Add(abstracts, key) < 0)) { + (is_abstract && PyFrozenSetWriter_Add(abstracts, key) < 0)) { Py_DECREF(it); Py_DECREF(key); goto error; @@ -454,7 +454,7 @@ compute_abstract_methods(PyObject *self) int is_abstract = _PyObject_IsAbstract(value); Py_DECREF(value); if (is_abstract < 0 || - (is_abstract && PySet_Add(abstracts, key) < 0)) + (is_abstract && PyFrozenSetWriter_Add(abstracts, key) < 0)) { Py_DECREF(key); Py_DECREF(iter); @@ -468,13 +468,20 @@ compute_abstract_methods(PyObject *self) } } - if (PyObject_SetAttr(self, &_Py_ID(__abstractmethods__), abstracts) < 0) { + PyObject *abstracts_set = PyFrozenSetWriter_Finish(abstracts); + abstracts = NULL; + if (abstracts_set == NULL) { goto error; } + if (PyObject_SetAttr(self, &_Py_ID(__abstractmethods__), abstracts_set) < 0) { + Py_DECREF(abstracts_set); + goto error; + } + Py_DECREF(abstracts_set); ret = 0; error: - Py_DECREF(abstracts); + PyFrozenSetWriter_Discard(abstracts); Py_XDECREF(ns); Py_XDECREF(items); Py_XDECREF(bases); diff --git a/Modules/_asynciomodule.c b/Modules/_asynciomodule.c index 5569f7a3670378b..77ecbaf55c766d3 100644 --- a/Modules/_asynciomodule.c +++ b/Modules/_asynciomodule.c @@ -1339,15 +1339,15 @@ _asyncio_Future__asyncio_awaited_by_get_impl(FutureObj *self) return PyFrozenSet_New(self->fut_awaited_by); } - PyObject *set = PyFrozenSet_New(NULL); + PyFrozenSetWriter *set = PyFrozenSetWriter_Create(1); if (set == NULL) { return NULL; } - if (PySet_Add(set, self->fut_awaited_by)) { - Py_DECREF(set); + if (PyFrozenSetWriter_Add(set, self->fut_awaited_by)) { + PyFrozenSetWriter_Discard(set); return NULL; } - return set; + return PyFrozenSetWriter_Finish(set); } diff --git a/Modules/_hashopenssl.c b/Modules/_hashopenssl.c index 9d7602b39d7b938..ba5272acd6e0082 100644 --- a/Modules/_hashopenssl.c +++ b/Modules/_hashopenssl.c @@ -2477,7 +2477,7 @@ PyType_Spec HMACtype_spec = { /* State for our callback function so that it can accumulate a result. */ typedef struct _internal_name_mapper_state { - PyObject *set; + PyFrozenSetWriter *set; int error; } _InternalNameMapperState; @@ -2505,11 +2505,11 @@ _openssl_hash_name_mapper(const EVP_MD *md, const char *from, py_name = name == NULL ? NULL : PyUnicode_FromString(name); if (py_name == NULL) { state->error = 1; - } else { - if (PySet_Add(state->set, py_name) != 0) { - state->error = 1; - } - Py_DECREF(py_name); + return; + } + + if (PyFrozenSetWriter_AddSteal(state->set, py_name) != 0) { + state->error = 1; } } @@ -2519,7 +2519,7 @@ static int hashlib_md_meth_names(PyObject *module) { _InternalNameMapperState state = { - .set = PyFrozenSet_New(NULL), + .set = PyFrozenSetWriter_Create(0), .error = 0 }; if (state.set == NULL) { @@ -2534,11 +2534,12 @@ hashlib_md_meth_names(PyObject *module) #endif if (state.error) { - Py_DECREF(state.set); + PyFrozenSetWriter_Discard(state.set); return -1; } - return PyModule_Add(module, "openssl_md_meth_names", state.set); + PyObject *names = PyFrozenSetWriter_Finish(state.set); + return PyModule_Add(module, "openssl_md_meth_names", names); } /*[clinic input] diff --git a/Modules/_ssl.c b/Modules/_ssl.c index ee0df146421d7f8..b2f7c4f9c04e34a 100644 --- a/Modules/_ssl.c +++ b/Modules/_ssl.c @@ -6506,7 +6506,6 @@ parseKeyUsage(PCCERT_CONTEXT pCertCtx, DWORD flags) { CERT_ENHKEY_USAGE *usage; DWORD size, error, i; - PyObject *retval; if (!CertGetEnhancedKeyUsage(pCertCtx, flags, NULL, &size)) { error = GetLastError(); @@ -6530,30 +6529,30 @@ parseKeyUsage(PCCERT_CONTEXT pCertCtx, DWORD flags) } return PyErr_SetFromWindowsErr(error); } - retval = PyFrozenSet_New(NULL); + PyFrozenSetWriter *retval = PyFrozenSetWriter_Create(usage->cUsageIdentifier); if (retval == NULL) { goto error; } for (i = 0; i < usage->cUsageIdentifier; ++i) { - if (usage->rgpszUsageIdentifier[i]) { - PyObject *oid; - int err; - oid = PyUnicode_FromString(usage->rgpszUsageIdentifier[i]); - if (oid == NULL) { - Py_CLEAR(retval); - goto error; - } - err = PySet_Add(retval, oid); - Py_DECREF(oid); - if (err == -1) { - Py_CLEAR(retval); - goto error; - } + if (usage->rgpszUsageIdentifier[i] == NULL) { + continue; + } + + PyObject *oid = PyUnicode_FromString(usage->rgpszUsageIdentifier[i]); + if (oid == NULL) { + goto error; + } + if (PyFrozenSetWriter_AddSteal(retval, oid) < 0) { + goto error; } } + PyMem_Free(usage); + return PyFrozenSetWriter_Finish(retval); + error: PyMem_Free(usage); - return retval; + PyFrozenSetWriter_Discard(retval); + return NULL; } static HCERTSTORE diff --git a/Modules/_testcapi/set.c b/Modules/_testcapi/set.c index 092715ab7d0aa40..8226203115eb4c4 100644 --- a/Modules/_testcapi/set.c +++ b/Modules/_testcapi/set.c @@ -31,14 +31,276 @@ test_set_type_size(PyObject *self, PyObject *Py_UNUSED(ignored)) } +typedef struct { + PyObject_HEAD + PyFrozenSetWriter *writer; +} WriterObject; + + +static inline int +writer_check(WriterObject *self) +{ + if (self->writer == NULL) { + PyErr_SetString(PyExc_ValueError, + "operation on finished/discarded writer"); + return -1; + } + return 0; +} + + +static PyObject * +writer_new(PyTypeObject *type, PyObject *args, PyObject *kwargs) +{ + WriterObject *self = (WriterObject *)type->tp_alloc(type, 0); + if (!self) { + return NULL; + } + self->writer = NULL; + return (PyObject*)self; +} + + +// Test PyFrozenSetWriter_Create() +static int +writer_init(PyObject *self_raw, PyObject *args, PyObject *kwargs) +{ + if (kwargs && PyDict_GET_SIZE(kwargs)) { + PyErr_Format(PyExc_TypeError, + "PyFrozenSetWriter() takes no keyword arguments"); + return -1; + } + + Py_ssize_t size; + if (!PyArg_ParseTuple(args, "n", &size)) { + return -1; + } + + WriterObject *self = (WriterObject *)self_raw; + PyFrozenSetWriter_Discard(self->writer); + + self->writer = PyFrozenSetWriter_Create(size); + if (self->writer == NULL) { + return -1; + } + + return 0; +} + + +// Test PyFrozenSetWriter_Discard() +static void +writer_dealloc(PyObject *self_raw) +{ + WriterObject *self = (WriterObject *)self_raw; + PyTypeObject *tp = Py_TYPE(self); + if (self->writer) { + PyFrozenSetWriter_Discard(self->writer); + } + tp->tp_free(self); + Py_DECREF(tp); +} + + +// Test PyFrozenSetWriter_Add() +static PyObject* +writer_add(PyObject *self_raw, PyObject *item) +{ + WriterObject *self = (WriterObject *)self_raw; + if (writer_check(self) < 0) { + return NULL; + } + + if (PyFrozenSetWriter_Add(self->writer, item) < 0) { + return NULL; + } + Py_RETURN_NONE; +} + + +// Test PyFrozenSetWriter_AddSteal() +static PyObject* +writer_add_steal(PyObject *self_raw, PyObject *item) +{ + WriterObject *self = (WriterObject *)self_raw; + if (writer_check(self) < 0) { + return NULL; + } + + if (PyFrozenSetWriter_AddSteal(self->writer, Py_NewRef(item)) < 0) { + return NULL; + } + Py_RETURN_NONE; +} + + +// Test PyFrozenSetWriter_Update() +static PyObject* +writer_update(PyObject *self_raw, PyObject *iterable) +{ + WriterObject *self = (WriterObject *)self_raw; + if (writer_check(self) < 0) { + return NULL; + } + + if (PyFrozenSetWriter_Update(self->writer, iterable) < 0) { + return NULL; + } + Py_RETURN_NONE; +} + + +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)) +{ + WriterObject *self = (WriterObject *)self_raw; + if (writer_check(self) < 0) { + return NULL; + } + + PySetObject *frozenset = _PySet_CAST(self->writer); + Py_ssize_t table_size = frozenset->mask + 1; + return PyLong_FromSsize_t(table_size); +} + + +// Test PyFrozenSetWriter_Finish() +static PyObject* +writer_finish(PyObject *self_raw, PyObject *Py_UNUSED(args)) +{ + WriterObject *self = (WriterObject *)self_raw; + if (writer_check(self) < 0) { + return NULL; + } + + PyObject *frozenset = PyFrozenSetWriter_Finish(self->writer); + self->writer = NULL; + return frozenset; +} + + +// Test PyFrozenSetWriter_Discard() +static PyObject* +writer_discard(PyObject *self_raw, PyObject *Py_UNUSED(args)) +{ + WriterObject *self = (WriterObject *)self_raw; + if (writer_check(self) < 0) { + return NULL; + } + + PyFrozenSetWriter_Discard(self->writer); + self->writer = NULL; + Py_RETURN_NONE; +} + + +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}, + {NULL, NULL} /* sentinel */ +}; + +static PyType_Slot Writer_Type_slots[] = { + {Py_tp_new, writer_new}, + {Py_tp_init, writer_init}, + {Py_tp_dealloc, writer_dealloc}, + {Py_tp_methods, writer_methods}, + {0, 0}, /* sentinel */ +}; + +static PyType_Spec Writer_spec = { + .name = "_testcapi.PyFrozenSetWriter", + .basicsize = sizeof(WriterObject), + .flags = Py_TPFLAGS_DEFAULT, + .slots = Writer_Type_slots, +}; + + static PyMethodDef test_methods[] = { {"set_get_size", set_get_size, METH_O}, {"test_set_type_size", test_set_type_size, METH_NOARGS}, {NULL}, }; + int _PyTestCapi_Init_Set(PyObject *m) { - return PyModule_AddFunctions(m, test_methods); + if (PyModule_AddFunctions(m, test_methods) < 0) { + return -1; + } + + PyTypeObject *writer_type = (PyTypeObject *)PyType_FromSpec(&Writer_spec); + if (writer_type == NULL) { + return -1; + } + if (PyModule_AddType(m, writer_type) < 0) { + Py_DECREF(writer_type); + return -1; + } + Py_DECREF(writer_type); + + if (PyModule_AddIntMacro(m, PySet_MINSIZE) < 0) { + return -1; + } + + return 0; } diff --git a/Objects/setobject.c b/Objects/setobject.c index a762edf5933e185..84c736204d47271 100644 --- a/Objects/setobject.c +++ b/Objects/setobject.c @@ -3223,3 +3223,179 @@ static PyTypeObject _PySetDummy_Type = { }; static PyObject _dummy_struct = _PyObject_HEAD_INIT(&_PySetDummy_Type); + + +// --- PyFrozenSetWriter API ------------------------------------------------ + +#ifndef NDEBUG +static int +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); + assert(!_PyObject_GC_IS_TRACKED(frozenset)); + return 1; +} +#endif + + +PyFrozenSetWriter* +PyFrozenSetWriter_Create(Py_ssize_t size) +{ + if (size < 0) { + PyErr_SetString(PyExc_ValueError, "size must be >= 0"); + return NULL; + } + + PyObject *frozenset_obj = make_new_set_untracked(&PyFrozenSet_Type, NULL); + if (frozenset_obj == NULL) { + return NULL; + } + PySetObject *frozenset = _PySet_CAST(frozenset_obj); + + if (size > 0 && size * 5 >= frozenset->mask * 3) { + if (set_table_resize(frozenset, size * 2) != 0) { + Py_DECREF(frozenset); + return NULL; + } + } + + PyFrozenSetWriter *writer = (PyFrozenSetWriter*)frozenset; + assert(frozensetwriter_check_consistency(writer)); + return writer; +} + + +void +PyFrozenSetWriter_Discard(PyFrozenSetWriter *writer) +{ + if (writer == NULL) { + return; + } + + assert(frozensetwriter_check_consistency(writer)); + PyObject *frozenset = _PyObject_CAST(writer); + Py_DECREF(frozenset); +} + + +PyObject* +PyFrozenSetWriter_Finish(PyFrozenSetWriter *writer) +{ + assert(frozensetwriter_check_consistency(writer)); + PyObject *frozenset = _PyObject_CAST(writer); + + // Now that the frozenset is fully initialized, it's safe to track it + // in the garbage collector. + _PyObject_GC_TRACK(frozenset); + + return frozenset; +} + + +int +PyFrozenSetWriter_AddSteal(PyFrozenSetWriter *writer, PyObject *item) +{ + assert(frozensetwriter_check_consistency(writer)); + + Py_hash_t hash = PyObject_Hash(item); + if (hash == -1) { + set_unhashable_type(item); + Py_DECREF(item); + return -1; + } + + PySetObject *frozenset = _PySet_CAST(writer); + return set_add_entry_takeref(frozenset, item, hash); +} + + +int +PyFrozenSetWriter_Add(PyFrozenSetWriter *writer, PyObject *item) +{ + return PyFrozenSetWriter_AddSteal(writer, Py_NewRef(item)); +} + + +int +PyFrozenSetWriter_Update(PyFrozenSetWriter *writer, PyObject *iterable) +{ + 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); +}