[pypy-svn] r56413 - in pypy/dist/pypy/module/itertools: . test

[adurdin at codespeak.net]

Author: adurdin
Date: Thu Jul 10 15:19:34 2008
New Revision: 56413

Modified:
   pypy/dist/pypy/module/itertools/__init__.py
   pypy/dist/pypy/module/itertools/interp_itertools.py
   pypy/dist/pypy/module/itertools/test/test_itertools.py
Log:
(adurdin, jlg) implemented repeat() in interp_itertools.

Modified: pypy/dist/pypy/module/itertools/__init__.py
==============================================================================
--- pypy/dist/pypy/module/itertools/__init__.py	(original)
+++ pypy/dist/pypy/module/itertools/__init__.py	Thu Jul 10 15:19:34 2008
@@ -6,6 +6,7 @@
 
     interpleveldefs = {
         'count'    : 'interp_itertools.W_Count',
+        'repeat'    : 'interp_itertools.W_Repeat',
     }
 
     appleveldefs = {

Modified: pypy/dist/pypy/module/itertools/interp_itertools.py
==============================================================================
--- pypy/dist/pypy/module/itertools/interp_itertools.py	(original)
+++ pypy/dist/pypy/module/itertools/interp_itertools.py	Thu Jul 10 15:19:34 2008
@@ -55,8 +55,54 @@
             n += 1
     """)
 
+
 class W_Repeat(Wrappable):
 
-    def __init__(self, space):
+    def __init__(self, space, w_obj, w_times):
         self.space = space
+        self.w_obj = w_obj
+        
+        if space.is_w(w_times, space.w_None):
+            self.counting = False
+            self.count = 0
+        else:
+            self.counting = True
+            self.count = self.space.int_w(w_times)
+
+    def next_w(self):
+        if self.counting:
+            if self.count <= 0:
+                raise OperationError(self.space.w_StopIteration, self.space.w_None)
+            self.count -= 1
+        return self.w_obj
+
+    def iter_w(self):
+        return self.space.wrap(self)
+
+def W_Repeat___new__(space, w_subtype, w_obj, w_times=None):
+    """
+    Create a new repeat object and call its initializer.
+    """
+    return space.wrap(W_Repeat(space, w_obj, w_times))
 
+W_Repeat.typedef = TypeDef(
+        'repeat',
+        __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']),
+        __doc__  = """Make an iterator that returns object over and over again.
+    Runs indefinitely unless the times argument is specified.  Used
+    as argument to imap() for invariant parameters to the called
+    function. Also used with izip() to create an invariant part of a
+    tuple record.
+
+    Equivalent to :
+
+    def repeat(object, times=None):
+        if times is None:
+            while True:
+                yield object
+        else:
+            for i in xrange(times):
+                yield object
+    """)

Modified: pypy/dist/pypy/module/itertools/test/test_itertools.py
==============================================================================
--- pypy/dist/pypy/module/itertools/test/test_itertools.py	(original)
+++ pypy/dist/pypy/module/itertools/test/test_itertools.py	Thu Jul 10 15:19:34 2008
@@ -1,26 +1,31 @@
 from pypy.conftest import gettestobjspace
 
 class AppTestItertools: 
-
     def setup_class(cls):
         cls.space = gettestobjspace(usemodules=['itertools'])
 
-    def test_count(self):
+    def test_iterables(self):
         import itertools
 
-        c = itertools.count()
-        for x in range(10):
-            assert c.next() == x
+        iterables = [
+            itertools.count(),
+            itertools.repeat(None),
+            ]
+
+        for it in iterables:
+            assert hasattr(it, '__iter__')
+            assert iter(it) is it
+            assert hasattr(it, 'next')
+            assert callable(it.next)
 
-    def test_count_iterable(self):
+    def test_count(self):
         import itertools
 
         c = itertools.count()
-        assert hasattr(c, '__iter__')
-        assert iter(c) is c
-        assert hasattr(c, 'next')
+        for x in range(10):
+            assert c.next() == x
 
-    def test_count_param(self):
+    def test_count_firstval(self):
         import itertools
 
         c = itertools.count(3)
@@ -35,8 +40,9 @@
         raises(OverflowError, c.next) 
         raises(OverflowError, c.next) 
 
+        raises(OverflowError, itertools.count, sys.maxint + 1)
+
     def test_repeat(self):
-        skip("for now")
         import itertools
 
         o = object()
@@ -45,3 +51,29 @@
         for x in range(10):
             assert o is r.next()
 
+    def test_repeat_times(self):
+        import itertools
+
+        times = 10
+        r = itertools.repeat(None, times=times)
+        for i in range(times):
+            r.next()
+        raises(StopIteration, r.next)
+
+        r = itertools.repeat(None, times=None)
+        for x in range(10):
+            r.next()    # Should be no StopIteration
+
+        r = itertools.repeat(None, times=0)
+        raises(StopIteration, r.next)
+        raises(StopIteration, r.next)
+
+        r = itertools.repeat(None, times=-1)
+        raises(StopIteration, r.next)
+        raises(StopIteration, r.next)
+
+    def test_repeat_overflow(self):
+        import itertools
+        import sys
+
+        raises(OverflowError, itertools.repeat, None, sys.maxint + 1)