[pypy-commit] pypy default: Same issue and same fix for sets.

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: 
Changeset: r54155:f16be5ae31e6
Date: 2012-04-02 18:22 +0200
http://bitbucket.org/pypy/pypy/changeset/f16be5ae31e6/

Log:	Same issue and same fix for sets.

diff --git a/pypy/objspace/std/setobject.py b/pypy/objspace/std/setobject.py
--- a/pypy/objspace/std/setobject.py
+++ b/pypy/objspace/std/setobject.py
@@ -359,7 +359,7 @@
         w_set.sstorage = w_other.get_storage_copy()
 
     def iter(self, w_set):
-        return EmptyIteratorImplementation(self.space, w_set)
+        return EmptyIteratorImplementation(self.space, self, w_set)
 
     def popitem(self, w_set):
         raise OperationError(self.space.w_KeyError,
@@ -784,8 +784,9 @@
             d_obj[w_item] = None
 
 class IteratorImplementation(object):
-    def __init__(self, space, implementation):
+    def __init__(self, space, strategy, implementation):
         self.space = space
+        self.strategy = strategy
         self.setimplementation = implementation
         self.len = implementation.length()
         self.pos = 0
@@ -801,7 +802,17 @@
         if self.pos < self.len:
             result = self.next_entry()
             self.pos += 1
-            return result
+            if self.strategy is self.setimplementation.strategy:
+                return result      # common case
+            else:
+                # waaa, obscure case: the strategy changed, but not the
+                # length of the set.  The 'result' might be out-of-date.
+                # We try to explicitly look it up in the set.
+                if not self.setimplementation.has_key(result):
+                    self.len = -1   # Make this error state sticky
+                    raise OperationError(self.space.w_RuntimeError,
+                        self.space.wrap("dictionary changed during iteration"))
+                return result
         # no more entries
         self.setimplementation = None
         return None
@@ -823,7 +834,7 @@
 
 class StringIteratorImplementation(IteratorImplementation):
     def __init__(self, space, strategy, w_set):
-        IteratorImplementation.__init__(self, space, w_set)
+        IteratorImplementation.__init__(self, space, strategy, w_set)
         d = strategy.unerase(w_set.sstorage)
         self.iterator = d.iterkeys()
 
@@ -835,9 +846,9 @@
 
 class IntegerIteratorImplementation(IteratorImplementation):
     #XXX same implementation in dictmultiobject on dictstrategy-branch
-    def __init__(self, space, strategy, dictimplementation):
-        IteratorImplementation.__init__(self, space, dictimplementation)
-        d = strategy.unerase(dictimplementation.sstorage)
+    def __init__(self, space, strategy, w_set):
+        IteratorImplementation.__init__(self, space, strategy, w_set)
+        d = strategy.unerase(w_set.sstorage)
         self.iterator = d.iterkeys()
 
     def next_entry(self):
@@ -848,9 +859,9 @@
             return None
 
 class RDictIteratorImplementation(IteratorImplementation):
-    def __init__(self, space, strategy, dictimplementation):
-        IteratorImplementation.__init__(self, space, dictimplementation)
-        d = strategy.unerase(dictimplementation.sstorage)
+    def __init__(self, space, strategy, w_set):
+        IteratorImplementation.__init__(self, space, strategy, w_set)
+        d = strategy.unerase(w_set.sstorage)
         self.iterator = d.iterkeys()
 
     def next_entry(self):
diff --git a/pypy/objspace/std/test/test_dictmultiobject.py b/pypy/objspace/std/test/test_dictmultiobject.py
--- a/pypy/objspace/std/test/test_dictmultiobject.py
+++ b/pypy/objspace/std/test/test_dictmultiobject.py
@@ -817,7 +817,7 @@
         class Foo(object):
             def __eq__(self, other):
                 return False
-        d.get(Foo())    # this changes the strategy of 'd'
+        assert d.get(Foo()) is None    # this changes the strategy of 'd'
         lst = list(it)  # but iterating still works
         assert sorted(lst) == [(1, 2), (3, 4), (5, 6)]
 
@@ -826,8 +826,10 @@
         it = d.iteritems()
         d['foo'] = 'bar'
         del d[1]
-        # 'd' is still length 3, but its strategy changed
-        raises(RuntimeError, it.next)
+        # 'd' is still length 3, but its strategy changed.  we are
+        # getting a RuntimeError because iterating over the old storage
+        # gives us (1, 2), but 1 is not in the dict any longer.
+        raises(RuntimeError, list, it)
 
 
 class FakeString(str):
diff --git a/pypy/objspace/std/test/test_setobject.py b/pypy/objspace/std/test/test_setobject.py
--- a/pypy/objspace/std/test/test_setobject.py
+++ b/pypy/objspace/std/test/test_setobject.py
@@ -907,3 +907,30 @@
                 return [5, 3, 4][i]
         s = set([10,3,2]).intersection(Obj())
         assert list(s) == [3]
+
+    def test_iter_set_length_change(self):
+        s = set([1, 3, 5])
+        it = iter(s)
+        s.add(7)
+        # 's' is now length 4
+        raises(RuntimeError, it.next)
+
+    def test_iter_set_strategy_only_change_1(self):
+        s = set([1, 3, 5])
+        it = iter(s)
+        class Foo(object):
+            def __eq__(self, other):
+                return False
+        assert Foo() not in s      # this changes the strategy of 'd'
+        lst = list(s)  # but iterating still works
+        assert sorted(lst) == [1, 3, 5]
+
+    def test_iter_set_strategy_only_change_2(self):
+        s = set([1, 3, 5])
+        it = iter(s)
+        s.add('foo')
+        s.remove(1)
+        # 's' is still length 3, but its strategy changed.  we are
+        # getting a RuntimeError because iterating over the old storage
+        # gives us 1, but 1 is not in the set any longer.
+        raises(RuntimeError, list, it)