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

[adurdin at codespeak.net]

Author: adurdin
Date: Fri Jul 11 17:14:18 2008
New Revision: 56456

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 - tee() done, only one left!


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 17:14:18 2008
@@ -37,6 +37,7 @@
         'repeat'        : 'interp_itertools.W_Repeat',
         'starmap'       : 'interp_itertools.W_StarMap',
         'takewhile'     : 'interp_itertools.W_TakeWhile',
+        'tee'           : 'interp_itertools.tee',
     }
 
     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 17:14:18 2008
@@ -617,7 +617,7 @@
 
         return self.space.call(self.w_fun, w_obj)
 
-def W_StarMap___new__(space, w_subtype, w_fun,w_iterable):
+def W_StarMap___new__(space, w_subtype, w_fun, w_iterable):
     return space.wrap(W_StarMap(space, w_fun, w_iterable))
 
 W_StarMap.typedef = TypeDef(
@@ -642,3 +642,73 @@
 W_StarMap.typedef.acceptable_as_base_class = False
 
 
+def tee(space, w_iterable, n=2):
+    """Return n independent iterators from a single iterable.
+    Note : once tee() has made a split, the original iterable
+    should not be used anywhere else; otherwise, the iterable could get
+    advanced without the tee objects being informed.
+    
+    Note : this member of the toolkit may require significant auxiliary
+    storage (depending on how much temporary data needs to be stored).
+    In general, if one iterator is going to use most or all of the
+    data before the other iterator, it is faster to use list() instead
+    of tee()
+    
+    Equivalent to :
+    
+    def tee(iterable, n=2):
+        def gen(next, data={}, cnt=[0]):
+            for i in count():
+                if i == cnt[0]:
+                    item = data[i] = next()
+                    cnt[0] += 1
+                else:
+                    item = data.pop(i)
+                yield item
+        it = iter(iterable)
+        return tuple([gen(it.next) for i in range(n)])
+    """
+    if n < 0:
+        raise OperationError(space.w_ValueError, space.wrap("n must be >= 0"))
+    
+    tee_state = TeeState(space, w_iterable)
+    iterators_w = [space.wrap(W_TeeIterable(space, tee_state)) for x in range(n)]
+    return space.newtuple(iterators_w)
+tee.unwrap_spec = [ObjSpace, W_Root, int]
+
+class TeeState(object):
+    def __init__(self, space, w_iterable):
+        self.space = space
+        self.w_iterable = self.space.iter(w_iterable)
+        self.num_saved = 0
+        self.saved_w = []
+
+    def get_next(self, index):
+        if index >= self.num_saved:
+            w_obj = self.space.next(self.w_iterable)
+            self.saved_w.append(w_obj)
+            self.num_saved += 1
+            return w_obj
+        else:
+            return self.saved_w[index]
+
+class W_TeeIterable(Wrappable):
+    def __init__(self, space, tee_state):
+        self.space = space
+        self.tee_state = tee_state
+        self.index = 0
+
+    def iter_w(self):
+        return self.space.wrap(self)
+
+    def next_w(self):
+        try:
+            w_obj = self.tee_state.get_next(self.index)
+            return w_obj
+        finally:
+            self.index += 1
+
+W_TeeIterable.typedef = TypeDef(
+        '_tee',
+        __iter__ = interp2app(W_TeeIterable.iter_w, unwrap_spec=['self']),
+        next     = interp2app(W_TeeIterable.next_w, unwrap_spec=['self']))

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 17:14:18 2008
@@ -379,7 +379,38 @@
         it = itertools.starmap(bool, [0])
         raises(TypeError, it.next)
 
-    
+    def test_tee(self):
+        import itertools
+
+        it1, it2 = itertools.tee([])
+        raises(StopIteration, it1.next)
+        raises(StopIteration, it2.next)
+
+        it1, it2 = itertools.tee([1, 2, 3])
+        for x in [1, 2]:
+            assert it1.next() == x
+        for x in [1, 2, 3]:
+            assert it2.next() == x
+        assert it1.next() == 3
+        raises(StopIteration, it1.next)
+        raises(StopIteration, it2.next)
+
+        assert itertools.tee([], 0) == ()
+
+        iterators = itertools.tee([1, 2, 3], 10)
+        for it in iterators:
+            for x in [1, 2, 3]:
+                assert it.next() == x
+            raises(StopIteration, it.next)
+
+    def test_iterables_wrongargs(self):
+        import itertools
+        
+        raises(TypeError, itertools.tee, 0)
+        raises(ValueError, itertools.tee, [], -1)
+        raises(TypeError, itertools.tee, [], None)
+
+
     def test_iterables(self):
         import itertools
     
@@ -396,6 +427,8 @@
             itertools.repeat(None),
             itertools.starmap(bool, []),
             itertools.takewhile(bool, []),
+            itertools.tee([])[0],
+            itertools.tee([])[1],
             ]
     
         for it in iterables:
@@ -403,7 +436,7 @@
             assert iter(it) is it
             assert hasattr(it, 'next')
             assert callable(it.next)
-    
+
     def test_docstrings(self):
         import itertools
         
@@ -421,6 +454,7 @@
             itertools.repeat,
             itertools.starmap,
             itertools.takewhile,
+            itertools.tee,
             ]
         for method in methods:
             assert method.__doc__