Issue #7105: Make WeakKeyDictionary and WeakValueDictionary robust against
the destruction of weakref'ed objects while iterating.
diff --git a/Lib/weakref.py b/Lib/weakref.py
index 5e6cc8b..66c4dc6 100644
--- a/Lib/weakref.py
+++ b/Lib/weakref.py
@@ -18,7 +18,7 @@
ProxyType,
ReferenceType)
-from _weakrefset import WeakSet
+from _weakrefset import WeakSet, _IterationGuard
import collections # Import after _weakref to avoid circular import.
@@ -46,11 +46,25 @@
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
+ # A list of keys to be removed
+ self._pending_removals = []
+ self._iterating = set()
self.data = d = {}
self.update(*args, **kw)
+ def _commit_removals(self):
+ l = self._pending_removals
+ d = self.data
+ # We shouldn't encounter any KeyError, because this method should
+ # always be called *before* mutating the dict.
+ while l:
+ del d[l.pop()]
+
def __getitem__(self, key):
o = self.data[key]()
if o is None:
@@ -59,6 +73,8 @@
return o
def __delitem__(self, key):
+ if self._pending_removals:
+ self._commit_removals()
del self.data[key]
def __len__(self):
@@ -75,6 +91,8 @@
return "<WeakValueDictionary at %s>" % id(self)
def __setitem__(self, key, value):
+ if self._pending_removals:
+ self._commit_removals()
self.data[key] = KeyedRef(value, self._remove, key)
def copy(self):
@@ -110,24 +128,19 @@
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 itervaluerefs(self):
"""Return an iterator that yields the weak references to the values.
@@ -139,15 +152,20 @@
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):
+ if self._pending_removals:
+ self._commit_removals()
while 1:
key, wr = self.data.popitem()
o = wr()
@@ -155,6 +173,8 @@
return key, o
def pop(self, key, *args):
+ if self._pending_removals:
+ self._commit_removals()
try:
o = self.data.pop(key)()
except KeyError:
@@ -170,12 +190,16 @@
try:
wr = self.data[key]
except KeyError:
+ if self._pending_removals:
+ self._commit_removals()
self.data[key] = KeyedRef(default, self._remove, key)
return default
else:
return wr()
def update(self, dict=None, **kwargs):
+ if self._pending_removals:
+ self._commit_removals()
d = self.data
if dict is not None:
if not hasattr(dict, "items"):
@@ -195,7 +219,7 @@
keep the values around longer than needed.
"""
- return self.data.values()
+ return list(self.data.values())
class KeyedRef(ref):
@@ -235,9 +259,29 @@
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)
+ # A list of dead weakrefs (keys to be removed)
+ self._pending_removals = []
+ self._iterating = set()
+ if dict is not None:
+ self.update(dict)
+
+ def _commit_removals(self):
+ # NOTE: We don't need to call this method before mutating the dict,
+ # because a dead weakref never compares equal to a live weakref,
+ # even if they happened to refer to equal objects.
+ # However, it means keys may already have been removed.
+ l = self._pending_removals
+ d = self.data
+ while l:
+ try:
+ del d[l.pop()]
+ except KeyError:
+ pass
def __delitem__(self, key):
del self.data[ref(key)]
@@ -284,34 +328,26 @@
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 +359,7 @@
keep the keys around longer than needed.
"""
- return self.data.keys()
+ return list(self.data)
def popitem(self):
while 1: