From c56bca31e9b7ac87a459a65680ee7ad454fd4f22 Mon Sep 17 00:00:00 2001
From: Antoine Pitrou <solipsis@pitrou.net>
Date: Thu, 1 Mar 2012 16:26:35 +0100
Subject: Issue #14159: Fix the len() of weak sets 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.

---
 Lib/_weakrefset.py       |  2 +-
 Lib/test/test_weakref.py | 60 ++++++++++++++++++++++++++++++++++++++++++++++++
 Lib/test/test_weakset.py | 47 +++++++++++++++++++++++++++++++++++++
 Misc/NEWS                |  4 ++++
 4 files changed, 112 insertions(+), 1 deletion(-)

diff --git a/Lib/_weakrefset.py b/Lib/_weakrefset.py
index ffa5e31..91077c5 100644
--- a/Lib/_weakrefset.py
+++ b/Lib/_weakrefset.py
@@ -63,7 +63,7 @@ class WeakSet(object):
                     yield item
 
     def __len__(self):
-        return sum(x() is not None for x in self.data)
+        return len(self.data) - len(self._pending_removals)
 
     def __contains__(self, item):
         try:
diff --git a/Lib/test/test_weakref.py b/Lib/test/test_weakref.py
index bc2982f..f2e328b 100644
--- a/Lib/test/test_weakref.py
+++ b/Lib/test/test_weakref.py
@@ -815,11 +815,71 @@ class Object:
     def __repr__(self):
         return "<Object %r>" % self.arg
 
+class RefCycle:
+    def __init__(self):
+        self.cycle = self
+
 
 class MappingTestCase(TestBase):
 
     COUNT = 10
 
+    def check_len_cycles(self, dict_type, cons):
+        N = 20
+        items = [RefCycle() for i in range(N)]
+        dct = dict_type(cons(o) for o in items)
+        # Keep an iterator alive
+        it = dct.iteritems()
+        try:
+            next(it)
+        except StopIteration:
+            pass
+        del items
+        gc.collect()
+        n1 = len(dct)
+        del it
+        gc.collect()
+        n2 = len(dct)
+        # one item may be kept alive inside the iterator
+        self.assertIn(n1, (0, 1))
+        self.assertEqual(n2, 0)
+
+    def test_weak_keyed_len_cycles(self):
+        self.check_len_cycles(weakref.WeakKeyDictionary, lambda k: (k, 1))
+
+    def test_weak_valued_len_cycles(self):
+        self.check_len_cycles(weakref.WeakValueDictionary, lambda k: (1, k))
+
+    def check_len_race(self, dict_type, cons):
+        # 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)]
+            dct = dict_type(cons(o) for o in items)
+            del items
+            # All items will be collected at next garbage collection pass
+            it = dct.iteritems()
+            try:
+                next(it)
+            except StopIteration:
+                pass
+            n1 = len(dct)
+            del it
+            n2 = len(dct)
+            self.assertGreaterEqual(n1, 0)
+            self.assertLessEqual(n1, N)
+            self.assertGreaterEqual(n2, 0)
+            self.assertLessEqual(n2, n1)
+
+    def test_weak_keyed_len_race(self):
+        self.check_len_race(weakref.WeakKeyDictionary, lambda k: (k, 1))
+
+    def test_weak_valued_len_race(self):
+        self.check_len_race(weakref.WeakValueDictionary, lambda k: (1, k))
+
     def test_weak_values(self):
         #
         #  This exercises d.copy(), d.items(), d[], del d[], len(d).
diff --git a/Lib/test/test_weakset.py b/Lib/test/test_weakset.py
index 89c2822..f981bdd 100644
--- a/Lib/test/test_weakset.py
+++ b/Lib/test/test_weakset.py
@@ -30,6 +30,10 @@ class SomeClass(object):
     def __hash__(self):
         return hash((SomeClass, self.value))
 
+class RefCycle(object):
+    def __init__(self):
+        self.cycle = self
+
 class TestWeakSet(unittest.TestCase):
 
     def setUp(self):
@@ -369,6 +373,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):
     test_support.run_unittest(TestWeakSet)
diff --git a/Misc/NEWS b/Misc/NEWS
index 8460cdd..6b3fbc2 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -101,6 +101,10 @@ Core and Builtins
 Library
 -------
 
+- Issue #14159: Fix the len() of weak sets to return a better approximation
+  when some objects are dead or dying.  Moreover, the implementation is now
+  O(1) rather than O(n).
+
 - Issue #2945: Make the distutils upload command aware of bdist_rpm products.
 
 - Issue #13447: Add a test file to host regression tests for bugs in the
-- 
cgit v0.12