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

[adurdin at codespeak.net]

Author: adurdin
Date: Fri Jul 11 15:40:15 2008
New Revision: 56450

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) interp_itertools - added cycle()


Modified: pypy/dist/pypy/module/itertools/__init__.py
==============================================================================
--- pypy/dist/pypy/module/itertools/__init__.py	(original)
+++ pypy/dist/pypy/module/itertools/__init__.py	Fri Jul 11 15:40:15 2008
@@ -25,16 +25,17 @@
     """
 
     interpleveldefs = {
-        'chain'     : 'interp_itertools.W_Chain',
-        'count'     : 'interp_itertools.W_Count',
-        'dropwhile' : 'interp_itertools.W_DropWhile',
-        'ifilter'   : 'interp_itertools.W_IFilter',
-        'ifilterfalse' : 'interp_itertools.W_IFilterFalse',
-        'imap'      : 'interp_itertools.W_IMap',
-        'islice'    : 'interp_itertools.W_ISlice',
-        'izip'      : 'interp_itertools.W_IZip',
-        'repeat'    : 'interp_itertools.W_Repeat',
-        'takewhile' : 'interp_itertools.W_TakeWhile',
+        'chain'         : 'interp_itertools.W_Chain',
+        'count'         : 'interp_itertools.W_Count',
+        'cycle'         : 'interp_itertools.W_Cycle',
+        'dropwhile'     : 'interp_itertools.W_DropWhile',
+        'ifilter'       : 'interp_itertools.W_IFilter',
+        'ifilterfalse'  : 'interp_itertools.W_IFilterFalse',
+        'imap'          : 'interp_itertools.W_IMap',
+        'islice'        : 'interp_itertools.W_ISlice',
+        'izip'          : 'interp_itertools.W_IZip',
+        'repeat'        : 'interp_itertools.W_Repeat',
+        'takewhile'     : 'interp_itertools.W_TakeWhile',
     }
 
     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	Fri Jul 11 15:40:15 2008
@@ -541,3 +541,46 @@
     """)
 W_IZip.typedef.acceptable_as_base_class = False
 
+
+class W_Cycle(Wrappable):
+
+    def __init__(self, space, w_iterable):
+        self.space = space
+        self.saved_w = []
+        self.w_iterable = space.iter(w_iterable)
+        self.saved_iterator = None
+        self.exhausted = False
+
+    def next_w(self):
+        if self.exhausted:
+            if not self.saved_w:
+                raise OperationError(self.space.w_StopIteration, self.space.w_None)
+            try:
+                w_obj = self.saved_iterator.next()
+            except StopIteration:
+                self.saved_iterator = iter(self.saved_w)
+                w_obj = self.saved_iterator.next()
+        else:
+            try:
+                w_obj = self.space.next(self.w_iterable)
+            except OperationError, e:
+                if e.match(self.space, self.space.w_StopIteration):
+                    self.exhausted = True
+                    if not self.saved_w:
+                        raise
+                    self.saved_iterator = iter(self.saved_w)
+                    w_obj = self.saved_iterator.next()
+                else:
+                    raise
+            else:
+                self.saved_w.append(w_obj)
+        return w_obj
+
+def W_Cycle___new__(space, w_subtype, w_iterable):
+    return space.wrap(W_Cycle(space, w_iterable))
+
+W_Cycle.typedef = TypeDef(
+        'cycle',
+        __new__  = interp2app(W_Cycle___new__, unwrap_spec=[ObjSpace, W_Root, W_Root]),
+        next     = interp2app(W_Cycle.next_w, unwrap_spec=['self']))
+W_Cycle.typedef.acceptable_as_base_class = False

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	Fri Jul 11 15:40:15 2008
@@ -368,6 +368,17 @@
             else:
                 fail("TypeError expected")
 
+    def test_cycle(self):
+        import itertools
+
+        it = itertools.cycle([])
+        raises(StopIteration, it.next)
+        
+        it = itertools.cycle([1, 2, 3])
+        for x in [1, 2, 3, 1, 2, 3, 1, 2, 3]:
+            assert it.next() == x
+
+
     def test_docstrings(self):
         import itertools