[pypy-svn] pypy default: Fix(?) corner cases of combinations_with_replacement().

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: 
Changeset: r41346:228591eb80ad
Date: 2011-01-26 14:19 +0100
http://bitbucket.org/pypy/pypy/changeset/228591eb80ad/

Log:	Fix(?) corner cases of combinations_with_replacement().

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
@@ -884,6 +884,15 @@
         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")]
 
+    def test_combinations_with_replacement_shortcases(self):
+        from itertools import combinations_with_replacement
+        assert list(combinations_with_replacement([-12], 2)) == [(-12, -12)]
+        assert list(combinations_with_replacement("AB", 3)) == [
+            ("A", "A", "A"), ("A", "A", "B"),
+            ("A", "B", "B"), ("B", "B", "B")]
+        assert list(combinations_with_replacement([], 2)) == []
+        assert list(combinations_with_replacement([], 0)) == [()]
+
     def test_izip_longest3(self):
         import itertools
         class Repeater(object):

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
@@ -1236,6 +1236,10 @@
 )
 
 class W_CombinationsWithReplacement(W_Combinations):
+    def __init__(self, space, pool_w, indices, r):
+        W_Combinations.__init__(self, space, pool_w, indices, r)
+        self.stopped = len(pool_w) == 0 and r > 0
+
     def get_maximum(self, i):
         return len(self.pool_w) - 1
 
@@ -1248,7 +1252,7 @@
     if r < 0:
         raise OperationError(space.w_ValueError,
                              space.wrap("r must be non-negative"))
-    indices = [0] * len(pool_w)
+    indices = [0] * r
     res = space.allocate_instance(W_CombinationsWithReplacement, w_subtype)
     res.__init__(space, pool_w, indices, r)
     return space.wrap(res)