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

[jlg at codespeak.net]

Author: jlg
Date: Fri Jul 11 16:25:59 2008
New Revision: 56452

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

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 16:25:59 2008
@@ -35,6 +35,7 @@
         'islice'        : 'interp_itertools.W_ISlice',
         'izip'          : 'interp_itertools.W_IZip',
         'repeat'        : 'interp_itertools.W_Repeat',
+        'starmap'       : 'interp_itertools.W_StarMap',
         'takewhile'     : 'interp_itertools.W_TakeWhile',
     }
 

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 16:25:59 2008
@@ -514,12 +514,8 @@
 class W_IZip(W_IMap):
     _error_name = "izip"
 
-    def __init__(self, space, args_w):
-        super(W_IZip, self).__init__(space, space.w_None, args_w)
-
-
 def W_IZip___new__(space, w_subtype, args_w):
-    return space.wrap(W_IZip(space, args_w))
+    return space.wrap(W_IZip(space, space.w_None, args_w))
 
 W_IZip.typedef = TypeDef(
         'izip',
@@ -603,3 +599,46 @@
                 yield element    
     """)
 W_Cycle.typedef.acceptable_as_base_class = False
+
+class W_StarMap(Wrappable):
+
+    def __init__(self, space, w_fun, w_iterable):
+        self.space = space
+        self.w_fun = w_fun
+        self.w_iterable = self.space.iter(w_iterable)
+
+    def iter_w(self):
+        return self.space.wrap(self)
+
+    def next_w(self):
+        w_obj = self.space.next(self.w_iterable)
+        if not self.space.is_true(self.space.isinstance(w_obj, self.space.w_tuple)):
+            raise OperationError(self.space.w_TypeError, self.space.wrap("iterator must return a tuple"))
+
+        return self.space.call(self.w_fun, w_obj)
+
+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(
+        'starmap',
+        __new__  = interp2app(W_StarMap___new__, unwrap_spec=[ObjSpace, W_Root, W_Root, W_Root]),
+        __iter__ = interp2app(W_StarMap.iter_w, unwrap_spec=['self']),
+        next     = interp2app(W_StarMap.next_w, unwrap_spec=['self']),
+        __doc__  = """Make an iterator that computes the function using arguments
+    tuples obtained from the iterable. Used instead of imap() when
+    argument parameters are already grouped in tuples from a single
+    iterable (the data has been ``pre-zipped''). The difference
+    between imap() and starmap() parallels the distinction between
+    function(a,b) and function(*c).
+
+    Equivalent to :
+
+    def starmap(function, iterable):
+        iterable = iter(iterable)
+        while True:
+            yield function(*iterable.next())    
+    """)
+W_StarMap.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 16:25:59 2008
@@ -356,6 +356,29 @@
         for x in [1, 2, 3, 1, 2, 3, 1, 2, 3]:
             assert it.next() == x
 
+    def test_starmap(self):
+        import itertools, operator
+
+        it = itertools.starmap(operator.add, [])
+        raises(StopIteration, it.next)
+
+        it = itertools.starmap(operator.add, [(0, 1), (2, 3), (4, 5)])
+        for x in [1, 5, 9]:
+            assert it.next() == x
+        raises(StopIteration, it.next)
+
+    def test_starmap_wrongargs(self):
+        import itertools
+
+        it = itertools.starmap(None, [(1, )])
+        raises(TypeError, it.next)
+
+        it = itertools.starmap(None, [])
+        raises(StopIteration, it.next)
+
+        it = itertools.starmap(bool, [0])
+        raises(TypeError, it.next)
+
     
     def test_iterables(self):
         import itertools
@@ -371,6 +394,7 @@
             itertools.islice([], 0),
             itertools.izip(),
             itertools.repeat(None),
+            itertools.starmap(bool, []),
             itertools.takewhile(bool, []),
             ]
     
@@ -395,6 +419,7 @@
             itertools.islice,
             itertools.izip,
             itertools.repeat,
+            itertools.starmap,
             itertools.takewhile,
             ]
         for method in methods:
@@ -414,6 +439,7 @@
             (itertools.islice, ([], 0)),
             (itertools.izip, ()),
             (itertools.repeat, (None,)),
+            (itertools.starmap, (bool, [])),
             (itertools.takewhile, (bool, [])),
             ]
         for cls, args in iterables: