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 @@ -4,6 +4,8 @@ import UserList import weakref import operator +import contextlib +import copy from test import test_support @@ -906,7 +908,7 @@ def check_len_cycles(self, dict_type, cons): N = 20 items = [RefCycle() for i in range(N)] - dct = dict_type(cons(o) for o in items) + dct = dict_type(cons(i, o) for i, o in enumerate(items)) # Keep an iterator alive it = dct.iteritems() try: @@ -916,6 +918,7 @@ del items gc.collect() n1 = len(dct) + list(it) del it gc.collect() n2 = len(dct) @@ -924,10 +927,10 @@ self.assertEqual(n2, 0) def test_weak_keyed_len_cycles(self): - self.check_len_cycles(weakref.WeakKeyDictionary, lambda k: (k, 1)) + self.check_len_cycles(weakref.WeakKeyDictionary, lambda n, k: (k, n)) def test_weak_valued_len_cycles(self): - self.check_len_cycles(weakref.WeakValueDictionary, lambda k: (1, k)) + self.check_len_cycles(weakref.WeakValueDictionary, lambda n, k: (n, k)) def check_len_race(self, dict_type, cons): # Extended sanity checks for len() in the face of cyclic collection @@ -1093,6 +1096,86 @@ self.assertTrue(len(values) == 0, "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, 'iterkeys') + self.check_weak_destroy_while_iterating(dict, objects, 'iteritems') + self.check_weak_destroy_while_iterating(dict, objects, 'itervalues') + self.check_weak_destroy_while_iterating(dict, objects, 'iterkeyrefs') + dict, objects = self.make_weak_keyed_dict() + @contextlib.contextmanager + def testcontext(): + try: + it = iter(dict.iteritems()) + 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, 'iterkeys') + self.check_weak_destroy_while_iterating(dict, objects, 'iteritems') + self.check_weak_destroy_while_iterating(dict, objects, 'itervalues') + self.check_weak_destroy_while_iterating(dict, objects, 'itervaluerefs') + dict, objects = self.make_weak_valued_dict() + @contextlib.contextmanager + def testcontext(): + try: + it = iter(dict.iteritems()) + 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}) diff --git a/Lib/test/test_weakset.py b/Lib/test/test_weakset.py --- a/Lib/test/test_weakset.py +++ b/Lib/test/test_weakset.py @@ -11,6 +11,7 @@ import collections import gc import contextlib +from UserString import UserString as ustr class Foo: @@ -448,6 +449,54 @@ self.assertGreaterEqual(n2, 0) self.assertLessEqual(n2, n1) + 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): test_support.run_unittest(TestWeakSet) diff --git a/Lib/weakref.py b/Lib/weakref.py --- a/Lib/weakref.py +++ b/Lib/weakref.py @@ -120,16 +120,18 @@ return L def iteritems(self): - for wr in self.data.itervalues(): + # To avoid problems with GC or unraveling of links during iteration, + # we simply avoid iterating over the unrelying data structures. + for wr in self.data.values(): #don't use 'itervalues' value = wr() if value is not None: yield wr.key, value def iterkeys(self): - return self.data.iterkeys() + for key in self.data.keys(): + yield key - def __iter__(self): - return self.data.iterkeys() + __iter__ = iterkeys def itervaluerefs(self): """Return an iterator that yields the weak references to the values. @@ -141,10 +143,12 @@ keep the values around longer than needed. """ - return self.data.itervalues() + # see 'iteritems' above as to why we don't use itervalues + for value in self.data.values(): + yield value def itervalues(self): - for wr in self.data.itervalues(): + for wr in self.data.values(): # do not 'iter', see above. obj = wr() if obj is not None: yield obj @@ -306,7 +310,7 @@ return L def iteritems(self): - for wr, value in self.data.iteritems(): + for wr, value in self.data.items(): # do not use 'iteritems', see above key = wr() if key is not None: yield key, value @@ -321,19 +325,22 @@ keep the keys around longer than needed. """ - return self.data.iterkeys() + for key in self.data.keys(): # do not use 'iterkeys', see above. + yield key def iterkeys(self): - for wr in self.data.iterkeys(): + for wr in self.data.keys(): # do not use 'iterkeys', see above. obj = wr() if obj is not None: yield obj def __iter__(self): - return self.iterkeys() + for key in self.keys(): + yield key def itervalues(self): - return self.data.itervalues() + for value in self.data.values(): # do not use itervalues. + yield value def keyrefs(self): """Return a list of weak references to the keys.