Issue #7105: Make WeakKeyDictionary and WeakValueDictionary robust against

the destruction of weakref'ed objects while iterating.

1 files changed, 85 insertions, 49 deletions

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 @@ from _weakref import (
      ProxyType,
      ReferenceType)
 
-from _weakrefset import WeakSet
+from _weakrefset import WeakSet, _IterationGuard
 
 import collections  # Import after _weakref to avoid circular import.
 
@@ -46,11 +46,25 @@ class WeakValueDictionary(collections.MutableMapping):
         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 @@ class WeakValueDictionary(collections.MutableMapping):
             return o
 
     def __delitem__(self, key):
+        if self._pending_removals:
+            self._commit_removals()
         del self.data[key]
 
     def __len__(self):
@@ -75,6 +91,8 @@ class WeakValueDictionary(collections.MutableMapping):
         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 @@ class WeakValueDictionary(collections.MutableMapping):
                 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 @@ class WeakValueDictionary(collections.MutableMapping):
         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 @@ class WeakValueDictionary(collections.MutableMapping):
                 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 @@ class WeakValueDictionary(collections.MutableMapping):
         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 @@ class WeakValueDictionary(collections.MutableMapping):
         keep the values around longer than needed.
 
         """
-        return self.data.values()
+        return list(self.data.values())
 
 
 class KeyedRef(ref):
@@ -235,9 +259,29 @@ class WeakKeyDictionary(collections.MutableMapping):
         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 @@ class WeakKeyDictionary(collections.MutableMapping):
         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 @@ class WeakKeyDictionary(collections.MutableMapping):
         keep the keys around longer than needed.
 
         """
-        return self.data.keys()
+        return list(self.data)
 
     def popitem(self):
         while 1: