[pypy-commit] pypy py3.3: Add support for pickling set iterator

[amauryfa]

Author: Amaury Forgeot d'Arc <amauryfa at gmail.com>
Branch: py3.3
Changeset: r76301:91188a5ae517
Date: 2015-03-06 16:10 +0100
http://bitbucket.org/pypy/pypy/changeset/91188a5ae517/

Log:	Add support for pickling set iterator

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
@@ -1520,10 +1520,23 @@
             return w_key
         raise OperationError(space.w_StopIteration, space.w_None)
 
+    def descr_reduce(self, space):
+        # copy the iterator state
+        w_set = self.iterimplementation.setimplementation
+        w_clone = W_SetIterObject(space, w_set.iter())
+        # spool until we have the same pos
+        for x in xrange(self.iterimplementation.pos):
+            w_clone.descr_next(space)
+        w_res = space.call_function(space.w_list, w_clone)
+        w_iter = space.builtin.get('iter')
+        return space.newtuple([w_iter, space.newtuple([w_res])])
+
+
 W_SetIterObject.typedef = TypeDef("setiterator",
     __length_hint__ = gateway.interp2app(W_SetIterObject.descr_length_hint),
     __iter__ = gateway.interp2app(W_SetIterObject.descr_iter),
-    __next__ = gateway.interp2app(W_SetIterObject.descr_next)
+    __next__ = gateway.interp2app(W_SetIterObject.descr_next),
+    __reduce__ = gateway.interp2app(W_SetIterObject.descr_reduce),
     )
 setiter_typedef = W_SetIterObject.typedef
 
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
@@ -1062,3 +1062,17 @@
         s = set([1, 2, 3])
         s.intersection_update(set())
         assert strategy(s) == "EmptySetStrategy"
+
+    def test_pickle(self):
+        d = {1, 2, 3}
+        it = iter(d)
+        first = next(it)
+        reduced = it.__reduce__()
+        rebuild, args = reduced
+        assert rebuild is iter
+        new = rebuild(*args)
+        items = set(new)
+        assert len(items) == 2
+        items.add(first)
+        assert items == set(d)        
+