Add size parameter to PyFrozenSetWriter_Create()

Author: vstinner

Doc/c-api/set.rst

@@ -241,10 +241,12 @@ See also the :c:func:`PyFrozenSet_New` function.
    success, or :c:func:`PyFrozenSetWriter_Discard` on error.
 
 
-.. c:function:: PyFrozenSetWriter* PyFrozenSetWriter_Create(void)
+.. c:function:: PyFrozenSetWriter* PyFrozenSetWriter_Create(Py_ssize_t size)
 
    Create a :c:type:`PyFrozenSetWriter`.
 
+   If *size* is greater than zero, preallocate *size* items.
+
    On error, set an exception and return ``NULL``.
 
    *size* must be positive or zero.

Include/cpython/setobject.h

@@ -75,7 +75,7 @@ static inline Py_ssize_t PySet_GET_SIZE(PyObject *so) {
 
 typedef struct PyFrozenSetWriter PyFrozenSetWriter;
 
-PyAPI_FUNC(PyFrozenSetWriter*) PyFrozenSetWriter_Create(void);
+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(

Lib/test/test_capi/test_set.py

@@ -278,30 +278,67 @@ def test_set_next_entry(self):
 
 
 class FrozenSetWriterTests(unittest.TestCase):
-    def create_writer(self):
+    PySet_MINSIZE = _testcapi.PySet_MINSIZE
+
+    def create_writer(self, size=0):
         # Test PyFrozenSetWriter_Create()
-        return _testcapi.PyFrozenSetWriter()
+        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 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 - 1, PySet_MINSIZE)
+        check_table_size(PySet_MINSIZE, PySet_MINSIZE * 2)
+        check_table_size(2 ** 15 - 1, 2 ** 15)
+
+    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_append(self):
         # Test PyFrozenSetWriter_Append()
-        items = (1, 2, 3)
-        writer = self.create_writer()
-        for item in items:
-            writer.add(item)
-        self.frozenset_equal(writer.finish(), items)
+        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()
+                for item in items:
+                    writer.add(item)
+                self.frozenset_equal(writer.finish(), items)
 
     def test_update(self):
         # Test PyFrozenSetWriter_Update()
-        writer = self.create_writer()
-        writer.update((2, 5))
-        writer.update((5, 7))
-        self.frozenset_equal(writer.finish(), {2, 5, 7})
+        items1 = 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)
+        total = len(expected)
+        self.assertEqual(total, len(items1) + len(items2) + len(common))
+
+        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)
+                self.frozenset_equal(writer.finish(), expected)
 
     def test_finish(self):
         # Test PyFrozenSetWriter_Finish(): track the frozenset in the GC if at

Modules/_abc.c

@@ -362,7 +362,7 @@ static int
 compute_abstract_methods(PyObject *self)
 {
     int ret = -1;
-    PyFrozenSetWriter *abstracts = PyFrozenSetWriter_Create();
+    PyFrozenSetWriter *abstracts = PyFrozenSetWriter_Create(0);
     if (abstracts == NULL) {
         return -1;
     }

Modules/_asynciomodule.c

@@ -1339,7 +1339,7 @@ _asyncio_Future__asyncio_awaited_by_get_impl(FutureObj *self)
         return PyFrozenSet_New(self->fut_awaited_by);
     }
 
-    PyFrozenSetWriter *set = PyFrozenSetWriter_Create();
+    PyFrozenSetWriter *set = PyFrozenSetWriter_Create(1);
     if (set == NULL) {
         return NULL;
     }

Modules/_hashopenssl.c

@@ -2519,7 +2519,7 @@ static int
 hashlib_md_meth_names(PyObject *module)
 {
     _InternalNameMapperState state = {
-        .set = PyFrozenSetWriter_Create(),
+        .set = PyFrozenSetWriter_Create(0),
         .error = 0
     };
     if (state.set == NULL) {

Modules/_testcapi/set.c

@@ -65,21 +65,21 @@ writer_new(PyTypeObject *type, PyObject *args, PyObject *kwargs)
 static int
 writer_init(PyObject *self_raw, PyObject *args, PyObject *kwargs)
 {
-    if (PyTuple_GET_SIZE(args)) {
-        PyErr_Format(PyExc_TypeError,
-                     "PyFrozenSetWriter() takes no positional arguments");
-        return -1;
-    }
     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();
+    self->writer = PyFrozenSetWriter_Create(size);
     if (self->writer == NULL) {
         return -1;
     }
@@ -134,6 +134,20 @@ writer_update(PyObject *self_raw, PyObject *iterable)
 }
 
 
+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))
@@ -167,6 +181,7 @@ writer_discard(PyObject *self_raw, PyObject *Py_UNUSED(args))
 static PyMethodDef writer_methods[] = {
     {"add", _PyCFunction_CAST(writer_add), METH_O},
     {"update", _PyCFunction_CAST(writer_update), 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 */
@@ -212,5 +227,9 @@ _PyTestCapi_Init_Set(PyObject *m)
     }
     Py_DECREF(writer_type);
 
+    if (PyModule_AddIntMacro(m, PySet_MINSIZE) < 0) {
+        return -1;
+    }
+
     return 0;
 }

Objects/setobject.c

@@ -3243,13 +3243,25 @@ frozsetwriter_check_consistency(PyFrozenSetWriter *writer)
 
 
 PyFrozenSetWriter*
-PyFrozenSetWriter_Create(void)
+PyFrozenSetWriter_Create(Py_ssize_t size)
 {
+    if (size < 0) {
+        PyErr_SetString(PyExc_ValueError, "size must be >= 0");
+        return NULL;
+    }
+
     PyObject *frozenset = make_new_set_untracked(&PyFrozenSet_Type, NULL);
     if (frozenset == NULL) {
         return NULL;
     }
 
+    if (size > 0) {
+        if (set_table_resize(_PySet_CAST(frozenset), size) != 0) {
+            Py_DECREF(frozenset);
+            return NULL;
+        }
+    }
+
     PyFrozenSetWriter *writer = (PyFrozenSetWriter*)frozenset;
     assert(frozsetwriter_check_consistency(writer));
     return writer;