Closes #21173: Fix len() on a WeakKeyDictionary when .clear() was called with an iterator alive.

3 files changed, 52 insertions, 0 deletions

diff --git a/Lib/test/test_weakref.py b/Lib/test/test_weakref.py
index 94c98ad..e735376 100644
--- a/Lib/test/test_weakref.py
+++ b/Lib/test/test_weakref.py
@@ -1310,6 +1310,36 @@ class MappingTestCase(TestBase):
             dict.clear()
         self.assertEqual(len(dict), 0)
 
+    def check_weak_del_and_len_while_iterating(self, dict, testcontext):
+        # Check that len() works when both iterating and removing keys
+        # explicitly through various means (.pop(), .clear()...), while
+        # implicit mutation is deferred because an iterator is alive.
+        # (each call to testcontext() should schedule one item for removal
+        #  for this test to work properly)
+        o = Object(123456)
+        with testcontext():
+            n = len(dict)
+            dict.popitem()
+            self.assertEqual(len(dict), n - 1)
+            dict[o] = o
+            self.assertEqual(len(dict), n)
+        with testcontext():
+            self.assertEqual(len(dict), n - 1)
+            dict.pop(next(dict.keys()))
+            self.assertEqual(len(dict), n - 2)
+        with testcontext():
+            self.assertEqual(len(dict), n - 3)
+            del dict[next(dict.keys())]
+            self.assertEqual(len(dict), n - 4)
+        with testcontext():
+            self.assertEqual(len(dict), n - 5)
+            dict.popitem()
+            self.assertEqual(len(dict), n - 6)
+        with testcontext():
+            dict.clear()
+            self.assertEqual(len(dict), 0)
+        self.assertEqual(len(dict), 0)
+
     def test_weak_keys_destroy_while_iterating(self):
         # Issue #7105: iterators shouldn't crash when a key is implicitly removed
         dict, objects = self.make_weak_keyed_dict()
@@ -1331,6 +1361,10 @@ class MappingTestCase(TestBase):
                 it = None           # should commit all removals
                 gc.collect()
         self.check_weak_destroy_and_mutate_while_iterating(dict, testcontext)
+        # Issue #21173: len() fragile when keys are both implicitly and
+        # explicitly removed.
+        dict, objects = self.make_weak_keyed_dict()
+        self.check_weak_del_and_len_while_iterating(dict, testcontext)
 
     def test_weak_values_destroy_while_iterating(self):
         # Issue #7105: iterators shouldn't crash when a key is implicitly removed
@@ -1354,6 +1388,8 @@ class MappingTestCase(TestBase):
                 it = None           # should commit all removals
                 gc.collect()
         self.check_weak_destroy_and_mutate_while_iterating(dict, testcontext)
+        dict, objects = self.make_weak_valued_dict()
+        self.check_weak_del_and_len_while_iterating(dict, testcontext)
 
     def test_make_weak_keyed_dict_from_dict(self):
         o = Object(3)
diff --git a/Lib/weakref.py b/Lib/weakref.py
index 33b8cbc..a4ecadc 100644
--- a/Lib/weakref.py
+++ b/Lib/weakref.py
@@ -322,6 +322,7 @@ class WeakKeyDictionary(collections.MutableMapping):
         # A list of dead weakrefs (keys to be removed)
         self._pending_removals = []
         self._iterating = set()
+        self._dirty_len = False
         if dict is not None:
             self.update(dict)
 
@@ -338,13 +339,23 @@ class WeakKeyDictionary(collections.MutableMapping):
             except KeyError:
                 pass
 
+    def _scrub_removals(self):
+        d = self.data
+        self._pending_removals = [k for k in self._pending_removals if k in d]
+        self._dirty_len = False
+
     def __delitem__(self, key):
+        self._dirty_len = True
         del self.data[ref(key)]
 
     def __getitem__(self, key):
         return self.data[ref(key)]
 
     def __len__(self):
+        if self._dirty_len and self._pending_removals:
+            # self._pending_removals may still contain keys which were
+            # explicitly removed, we have to scrub them (see issue #21173).
+            self._scrub_removals()
         return len(self.data) - len(self._pending_removals)
 
     def __repr__(self):
@@ -417,6 +428,7 @@ class WeakKeyDictionary(collections.MutableMapping):
         return list(self.data)
 
     def popitem(self):
+        self._dirty_len = True
         while True:
             key, value = self.data.popitem()
             o = key()
@@ -424,6 +436,7 @@ class WeakKeyDictionary(collections.MutableMapping):
                 return o, value
 
     def pop(self, key, *args):
+        self._dirty_len = True
         return self.data.pop(ref(key), *args)
 
     def setdefault(self, key, default=None):
diff --git a/Misc/NEWS b/Misc/NEWS
index 3aa92e8..acbe040 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -166,6 +166,9 @@ Core and Builtins
 Library
 -------
 
+- Issue #21173: Fix len() on a WeakKeyDictionary when .clear() was called
+  with an iterator alive.
+
 - Issue #11866: Eliminated race condition in the computation of names
   for new threads.