Complete test_discard()

Some cleanups.

Author: vstinner

Doc/c-api/set.rst

@@ -245,7 +245,8 @@ See also the :c:func:`PyFrozenSet_New` function.
 
    Create a :c:type:`PyFrozenSetWriter`.
 
-   If *size* is greater than zero, preallocate *size* items.
+   If *size* is greater than zero, allocate an internal table big enough to
+   store *size* items.
 
    On error, set an exception and return ``NULL``.
 

Lib/test/test_capi/test_set.py

@@ -323,7 +323,7 @@ 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):
+            with self.subTest(prealloc=prealloc):
                 writer = self.create_writer()
                 add = getattr(writer, add_method)
                 for item in items:
@@ -356,7 +356,7 @@ def check_update_method(self, update_method):
         empty = ()
 
         for prealloc in (0, total // 2, total, total * 10):
-            with self.subTest(size=prealloc):
+            with self.subTest(prealloc=prealloc):
                 writer = self.create_writer(prealloc)
                 update = getattr(writer, update_method)
                 update(items + common)
@@ -412,10 +412,21 @@ class MyObject:
 
     def test_discard(self):
         # Test PyFrozenSetWriter_Discard(): must not leak memory
-        writer = self.create_writer()
-        writer.update([2, 5, 7])
+
+        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

Objects/setobject.c

@@ -3287,13 +3287,13 @@ PyObject*
 PyFrozenSetWriter_Finish(PyFrozenSetWriter *writer)
 {
     assert(frozensetwriter_check_consistency(writer));
-    PySetObject *frozenset = _PySet_CAST(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 _PyObject_CAST(frozenset);
+    return frozenset;
 }