[pypy-commit] pypy py3.3: Pickling of iterators produced by zip(seq, ...).

[kvas]

Author: Vasily Kuznetsov <kvas.it at gmail.com>
Branch: py3.3
Changeset: r72535:1f1711d14798
Date: 2014-07-26 16:38 +0200
http://bitbucket.org/pypy/pypy/changeset/1f1711d14798/

Log:	Pickling of iterators produced by zip(seq, ...).

diff --git a/pypy/module/__builtin__/functional.py b/pypy/module/__builtin__/functional.py
--- a/pypy/module/__builtin__/functional.py
+++ b/pypy/module/__builtin__/functional.py
@@ -753,6 +753,12 @@
             raise OperationError(self.space.w_StopIteration, self.space.w_None)
         return W_Map.next_w(self)
 
+    def descr_reduce(self, space):
+        w_zip = space.getattr(space.getbuiltinmodule('builtins'),
+                space.wrap('zip'))
+        return space.newtuple([w_zip, space.newtuple(self.iterators_w)])
+
+
 def W_Zip___new__(space, w_subtype, args_w):
     r = space.allocate_instance(W_Zip, w_subtype)
     r.__init__(space, None, args_w)
@@ -763,6 +769,7 @@
         __new__  = interp2app(W_Zip___new__),
         __iter__ = interp2app(W_Zip.iter_w),
         __next__ = interp2app(W_Zip.next_w),
+        __reduce__ = interp2app(W_Zip.descr_reduce),
         __doc__  = """\
 Return a zip object whose .__next__() method returns a tuple where
 the i-th element comes from the i-th iterable argument.  The .__next__()
diff --git a/pypy/module/__builtin__/test/test_zip_pickle.py b/pypy/module/__builtin__/test/test_zip_pickle.py
new file mode 100644
--- /dev/null
+++ b/pypy/module/__builtin__/test/test_zip_pickle.py
@@ -0,0 +1,14 @@
+class AppTestZipPickle:
+
+    def test_zip_pickle(self):
+        import pickle
+
+        def pickle_unpickle(obj):
+            d = pickle.dumps(obj)
+            return pickle.loads(d)
+
+        z1 = zip([1, 2, 3], [4, 5, 6])
+        z1_ = pickle_unpickle(z1)
+        l1, l1_ = list(z1), list(z1_)
+
+        assert l1 == l1_