[pypy-commit] pypy py3.3: pickle support for itertools.takewhile and dropwhile.

[amauryfa]

Author: Amaury Forgeot d'Arc <amauryfa at gmail.com>
Branch: py3.3
Changeset: r82011:6e3ddc38f7c0
Date: 2016-01-28 21:21 +0100
http://bitbucket.org/pypy/pypy/changeset/6e3ddc38f7c0/

Log:	pickle support for itertools.takewhile and dropwhile.

diff --git a/pypy/module/itertools/interp_itertools.py b/pypy/module/itertools/interp_itertools.py
--- a/pypy/module/itertools/interp_itertools.py
+++ b/pypy/module/itertools/interp_itertools.py
@@ -148,7 +148,7 @@
     def __init__(self, space, w_predicate, w_iterable):
         self.space = space
         self.w_predicate = w_predicate
-        self.iterable = space.iter(w_iterable)
+        self.w_iterable = space.iter(w_iterable)
         self.stopped = False
 
     def iter_w(self):
@@ -158,7 +158,7 @@
         if self.stopped:
             raise OperationError(self.space.w_StopIteration, self.space.w_None)
 
-        w_obj = self.space.next(self.iterable)  # may raise a w_StopIteration
+        w_obj = self.space.next(self.w_iterable)  # may raise a w_StopIteration
         w_bool = self.space.call_function(self.w_predicate, w_obj)
         if not self.space.is_true(w_bool):
             self.stopped = True
@@ -166,6 +166,16 @@
 
         return w_obj
 
+    def descr_reduce(self, space):
+        return space.newtuple([
+            space.type(self),
+            space.newtuple([self.w_predicate, self.w_iterable]),
+            space.wrap(self.stopped)
+        ])
+
+    def descr_setstate(self, space, w_state):
+        self.stopped = space.bool_w(w_state)
+
 def W_TakeWhile___new__(space, w_subtype, w_predicate, w_iterable):
     r = space.allocate_instance(W_TakeWhile, w_subtype)
     r.__init__(space, w_predicate, w_iterable)
@@ -177,6 +187,8 @@
         __new__  = interp2app(W_TakeWhile___new__),
         __iter__ = interp2app(W_TakeWhile.iter_w),
         __next__ = interp2app(W_TakeWhile.next_w),
+        __reduce__ = interp2app(W_TakeWhile.descr_reduce),
+        __setstate__ = interp2app(W_TakeWhile.descr_setstate),
         __doc__  = """Make an iterator that returns elements from the iterable as
     long as the predicate is true.
 
@@ -195,7 +207,7 @@
     def __init__(self, space, w_predicate, w_iterable):
         self.space = space
         self.w_predicate = w_predicate
-        self.iterable = space.iter(w_iterable)
+        self.w_iterable = space.iter(w_iterable)
         self.started = False
 
     def iter_w(self):
@@ -203,10 +215,10 @@
 
     def next_w(self):
         if self.started:
-            w_obj = self.space.next(self.iterable)  # may raise w_StopIteration
+            w_obj = self.space.next(self.w_iterable)  # may raise w_StopIter
         else:
             while True:
-                w_obj = self.space.next(self.iterable)  # may raise w_StopIter
+                w_obj = self.space.next(self.w_iterable)  # may raise w_StopIter
                 w_bool = self.space.call_function(self.w_predicate, w_obj)
                 if not self.space.is_true(w_bool):
                     self.started = True
@@ -214,6 +226,16 @@
 
         return w_obj
 
+    def descr_reduce(self, space):
+        return space.newtuple([
+            space.type(self),
+            space.newtuple([self.w_predicate, self.w_iterable]),
+            space.wrap(self.started)
+        ])
+
+    def descr_setstate(self, space, w_state):
+        self.started = space.bool_w(w_state)
+
 def W_DropWhile___new__(space, w_subtype, w_predicate, w_iterable):
     r = space.allocate_instance(W_DropWhile, w_subtype)
     r.__init__(space, w_predicate, w_iterable)
@@ -225,6 +247,8 @@
         __new__  = interp2app(W_DropWhile___new__),
         __iter__ = interp2app(W_DropWhile.iter_w),
         __next__ = interp2app(W_DropWhile.next_w),
+        __reduce__ = interp2app(W_DropWhile.descr_reduce),
+        __setstate__ = interp2app(W_DropWhile.descr_setstate),
         __doc__  = """Make an iterator that drops elements from the iterable as long
     as the predicate is true; afterwards, returns every
     element. Note, the iterator does not produce any output until the
diff --git a/pypy/module/itertools/test/test_itertools.py b/pypy/module/itertools/test/test_itertools.py
--- a/pypy/module/itertools/test/test_itertools.py
+++ b/pypy/module/itertools/test/test_itertools.py
@@ -836,12 +836,6 @@
         "usemodules": ['itertools', 'struct', 'binascii'],
     }
 
-    def setup_class(cls):
-        if cls.space.is_true(cls.space.appexec([], """():
-            import sys; return sys.version_info < (2, 7)
-            """)):
-            py.test.skip("Requires Python 2.7")
-
     def test_compress(self):
         import itertools
         it = itertools.compress(['a', 'b', 'c'], [0, 1, 0])
@@ -1001,6 +995,16 @@
         assert list(itertools.islice(
             pickle.loads(pickle.dumps(c)), 10)) == list('bcabcabcab')
 
+    def test_takewhile_pickle(self):
+        data = [1, 2, 3, 0, 4, 5, 6]
+        import itertools, pickle
+        t = itertools.takewhile(bool, data)
+        next(t)
+        assert list(pickle.loads(pickle.dumps(t))) == [2, 3]
+        t = itertools.dropwhile(bool, data)
+        next(t)
+        assert list(pickle.loads(pickle.dumps(t))) == [4, 5, 6]
+
 
 class AppTestItertools32:
     spaceconfig = dict(usemodules=['itertools'])