[pypy-commit] pypy py3.3: Fixed pickling of iterators produced by map(func, iter, ...).

[kvas]

Author: Vasily Kuznetsov <kvas.it at gmail.com>
Branch: py3.3
Changeset: r72516:b30c4330c8b3
Date: 2014-07-26 15:57 +0200
http://bitbucket.org/pypy/pypy/changeset/b30c4330c8b3/

Log:	Fixed pickling of iterators produced by map(func, iter, ...).

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
@@ -667,6 +667,12 @@
         # the loop is out of the way of the JIT
         return [self.space.next(w_elem) for w_elem in self.iterators_w]
 
+    def descr_reduce(self, space):
+        w_map = space.getattr(space.getbuiltinmodule('builtins'),
+                space.wrap('map'))
+        args = [self.w_fun] + self.iterators_w
+        return space.newtuple([w_map, space.newtuple(args)])
+
 
 def W_Map___new__(space, w_subtype, w_fun, args_w):
     if len(args_w) == 0:
@@ -681,6 +687,7 @@
         __new__  = interp2app(W_Map___new__),
         __iter__ = interp2app(W_Map.iter_w),
         __next__ = interp2app(W_Map.next_w),
+        __reduce__ = interp2app(W_Map.descr_reduce),
         __doc__ = """\ 
 Make an iterator that computes the function using arguments from
 each of the iterables.  Stops when the shortest iterable is exhausted.""")
diff --git a/pypy/module/__builtin__/test/test_map_pickle.py b/pypy/module/__builtin__/test/test_map_pickle.py
new file mode 100644
--- /dev/null
+++ b/pypy/module/__builtin__/test/test_map_pickle.py
@@ -0,0 +1,28 @@
+class AppTestMapPickle:
+
+    def test_map_pickle(self):
+        """Pickle a map with one sequence."""
+        import pickle
+
+        def pickle_unpickle(obj):
+            d = pickle.dumps(obj)
+            return pickle.loads(d)
+
+        m1 = map(ord, "Is this the real life?")
+        m1_ = pickle_unpickle(m1)
+
+        assert list(m1) == list(m1_)
+
+    def test_map2_pickle(self):
+        """Pickle a map with multiple sequences."""
+        import pickle
+
+        def pickle_unpickle(obj):
+            d = pickle.dumps(obj)
+            return pickle.loads(d)
+
+        m1 = map(max, "abc", "def")
+        m1_ = pickle_unpickle(m1)
+
+        assert list(m1) == list(m1_)
+