Issue #7105: Make WeakKeyDictionary and WeakValueDictionary robust against

the destruction of weakref'ed objects while iterating.

Author: pitrou

Doc/library/weakref.rst

@@ -159,7 +159,7 @@ than needed.
 
 .. method:: WeakKeyDictionary.keyrefs()
 
-   Return an :term:`iterator` that yields the weak references to the keys.
+   Return an iterable of the weak references to the keys.
 
 
 .. class:: WeakValueDictionary([dict])
@@ -182,7 +182,7 @@ These method have the same issues as the and :meth:`keyrefs` method of
 
 .. method:: WeakValueDictionary.valuerefs()
 
-   Return an :term:`iterator` that yields the weak references to the values.
+   Return an iterable of the weak references to the values.
 
 
 .. class:: WeakSet([elements])

Lib/_weakrefset.py

@@ -6,22 +6,61 @@
 
 __all__ = ['WeakSet']
 
+
+class _IterationGuard:
+    # This context manager registers itself in the current iterators of the
+    # weak container, such as to delay all removals until the context manager
+    # exits.
+    # This technique should be relatively thread-safe (since sets are).
+
+    def __init__(self, weakcontainer):
+        # Don't create cycles
+        self.weakcontainer = ref(weakcontainer)
+
+    def __enter__(self):
+        w = self.weakcontainer()
+        if w is not None:
+            w._iterating.add(self)
+        return self
+
+    def __exit__(self, e, t, b):
+        w = self.weakcontainer()
+        if w is not None:
+            s = w._iterating
+            s.remove(self)
+            if not s:
+                w._commit_removals()
+
+
 class WeakSet:
     def __init__(self, data=None):
         self.data = set()
         def _remove(item, selfref=ref(self)):
             self = selfref()
             if self is not None:
-                self.data.discard(item)
+                if self._iterating:
+                    self._pending_removals.append(item)
+                else:
+                    self.data.discard(item)
         self._remove = _remove
+        # A list of keys to be removed
+        self._pending_removals = []
+        self._iterating = set()
         if data is not None:
             self.update(data)
 
+    def _commit_removals(self):
+        l = self._pending_removals
+        discard = self.data.discard
+        while l:
+            discard(l.pop())
+
     def __iter__(self):
-        for itemref in self.data:
-            item = itemref()
-            if item is not None:
-                yield item
+        with _IterationGuard(self):
+            for itemref in self.data:
+                item = itemref()
+                if item is not None:
+                    yield item
 
     def __len__(self):
         return sum(x() is not None for x in self.data)
@@ -34,15 +73,21 @@ def __reduce__(self):
                 getattr(self, '__dict__', None))
 
     def add(self, item):
+        if self._pending_removals:
+            self._commit_removals()
         self.data.add(ref(item, self._remove))
 
     def clear(self):
+        if self._pending_removals:
+            self._commit_removals()
         self.data.clear()
 
     def copy(self):
         return self.__class__(self)
 
     def pop(self):
+        if self._pending_removals:
+            self._commit_removals()
         while True:
             try:
                 itemref = self.data.pop()
@@ -53,17 +98,24 @@ def pop(self):
                 return item
 
     def remove(self, item):
+        if self._pending_removals:
+            self._commit_removals()
         self.data.remove(ref(item))
 
     def discard(self, item):
+        if self._pending_removals:
+            self._commit_removals()
         self.data.discard(ref(item))
 
     def update(self, other):
+        if self._pending_removals:
+            self._commit_removals()
         if isinstance(other, self.__class__):
             self.data.update(other.data)
         else:
             for element in other:
                 self.add(element)
+
     def __ior__(self, other):
         self.update(other)
         return self
@@ -82,11 +134,15 @@ def difference(self, other):
     __sub__ = difference
 
     def difference_update(self, other):
+        if self._pending_removals:
+            self._commit_removals()
         if self is other:
             self.data.clear()
         else:
             self.data.difference_update(ref(item) for item in other)
     def __isub__(self, other):
+        if self._pending_removals:
+            self._commit_removals()
         if self is other:
             self.data.clear()
         else:
@@ -98,8 +154,12 @@ def intersection(self, other):
     __and__ = intersection
 
     def intersection_update(self, other):
+        if self._pending_removals:
+            self._commit_removals()
         self.data.intersection_update(ref(item) for item in other)
     def __iand__(self, other):
+        if self._pending_removals:
+            self._commit_removals()
         self.data.intersection_update(ref(item) for item in other)
         return self
 
@@ -127,11 +187,15 @@ def symmetric_difference(self, other):
     __xor__ = symmetric_difference
 
     def symmetric_difference_update(self, other):
+        if self._pending_removals:
+            self._commit_removals()
         if self is other:
             self.data.clear()
         else:
             self.data.symmetric_difference_update(ref(item) for item in other)
     def __ixor__(self, other):
+        if self._pending_removals:
+            self._commit_removals()
         if self is other:
             self.data.clear()
         else:

Lib/test/test_weakref.py

@@ -4,6 +4,8 @@
 import collections
 import weakref
 import operator
+import contextlib
+import copy
 
 from test import support
 
