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

[jlg at codespeak.net]

Author: jlg
Date: Thu Jul 10 16:11:51 2008
New Revision: 56417

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 - takewhile implemented

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 16:11:51 2008
@@ -5,8 +5,9 @@
     """An itertools module."""
 
     interpleveldefs = {
-        'count'    : 'interp_itertools.W_Count',
+        'count'     : 'interp_itertools.W_Count',
         '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	Thu Jul 10 16:11:51 2008
@@ -106,3 +106,52 @@
             for i in xrange(times):
                 yield object
     """)
+
+class W_TakeWhile(Wrappable):
+
+    def __init__(self, space, w_predicate, w_iterable):
+        self.space = space
+        self.w_predicate = w_predicate
+        self.iterable = space.iter(w_iterable)
+        self.stopped = False
+
+    def iter_w(self):
+        return self.space.wrap(self)
+
+    def next_w(self):
+        if self.stopped:
+            raise OperationError(self.space.w_StopIteration, self.space.w_None)
+
+        try:
+            w_obj = self.space.next(self.iterable)
+        except StopIteration:
+            raise OperationError(self.space.w_StopIteration, self.space.w_None)
+
+        w_bool = self.space.call_function(self.w_predicate, w_obj)
+        if not self.space.is_true(w_bool):
+            self.stopped = True
+            raise OperationError(self.space.w_StopIteration, self.space.w_None)
+
+        return w_obj
+
+def W_TakeWhile___new__(space, w_subtype, w_predicate, w_iterable):
+    return space.wrap(W_TakeWhile(space, w_predicate, w_iterable))
+
+
+W_TakeWhile.typedef = TypeDef(
+        'takewhile',
+        __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']),
+        __doc__  = """Make an iterator that returns elements from the iterable as
+    long as the predicate is true.
+
+    Equivalent to :
+    
+    def takewhile(predicate, iterable):
+        for x in iterable:
+            if predicate(x):
+                yield x
+            else:
+                break
+    """)

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 16:11:51 2008
@@ -10,6 +10,7 @@
         iterables = [
             itertools.count(),
             itertools.repeat(None),
+            itertools.takewhile(bool, []),
             ]
 
         for it in iterables:
@@ -77,3 +78,26 @@
         import sys
 
         raises(OverflowError, itertools.repeat, None, sys.maxint + 1)
+
+    def test_takewhile(self):
+        import itertools
+
+        tw = itertools.takewhile(bool, [])
+        raises(StopIteration, tw.next)
+
+        tw = itertools.takewhile(bool, [False, True, True])
+        raises(StopIteration, tw.next)
+
+        tw = itertools.takewhile(bool, [1, 2, 3, 0, 1, 1])
+        for x in range(3):
+            assert tw.next() == x + 1
+
+        raises(StopIteration, tw.next)
+
+    def test_takewhile_wrongargs(self):
+        import itertools
+
+        tw = itertools.takewhile(None, [1])
+        raises(TypeError, tw.next)
+
+        raises(TypeError, itertools.takewhile, bool, None)