[pypy-svn] r78878 - in pypy/branch/fast-forward: lib_pypy lib_pypy/pypy_test pypy/module/itertools

dcolish at codespeak.net dcolish at codespeak.net
Mon Nov 8 17:07:06 CET 2010


Author: dcolish
Date: Mon Nov  8 17:07:04 2010
New Revision: 78878

Added:
   pypy/branch/fast-forward/lib_pypy/pypy_test/test_itertools.py
Modified:
   pypy/branch/fast-forward/lib_pypy/itertools.py
   pypy/branch/fast-forward/pypy/module/itertools/interp_itertools.py
Log:
add non-module equiv of compress and product. simplify W_Compress.next_w

Modified: pypy/branch/fast-forward/lib_pypy/itertools.py
==============================================================================
--- pypy/branch/fast-forward/lib_pypy/itertools.py	(original)
+++ pypy/branch/fast-forward/lib_pypy/itertools.py	Mon Nov  8 17:07:04 2010
@@ -60,6 +60,23 @@
                 raise TypeError('%s has no next() method' % \
                                 (self._cur_iterable_iter))
 
+
+class compress(object):
+    def __init__(self, data, selectors):
+        self.data = iter(data)
+        self.selectors = iter(selectors)
+
+    def __iter__(self):
+        return self
+
+    def next(self):
+        while True:
+            next_item = self.data.next()
+            next_selector = self.selectors.next()
+            if bool(next_selector):
+                return next_item
+
+
 class count(object):
     """Make an iterator that returns consecutive integers starting
     with n.  If not specified n defaults to zero. Does not currently
@@ -409,7 +426,56 @@
         except AttributeError:
             # CPython raises a TypeError when next() is not defined
             raise TypeError('%s has no next() method' % (i))
-    
+
+
+class product(object):
+
+    def __init__(self, *args, **kw):
+        if len(kw) > 1:
+            raise TypeError("product() takes at most 1 argument (%d given)" %
+                             len(kw))
+        self.repeat = kw.get('repeat', 1)
+        self.gears = [x for x in args] * self.repeat
+        self.num_gears = len(self.gears)
+        # initialization of indicies to loop over
+        self.indicies = [(0, len(self.gears[x]))
+                         for x in range(0, self.num_gears)]
+        self.cont = True
+
+    def roll_gears(self):
+        # Starting from the end of the gear indicies work to the front
+        # incrementing the gear until the limit is reached. When the limit
+        # is reached carry operation to the next gear
+        should_carry = True
+        for n in range(0, self.num_gears):
+            nth_gear = self.num_gears - n - 1
+            if should_carry:
+                count, lim = self.indicies[nth_gear]
+                count += 1
+                if count == lim and nth_gear == 0:
+                    self.cont = False
+                if count == lim:
+                    should_carry = True
+                    count = 0
+                else:
+                    should_carry = False
+                self.indicies[nth_gear] = (count, lim)
+            else:
+                break
+
+    def __iter__(self):
+        return self
+
+    def next(self):
+        if not self.cont:
+            raise StopIteration
+        l = []
+        for x in range(0, self.num_gears):
+            index, limit = self.indicies[x]
+            l.append(self.gears[x][index])
+        self.roll_gears()
+        return tuple(l)
+
 
 class repeat(object):
     """Make an iterator that returns object over and over again.

Added: pypy/branch/fast-forward/lib_pypy/pypy_test/test_itertools.py
==============================================================================
--- (empty file)
+++ pypy/branch/fast-forward/lib_pypy/pypy_test/test_itertools.py	Mon Nov  8 17:07:04 2010
@@ -0,0 +1,50 @@
+from py.test import raises
+from lib_pypy import itertools
+
+class TestItertools(object):
+
+    def test_compress(self):
+        it = itertools.compress(['a', 'b', 'c'], [0, 1, 0])
+
+        assert list(it) == ['b']
+
+    def test_compress_diff_len(self):
+        it = itertools.compress(['a'], [])
+        raises(StopIteration, it.next)
+
+    def test_product(self):
+        l = [1, 2]
+        m = ['a', 'b']
+
+        prodlist = itertools.product(l, m)
+        assert list(prodlist) == [(1, 'a'), (1, 'b'), (2, 'a'), (2, 'b')]
+
+    def test_product_repeat(self):
+        l = [1, 2]
+        m = ['a', 'b']
+
+        prodlist = itertools.product(l, m, repeat=2)
+        ans = [(1, 'a', 1, 'a'), (1, 'a', 1, 'b'), (1, 'a', 2, 'a'),
+               (1, 'a', 2, 'b'), (1, 'b', 1, 'a'), (1, 'b', 1, 'b'),
+               (1, 'b', 2, 'a'), (1, 'b', 2, 'b'), (2, 'a', 1, 'a'),
+               (2, 'a', 1, 'b'), (2, 'a', 2, 'a'), (2, 'a', 2, 'b'),
+               (2, 'b', 1, 'a'), (2, 'b', 1, 'b'), (2, 'b', 2, 'a'),
+               (2, 'b', 2, 'b')]
+        assert list(prodlist) == ans
+
+    def test_product_diff_sizes(self):
+        l = [1, 2]
+        m = ['a']
+
+        prodlist = itertools.product(l, m)
+        assert list(prodlist) == [(1, 'a'), (2, 'a')]
+
+        l = [1]
+        m = ['a', 'b']
+        prodlist = itertools.product(l, m)
+        assert list(prodlist) == [(1, 'a'), (1, 'b')]
+
+    def test_product_toomany_args(self):
+        l = [1, 2]
+        m = ['a']
+        raises(TypeError, itertools.product, l, m, repeat=1, foo=2)

Modified: pypy/branch/fast-forward/pypy/module/itertools/interp_itertools.py
==============================================================================
--- pypy/branch/fast-forward/pypy/module/itertools/interp_itertools.py	(original)
+++ pypy/branch/fast-forward/pypy/module/itertools/interp_itertools.py	Mon Nov  8 17:07:04 2010
@@ -954,9 +954,8 @@
         return self.space.wrap(self)
 
     def next_w(self):
-        if not self.w_data or not self.w_selectors:
-            raise OperationError(self.space.w_StopIteration, self.space.w_None)
-
+        # No need to check for StopIteration since either w_data
+        # or w_selectors will raise this. The shortest one stops first.
         while True:
             w_next_item = self.space.next(self.w_data)
             w_next_selector = self.space.next(self.w_selectors)



More information about the Pypy-commit mailing list