@@ -788,6 +790,10 @@ def __init__(self, arg):
         self.arg = arg
     def __repr__(self):
         return "<Object %r>" % self.arg
+    def __eq__(self, other):
+        if isinstance(other, Object):
+            return self.arg == other.arg
+        return NotImplemented
     def __lt__(self, other):
         if isinstance(other, Object):
             return self.arg < other.arg
@@ -935,6 +941,87 @@ def check_iters(self, dict):
         self.assertFalse(values,
                      "itervalues() did not touch all values")
 
+    def check_weak_destroy_while_iterating(self, dict, objects, iter_name):
+        n = len(dict)
+        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])
+        del it
+        # The removal has been committed
+        self.assertEqual(len(dict), n - 1)
+
+    def check_weak_destroy_and_mutate_while_iterating(self, dict, testcontext):
+        # Check that we can explicitly mutate the weak dict without
+        # interfering with delayed removal.
+        # `testcontext` should create an iterator, destroy one of the
+        # weakref'ed objects and then return a new key/value pair corresponding
+        # to the destroyed object.
+        with testcontext() as (k, v):
+            self.assertFalse(k in dict)
+        with testcontext() as (k, v):
+            self.assertRaises(KeyError, dict.__delitem__, k)
+        self.assertFalse(k in dict)
+        with testcontext() as (k, v):
+            self.assertRaises(KeyError, dict.pop, k)
+        self.assertFalse(k in dict)
+        with testcontext() as (k, v):
+            dict[k] = v
+        self.assertEqual(dict[k], v)
+        ddict = copy.copy(dict)
+        with testcontext() as (k, v):
+            dict.update(ddict)
+        self.assertEqual(dict, ddict)
+        with testcontext() as (k, v):
+            dict.clear()
+        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()
+        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')
+        dict, objects = self.make_weak_keyed_dict()
+        @contextlib.contextmanager
+        def testcontext():
+            try:
+                it = iter(dict.items())
+                next(it)
+                # Schedule a key/value for removal and recreate it
+                v = objects.pop().arg
+                gc.collect()      # just in case
+                yield Object(v), v
+            finally:
+                it = None           # should commit all removals
+        self.check_weak_destroy_and_mutate_while_iterating(dict, testcontext)
+
+    def test_weak_values_destroy_while_iterating(self):
+        # Issue #7105: iterators shouldn't crash when a key is implicitly removed
+        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')
+        dict, objects = self.make_weak_valued_dict()
+        @contextlib.contextmanager
+        def testcontext():
+            try:
+                it = iter(dict.items())
+                next(it)
+                # Schedule a key/value for removal and recreate it
+                k = objects.pop().arg
+                gc.collect()      # just in case
+                yield k, Object(k)
+            finally:
+                it = None           # should commit all removals
+        self.check_weak_destroy_and_mutate_while_iterating(dict, testcontext)
+
     def test_make_weak_keyed_dict_from_dict(self):
         o = Object(3)
         dict = weakref.WeakKeyDictionary({o:364})

Lib/test/test_weakset.py

@@ -10,6 +10,8 @@
 import warnings
 import collections
 from collections import UserString as ustr
+import gc
+import contextlib
 
 
 class Foo:
@@ -307,6 +309,54 @@ def test_eq(self):
         self.assertFalse(self.s == WeakSet([Foo]))
         self.assertFalse(self.s == 1)
 
+    def test_weak_destroy_while_iterating(self):
+        # Issue #7105: iterators shouldn't crash when a key is implicitly removed
+        # Create new items to be sure no-one else holds a reference
+        items = [ustr(c) for c in ('a', 'b', 'c')]
+        s = WeakSet(items)
+        it = iter(s)
+        next(it)             # Trigger internal iteration
+        # Destroy an item
+        del items[-1]
+        gc.collect()    # just in case
+        # We have removed either the first consumed items, or another one
+        self.assertIn(len(list(it)), [len(items), len(items) - 1])
+        del it
+        # The removal has been committed
+        self.assertEqual(len(s), len(items))
+
+    def test_weak_destroy_and_mutate_while_iterating(self):
+        # Issue #7105: iterators shouldn't crash when a key is implicitly removed
+        items = [ustr(c) for c in string.ascii_letters]
+        s = WeakSet(items)
+        @contextlib.contextmanager
+        def testcontext():
+            try:
+                it = iter(s)
+                next(it)
+                # Schedule an item for removal and recreate it
+                u = ustr(str(items.pop()))
+                gc.collect()      # just in case
+                yield u
+            finally:
+                it = None           # should commit all removals
+
+        with testcontext() as u:
+            self.assertFalse(u in s)
+        with testcontext() as u:
+            self.assertRaises(KeyError, s.remove, u)
+        self.assertFalse(u in s)
+        with testcontext() as u:
+            s.add(u)
+        self.assertTrue(u in s)
+        t = s.copy()
+        with testcontext() as u:
+            s.update(t)
+        self.assertEqual(len(s), len(t))
+        with testcontext() as u:
+            s.clear()
+        self.assertEqual(len(s), 0)
+
 
 def test_main(verbose=None):
     support.run_unittest(TestWeakSet)