[pypy-svn] pypy default: Implemented itertools.combinations_with_replacement (look ma, codereuse!).

alex_gaynor commits-noreply at bitbucket.org
Fri Jan 21 01:00:05 CET 2011


Author: Alex Gaynor <alex.gaynor at gmail.com>
Branch: 
Changeset: r41103:eb75eb40f65c
Date: 2011-01-20 17:59 -0600
http://bitbucket.org/pypy/pypy/changeset/eb75eb40f65c/

Log:	Implemented itertools.combinations_with_replacement (look ma,
	codereuse!).

diff --git a/pypy/module/itertools/__init__.py b/pypy/module/itertools/__init__.py
--- a/pypy/module/itertools/__init__.py
+++ b/pypy/module/itertools/__init__.py
@@ -27,6 +27,7 @@
     interpleveldefs = {
         'chain'         : 'interp_itertools.W_Chain',
         'combinations'  : 'interp_itertools.W_Combinations',
+        'combinations_with_replacement' : 'interp_itertools.W_CombinationsWithReplacement',
         'compress'      : 'interp_itertools.W_Compress',
         'count'         : 'interp_itertools.W_Count',
         'cycle'         : 'interp_itertools.W_Cycle',

diff --git a/pypy/module/itertools/test/test_itertools.py b/pypy/module/itertools/test/test_itertools.py
--- a/pypy/module/itertools/test/test_itertools.py
+++ b/pypy/module/itertools/test/test_itertools.py
@@ -766,7 +766,16 @@
 
         raises(TypeError, combinations, "abc")
         raises(TypeError, combinations, "abc", 2, 1)
-        raises(TypeError, None)
+        raises(TypeError, combinations, None)
         raises(ValueError, combinations, "abc", -2)
         assert list(combinations("abc", 32)) == []
         assert list(combinations(range(4), 3)) == [(0, 1, 2), (0, 1, 3), (0, 2, 3), (1, 2, 3)]
+
+    def test_combinations_with_replacement(self):
+        from itertools import combinations_with_replacement
+
+        raises(TypeError, combinations_with_replacement, "abc")
+        raises(TypeError, combinations_with_replacement, "abc", 2, 1)
+        raises(TypeError, combinations_with_replacement, None)
+        raises(ValueError, combinations_with_replacement, "abc", -2)
+        assert list(combinations_with_replacement("ABC", 2)) == [("A", "A"), ("A", 'B'), ("A", "C"), ("B", "B"), ("B", "C"), ("C", "C")]

diff --git a/pypy/module/itertools/interp_itertools.py b/pypy/module/itertools/interp_itertools.py
--- a/pypy/module/itertools/interp_itertools.py
+++ b/pypy/module/itertools/interp_itertools.py
@@ -1109,15 +1109,20 @@
         self.last_result_w = None
         self.stopped = r > len(pool_w)
 
+    def get_maximum(self, i):
+        return i + len(self.pool_w) - self.r
+
+    def max_index(self, j):
+        return self.indices[j - 1] + 1
+
     @unwrap_spec(ObjSpace, W_Root, W_Root, int)
     def descr__new__(space, w_subtype, w_iterable, r):
         pool_w = space.fixedview(w_iterable)
-        n = len(pool_w)
         if r < 0:
             raise OperationError(space.w_ValueError,
                 space.wrap("r must be non-negative")
             )
-        indices = range(r)
+        indices = range(len(pool_w))
         return W_Combinations(space, pool_w, indices, r)
 
     @unwrap_spec("self", ObjSpace)
@@ -1138,9 +1143,9 @@
             # Copy the previous result
             result_w = self.last_result_w[:]
             # Scan indices right-to-left until finding one that is not at its
-            # maximum (i + n - r).
+            # maximum
             i = self.r - 1
-            while i >= 0 and self.indices[i] == i + len(self.pool_w) - self.r:
+            while i >= 0 and self.indices[i] == self.get_maximum(i):
                 i -= 1
 
             # If i is negative, then the indices are all at their maximum value
@@ -1151,11 +1156,10 @@
 
             # Increment the current index which we know is not at its maximum.
             # Then move back to the right setting each index to its lowest
-            # possible value (one higher than the index to its left -- this
-            # maintains the sort order invariant).
+            # possible value
             self.indices[i] += 1
             for j in xrange(i + 1, self.r):
-                self.indices[j] = self.indices[j-1] + 1
+                self.indices[j] = self.max_index(j)
 
             # Update the result for the new indices starting with i, the
             # leftmost index that changed
@@ -1171,3 +1175,26 @@
     __iter__ = interp2app(W_Combinations.descr__iter__),
     next = interp2app(W_Combinations.descr_next),
 )
+
+class W_CombinationsWithReplacement(W_Combinations):
+    def get_maximum(self, i):
+        return len(self.pool_w) - 1
+
+    def max_index(self, j):
+        return self.indices[j - 1]
+
+    @unwrap_spec(ObjSpace, W_Root, W_Root, int)
+    def descr__new__(space, w_subtype, w_iterable, r):
+        pool_w = space.fixedview(w_iterable)
+        if r < 0:
+            raise OperationError(space.w_ValueError,
+                space.wrap("r must be non-negative")
+            )
+        indices = [0] * len(pool_w)
+        return W_CombinationsWithReplacement(space, pool_w, indices, r)
+
+W_CombinationsWithReplacement.typedef = TypeDef("combinations_with_replacement",
+    __new__ = interp2app(W_CombinationsWithReplacement.descr__new__.im_func),
+    __iter__ = interp2app(W_CombinationsWithReplacement.descr__iter__),
+    next = interp2app(W_CombinationsWithReplacement.descr_next),
+)


More information about the Pypy-commit mailing list