diff --git a/Lib/test/test_weakref.py b/Lib/test/test_weakref.py --- a/Lib/test/test_weakref.py +++ b/Lib/test/test_weakref.py @@ -1298,6 +1298,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() @@ -1319,6 +1349,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 @@ -1342,6 +1376,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 --- a/Lib/weakref.py +++ b/Lib/weakref.py @@ -322,6 +322,7 @@ class WeakKeyDictionary(collections.Muta # 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.Muta 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.Muta 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.Muta 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):