[pypy-svn] pypy default: Make itertools.count() objects picklable and copy'able.

arigo commits-noreply at bitbucket.org
Sat Jan 29 15:36:46 CET 2011


Author: Armin Rigo <arigo at tunes.org>
Branch: 
Changeset: r41444:099b99fd4a0c
Date: 2011-01-29 15:35 +0100
http://bitbucket.org/pypy/pypy/changeset/099b99fd4a0c/

Log:	Make itertools.count() objects picklable and copy'able.

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
@@ -959,3 +959,18 @@
         assert type(A('', 0)) is A
         class A(itertools.combinations_with_replacement): pass
         assert type(A('', 0)) is A
+
+    def test_copy_pickle(self):
+        import itertools, copy, pickle, sys
+        for value in [42, -sys.maxint*99]:
+            for step in [1, sys.maxint*42, 5.5]:
+                expected = [value, value+step, value+2*step]
+                c = itertools.count(value, step)
+                assert list(itertools.islice(c, 3)) == expected
+                c = itertools.count(value, step)
+                c1 = copy.copy(c)
+                assert list(itertools.islice(c1, 3)) == expected
+                c2 = copy.deepcopy(c)
+                assert list(itertools.islice(c2, 3)) == expected
+                c3 = pickle.loads(pickle.dumps(c))
+                assert list(itertools.islice(c3, 3)) == expected

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
@@ -20,17 +20,29 @@
         self.w_c = self.space.add(w_c, self.w_step)
         return w_c
 
+    def single_argument(self):
+        space = self.space
+        return (space.isinstance_w(self.w_step, space.w_int) and
+                space.eq_w(self.w_step, space.wrap(1)))
+
     def repr_w(self):
         space = self.space
         c = space.str_w(space.repr(self.w_c))
-        if (space.isinstance_w(self.w_step, space.w_int) and
-            space.eq_w(self.w_step, space.wrap(1))):
+        if self.single_argument():
             s = 'count(%s)' % (c,)
         else:
             step = space.str_w(space.repr(self.w_step))
             s = 'count(%s, %s)' % (c, step)
         return self.space.wrap(s)
 
+    def reduce_w(self):
+        space = self.space
+        if self.single_argument():
+            args_w = [self.w_c]
+        else:
+            args_w = [self.w_c, self.w_step]
+        return space.newtuple([space.gettypefor(W_Count),
+                               space.newtuple(args_w)])
 
 def check_number(space, w_obj):
     if (space.lookup(w_obj, '__add__') is None or
@@ -49,9 +61,11 @@
 
 W_Count.typedef = TypeDef(
         'count',
+        __module__ = 'itertools',
         __new__ = interp2app(W_Count___new__),
         __iter__ = interp2app(W_Count.iter_w, unwrap_spec=['self']),
         next = interp2app(W_Count.next_w, unwrap_spec=['self']),
+        __reduce__ = interp2app(W_Count.reduce_w, unwrap_spec=['self']),
         __repr__ = interp2app(W_Count.repr_w, unwrap_spec=['self']),
         __doc__ = """Make an iterator that returns consecutive integers starting
     with n.  If not specified n defaults to zero. Does not currently
@@ -108,6 +122,7 @@
 
 W_Repeat.typedef = TypeDef(
         'repeat',
+        __module__ = 'itertools',
         __new__  = interp2app(W_Repeat___new__, unwrap_spec=[ObjSpace, W_Root, W_Root, W_Root]),
         __iter__ = interp2app(W_Repeat.iter_w, unwrap_spec=['self']),
         next     = interp2app(W_Repeat.next_w, unwrap_spec=['self']),
@@ -160,6 +175,7 @@
 
 W_TakeWhile.typedef = TypeDef(
         'takewhile',
+        __module__ = 'itertools',
         __new__  = interp2app(W_TakeWhile___new__, unwrap_spec=[ObjSpace, W_Root, W_Root, W_Root]),
         __iter__ = interp2app(W_TakeWhile.iter_w, unwrap_spec=['self']),
         next     = interp2app(W_TakeWhile.next_w, unwrap_spec=['self']),
@@ -208,6 +224,7 @@
 
 W_DropWhile.typedef = TypeDef(
         'dropwhile',
+        __module__ = 'itertools',
         __new__  = interp2app(W_DropWhile___new__, unwrap_spec=[ObjSpace, W_Root, W_Root, W_Root]),
         __iter__ = interp2app(W_DropWhile.iter_w, unwrap_spec=['self']),
         next     = interp2app(W_DropWhile.next_w, unwrap_spec=['self']),
@@ -264,6 +281,7 @@
 
 W_IFilter.typedef = TypeDef(
         'ifilter',
+        __module__ = 'itertools',
         __new__  = interp2app(W_IFilter___new__, unwrap_spec=[ObjSpace, W_Root, W_Root, W_Root]),
         __iter__ = interp2app(W_IFilter.iter_w, unwrap_spec=['self']),
         next     = interp2app(W_IFilter.next_w, unwrap_spec=['self']),
@@ -291,6 +309,7 @@
 
 W_IFilterFalse.typedef = TypeDef(
         'ifilterfalse',
+        __module__ = 'itertools',
         __new__  = interp2app(W_IFilterFalse___new__, unwrap_spec=[ObjSpace, W_Root, W_Root, W_Root]),
         __iter__ = interp2app(W_IFilterFalse.iter_w, unwrap_spec=['self']),
         next     = interp2app(W_IFilterFalse.next_w, unwrap_spec=['self']),
@@ -381,6 +400,7 @@
 
 W_ISlice.typedef = TypeDef(
         'islice',
+        __module__ = 'itertools',
         __new__  = interp2app(W_ISlice___new__, unwrap_spec=[ObjSpace, W_Root, W_Root, W_Root, 'args_w']),
         __iter__ = interp2app(W_ISlice.iter_w, unwrap_spec=['self']),
         next     = interp2app(W_ISlice.next_w, unwrap_spec=['self']),
@@ -448,6 +468,7 @@
 
 W_Chain.typedef = TypeDef(
         'chain',
+        __module__ = 'itertools',
         __new__  = interp2app(W_Chain___new__, unwrap_spec=[ObjSpace, W_Root, 'args_w']),
         __iter__ = interp2app(W_Chain.iter_w, unwrap_spec=['self']),
         next     = interp2app(W_Chain.next_w, unwrap_spec=['self']),
@@ -528,6 +549,7 @@
 
 W_IMap.typedef = TypeDef(
         'imap',
+        __module__ = 'itertools',
         __new__  = interp2app(W_IMap___new__, unwrap_spec=[ObjSpace, W_Root, W_Root, 'args_w']),
         __iter__ = interp2app(W_IMap.iter_w, unwrap_spec=['self']),
         next     = interp2app(W_IMap.next_w, unwrap_spec=['self']),
@@ -572,6 +594,7 @@
 
 W_IZip.typedef = TypeDef(
         'izip',
+        __module__ = 'itertools',
         __new__  = interp2app(W_IZip___new__, unwrap_spec=[ObjSpace, W_Root, 'args_w']),
         __iter__ = interp2app(W_IZip.iter_w, unwrap_spec=['self']),
         next     = interp2app(W_IZip.next_w, unwrap_spec=['self']),
@@ -641,6 +664,7 @@
 
 W_IZipLongest.typedef = TypeDef(
         'izip_longest',
+        __module__ = 'itertools',
         __new__  = interp2app(W_IZipLongest___new__),
         __iter__ = interp2app(W_IZipLongest.iter_w, unwrap_spec=['self']),
         next     = interp2app(W_IZipLongest.next_w, unwrap_spec=['self']),
@@ -700,6 +724,7 @@
 
 W_Cycle.typedef = TypeDef(
         'cycle',
+        __module__ = 'itertools',
         __new__  = interp2app(W_Cycle___new__, unwrap_spec=[ObjSpace, W_Root, W_Root]),
         __iter__ = interp2app(W_Cycle.iter_w, unwrap_spec=['self']),
         next     = interp2app(W_Cycle.next_w, unwrap_spec=['self']),
@@ -740,6 +765,7 @@
 
 W_StarMap.typedef = TypeDef(
         'starmap',
+        __module__ = 'itertools',
         __new__  = interp2app(W_StarMap___new__, unwrap_spec=[ObjSpace, W_Root, W_Root, W_Root]),
 
         __iter__ = interp2app(W_StarMap.iter_w, unwrap_spec=['self']),
@@ -845,6 +871,7 @@
 
 W_TeeIterable.typedef = TypeDef(
         '_tee',
+        __module__ = 'itertools',
         __new__ = interp2app(W_TeeIterable___new__, unwrap_spec=[ObjSpace,
                                                                  W_Root,
                                                                  W_Root]),
@@ -949,6 +976,7 @@
 
 W_GroupBy.typedef = TypeDef(
         'groupby',
+        __module__ = 'itertools',
         __new__  = interp2app(W_GroupBy___new__, unwrap_spec=[ObjSpace, W_Root, W_Root, W_Root]),
         __iter__ = interp2app(W_GroupBy.iter_w, unwrap_spec=['self']),
         next     = interp2app(W_GroupBy.next_w, unwrap_spec=['self']),
@@ -995,6 +1023,7 @@
 
 W_GroupByIterator.typedef = TypeDef(
         '_groupby',
+        __module__ = 'itertools',
         __iter__ = interp2app(W_GroupByIterator.iter_w, unwrap_spec=['self']),
         next     = interp2app(W_GroupByIterator.next_w, unwrap_spec=['self']))
 W_GroupByIterator.typedef.acceptable_as_base_class = False
@@ -1026,6 +1055,7 @@
 
 W_Compress.typedef = TypeDef(
     'compress',
+    __module__ = 'itertools',
     __new__ = interp2app(W_Compress__new__,
                          unwrap_spec=[ObjSpace, W_Root, W_Root, W_Root]),
     __iter__ = interp2app(W_Compress.iter_w, unwrap_spec=['self']),
@@ -1119,6 +1149,7 @@
 
 W_Product.typedef = TypeDef(
     'product',
+    __module__ = 'itertools',
     __new__ = interp2app(W_Product__new__),
     __iter__ = interp2app(W_Product.iter_w),
     next = interp2app(W_Product.next_w),
@@ -1224,6 +1255,7 @@
     return space.wrap(res)
 
 W_Combinations.typedef = TypeDef("combinations",
+    __module__ = 'itertools',
     __new__ = interp2app(W_Combinations__new__),
     __iter__ = interp2app(W_Combinations.descr__iter__),
     next = interp2app(W_Combinations.descr_next),
@@ -1258,6 +1290,7 @@
     return space.wrap(res)
 
 W_CombinationsWithReplacement.typedef = TypeDef("combinations_with_replacement",
+    __module__ = 'itertools',
     __new__ = interp2app(W_CombinationsWithReplacement__new__),
     __iter__ = interp2app(W_CombinationsWithReplacement.descr__iter__),
     next = interp2app(W_CombinationsWithReplacement.descr_next),
@@ -1326,6 +1359,7 @@
     return space.wrap(res)
 
 W_Permutations.typedef = TypeDef("permutations",
+    __module__ = 'itertools',
     __new__ = interp2app(W_Permutations__new__),
     __iter__ = interp2app(W_Permutations.descr__iter__),
     next = interp2app(W_Permutations.descr_next),


More information about the Pypy-commit mailing list