Issue #14159: Fix the len() of weak containers (WeakSet, WeakKeyDictionary, WeakValueDictionary) to return a better approximation when some objects are dead or dying.

Moreover, the implementation is now O(1) rather than O(n).
Thanks to Yury Selivanov for reporting.

1 files changed, 47 insertions, 0 deletions

diff --git a/Lib/test/test_weakset.py b/Lib/test/test_weakset.py
index 3cddf40..35db7a6 100644
--- a/Lib/test/test_weakset.py
+++ b/Lib/test/test_weakset.py
@@ -17,6 +17,10 @@ import contextlib
 class Foo:
     pass
 
+class RefCycle:
+    def __init__(self):
+        self.cycle = self
+
 
 class TestWeakSet(unittest.TestCase):
 
@@ -359,6 +363,49 @@ class TestWeakSet(unittest.TestCase):
             s.clear()
         self.assertEqual(len(s), 0)
 
+    def test_len_cycles(self):
+        N = 20
+        items = [RefCycle() for i in range(N)]
+        s = WeakSet(items)
+        del items
+        it = iter(s)
+        try:
+            next(it)
+        except StopIteration:
+            pass
+        gc.collect()
+        n1 = len(s)
+        del it
+        gc.collect()
+        n2 = len(s)
+        # one item may be kept alive inside the iterator
+        self.assertIn(n1, (0, 1))
+        self.assertEqual(n2, 0)
+
+    def test_len_race(self):
+        # Extended sanity checks for len() in the face of cyclic collection
+        self.addCleanup(gc.set_threshold, *gc.get_threshold())
+        for th in range(1, 100):
+            N = 20
+            gc.collect(0)
+            gc.set_threshold(th, th, th)
+            items = [RefCycle() for i in range(N)]
+            s = WeakSet(items)
+            del items
+            # All items will be collected at next garbage collection pass
+            it = iter(s)
+            try:
+                next(it)
+            except StopIteration:
+                pass
+            n1 = len(s)
+            del it
+            n2 = len(s)
+            self.assertGreaterEqual(n1, 0)
+            self.assertLessEqual(n1, N)
+            self.assertGreaterEqual(n2, 0)
+            self.assertLessEqual(n2, n1)
+
 
 def test_main(verbose=None):
     support.run_unittest(TestWeakSet)