Commits

cbjadwani committed 2f592ec Draft

Add lib_pypy.itertools tests, fix check_number in interp_itertools

Comments (0)

Files changed (2)

lib_pypy/pypy_test/test_itertools.py

         l = [1, 2]
         m = ['a']
         raises(TypeError, itertools.product, l, m, repeat=1, foo=2)
+
+    def test_tee_copy_constructor(self):
+        a, b = itertools.tee(range(10))
+        next(a)
+        next(a)
+        c, d = itertools.tee(a)
+        assert list(a) == list(d)
+
+    def test_product_kwargs(self):
+        raises(TypeError, itertools.product, range(10), garbage=1)
+
+    def test_takewhile_stops(self):
+        tw = itertools.takewhile(lambda x: bool(x), [1, 1, 0, 1, 1])
+        next(tw)
+        next(tw)
+        raises(StopIteration, next, tw)
+        raises(StopIteration, next, tw)

pypy/module/itertools/interp_itertools.py

                                space.newtuple(args_w)])
 
 def check_number(space, w_obj):
-    if (space.lookup(w_obj, '__add__') is None or
-        space.is_true(space.isinstance(w_obj, space.w_str)) or
-        space.is_true(space.isinstance(w_obj, space.w_unicode))):
+    if (space.lookup(w_obj, '__int__') is None and
+            space.lookup(w_obj, '__float__') is None):
         raise OperationError(space.w_TypeError,
                              space.wrap("expected a number"))
 
         next = interp2app(W_Count.next_w),
         __reduce__ = interp2app(W_Count.reduce_w),
         __repr__ = interp2app(W_Count.repr_w),
-        __doc__ = """Make an iterator that returns consecutive integers starting
-    with n.  If not specified n defaults to zero. Does not currently
-    support python long integers. Often used as an argument to imap()
-    to generate consecutive data points.  Also, used with izip() to
-    add sequence numbers.
+        __doc__ = """Make an iterator that returns evenly spaced values starting
+    with n.  If not specified n defaults to zero.  Often used as an
+    argument to imap() to generate consecutive data points.  Also,
+    used with izip() to add sequence numbers.
 
-    Equivalent to :
+    Equivalent to:
 
-    def count(n=0):
-        if not isinstance(n, int):
-            raise TypeError("%s is not a regular integer" % n)
+    def count(start=0, step=1):
+        n = start
         while True:
             yield n
-            n += 1
+            n += step
     """)