diff -r 8c3e7d519b8b Lib/test/test_weakref.py --- a/Lib/test/test_weakref.py Sun Oct 11 16:50:57 2009 +0200 +++ b/Lib/test/test_weakref.py Sun Oct 11 19:32:23 2009 +0200 @@ -924,6 +924,30 @@ class MappingTestCase(TestBase): self.assertFalse(values, "itervalues() did not touch all values") + def check_weak_destroy_while_iterating(self, dict, objects, iter_name): + it = iter(getattr(dict, iter_name)()) + next(it) # Trigger internal iteration + # Destroy an object + del objects[-1] + gc.collect() # just in case + # We have removed either the first consumed object, or another one + self.assertIn(len(list(it)), [len(objects), len(objects) - 1]) + + def test_weak_keys_destroy_while_iterating(self): + dict, objects = self.make_weak_keyed_dict() + self.check_weak_destroy_while_iterating(dict, objects, 'keys') + self.check_weak_destroy_while_iterating(dict, objects, 'items') + self.check_weak_destroy_while_iterating(dict, objects, 'values') + self.check_weak_destroy_while_iterating(dict, objects, 'keyrefs') + + def test_weak_values_destroy_while_iterating(self): + dict, objects = self.make_weak_valued_dict() + self.check_weak_destroy_while_iterating(dict, objects, 'keys') + self.check_weak_destroy_while_iterating(dict, objects, 'items') + self.check_weak_destroy_while_iterating(dict, objects, 'values') + self.check_weak_destroy_while_iterating(dict, objects, 'itervaluerefs') + self.check_weak_destroy_while_iterating(dict, objects, 'valuerefs') + def test_make_weak_keyed_dict_from_dict(self): o = Object(3) dict = weakref.WeakKeyDictionary({o:364}) diff -r 8c3e7d519b8b Lib/weakref.py --- a/Lib/weakref.py Sun Oct 11 16:50:57 2009 +0200 +++ b/Lib/weakref.py Sun Oct 11 19:32:23 2009 +0200 @@ -30,6 +30,31 @@ __all__ = ["ref", "proxy", "getweakrefco "WeakSet"] +class _IterationGuard: + # This context manager registers itself in the current iterators of the + # weak dictionary, such as to delay all removals until the context manager + # exits. + # This technique should be relatively thread-safe (since sets are). + + def __init__(self, weakdict): + # Don't create cycles + self.weakdict = ref(weakdict) + + def __enter__(self): + wd = self.weakdict() + if wd is not None: + wd._iterating.add(self) + return self + + def __exit__(self, e, t, b): + wd = self.weakdict() + if wd is not None: + s = wd._iterating + s.remove(self) + if not s: + wd._commit_removals() + + class WeakValueDictionary(collections.MutableMapping): """Mapping class that references values weakly. @@ -46,11 +71,21 @@ class WeakValueDictionary(collections.Mu def remove(wr, selfref=ref(self)): self = selfref() if self is not None: - del self.data[wr.key] + if self._iterating: + self._pending_removals.append(wr.key) + else: + del self.data[wr.key] self._remove = remove + self._pending_removals = [] + self._iterating = set() self.data = d = {} self.update(*args, **kw) + def _commit_removals(self): + l = self._pending_removals + while l: + del self.data[l.pop()] + def __getitem__(self, key): o = self.data[key]() if o is None: @@ -110,24 +145,26 @@ class WeakValueDictionary(collections.Mu return o def items(self): - L = [] - for key, wr in self.data.items(): - o = wr() - if o is not None: - L.append((key, o)) - return L - - def items(self): - for wr in self.data.values(): - value = wr() - if value is not None: - yield wr.key, value + with _IterationGuard(self): + for k, wr in self.data.items(): + v = wr() + if v is not None: + yield k, v def keys(self): - return iter(self.data.keys()) + with _IterationGuard(self): + for k, wr in self.data.items(): + if wr() is not None: + yield k - def __iter__(self): - return iter(self.data.keys()) + __iter__ = keys + + def values(self): + with _IterationGuard(self): + for wr in self.data.values(): + v = wr() + if v is not None: + yield v def itervaluerefs(self): """Return an iterator that yields the weak references to the values. @@ -139,13 +176,16 @@ class WeakValueDictionary(collections.Mu keep the values around longer than needed. """ - return self.data.values() + with _IterationGuard(self): + for wr in self.data.values(): + yield wr def values(self): - for wr in self.data.values(): - obj = wr() - if obj is not None: - yield obj + with _IterationGuard(self): + for wr in self.data.values(): + obj = wr() + if obj is not None: + yield obj def popitem(self): while 1: @@ -195,7 +235,7 @@ class WeakValueDictionary(collections.Mu keep the values around longer than needed. """ - return self.data.values() + return list(self.data.values()) class KeyedRef(ref): @@ -235,9 +275,20 @@ class WeakKeyDictionary(collections.Muta def remove(k, selfref=ref(self)): self = selfref() if self is not None: - del self.data[k] + if self._iterating: + self._pending_removals.append(k) + else: + del self.data[k] self._remove = remove - if dict is not None: self.update(dict) + self._pending_removals = [] + self._iterating = set() + if dict is not None: + self.update(dict) + + def _commit_removals(self): + l = self._pending_removals + while l: + del self.data[l.pop()] def __delitem__(self, key): del self.data[ref(key)] @@ -284,34 +335,26 @@ class WeakKeyDictionary(collections.Muta return wr in self.data def items(self): - for wr, value in self.data.items(): - key = wr() - if key is not None: - yield key, value - - def keyrefs(self): - """Return an iterator that yields the weak references to the keys. - - The references are not guaranteed to be 'live' at the time - they are used, so the result of calling the references needs - to be checked before being used. This can be used to avoid - creating references that will cause the garbage collector to - keep the keys around longer than needed. - - """ - return self.data.keys() + with _IterationGuard(self): + for wr, value in self.data.items(): + key = wr() + if key is not None: + yield key, value def keys(self): - for wr in self.data.keys(): - obj = wr() - if obj is not None: - yield obj + with _IterationGuard(self): + for wr in self.data: + obj = wr() + if obj is not None: + yield obj - def __iter__(self): - return iter(self.keys()) + __iter__ = keys def values(self): - return iter(self.data.values()) + with _IterationGuard(self): + for wr, value in self.data.items(): + if wr() is not None: + yield value def keyrefs(self): """Return a list of weak references to the keys. @@ -323,7 +366,7 @@ class WeakKeyDictionary(collections.Muta keep the keys around longer than needed. """ - return self.data.keys() + return list(self.data) def popitem(self): while 1: