diff options
Diffstat (limited to 'Lib/test/test_itertools.py')
| -rw-r--r-- | Lib/test/test_itertools.py | 40 |
1 files changed, 32 insertions, 8 deletions
diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index 2182cb9..a62bad2 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -191,6 +191,21 @@ class TestBasicOps(unittest.TestCase): self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1) self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1) + def test_compress(self): + self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF')) + self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list('')) + self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF')) + self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC')) + self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC')) + n = 10000 + data = chain.from_iterable(repeat(range(6), n)) + selectors = chain.from_iterable(repeat((0, 1))) + self.assertEqual(list(compress(data, selectors)), [1,3,5] * n) + self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable + self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable + self.assertRaises(TypeError, compress, range(6)) # too few args + self.assertRaises(TypeError, compress, range(6), None) # too many args + def test_count(self): self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)]) self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)]) @@ -701,6 +716,9 @@ class TestExamples(unittest.TestCase): self.assertEqual(list(combinations(range(4), 3)), [(0,1,2), (0,1,3), (0,2,3), (1,2,3)]) + def test_compress(self): + self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF')) + def test_count(self): self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14]) @@ -781,6 +799,10 @@ class TestGC(unittest.TestCase): a = [] self.makecycle(combinations([1,2,a,3], 3), a) + def test_compress(self): + a = [] + self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a) + def test_cycle(self): a = [] self.makecycle(cycle([a]*2), a) @@ -934,6 +956,15 @@ class TestVariousIteratorArgs(unittest.TestCase): self.assertRaises(TypeError, list, chain(N(s))) self.assertRaises(ZeroDivisionError, list, chain(E(s))) + def test_compress(self): + for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)): + n = len(s) + for g in (G, I, Ig, S, L, R): + self.assertEqual(list(compress(g(s), repeat(1))), list(g(s))) + self.assertRaises(TypeError, compress, X(s), repeat(1)) + self.assertRaises(TypeError, list, compress(N(s), repeat(1))) + self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1))) + def test_product(self): for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)): self.assertRaises(TypeError, product, X(s)) @@ -1125,7 +1156,7 @@ class SubclassWithKwargsTest(unittest.TestCase): def test_keywords_in_subclass(self): # count is not subclassable... for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap, - starmap, islice, takewhile, dropwhile, cycle): + starmap, islice, takewhile, dropwhile, cycle, compress): class Subclass(cls): def __init__(self, newarg=None, *args): cls.__init__(self, *args) @@ -1262,10 +1293,6 @@ Samuele ... for n in xrange(2**len(pairs)): ... yield set(x for m, x in pairs if m&n) ->>> def compress(data, selectors): -... "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F" -... return (d for d, s in izip(data, selectors) if s) - >>> def combinations_with_replacement(iterable, r): ... "combinations_with_replacement('ABC', 3) --> AA AB AC BB BC CC" ... pool = tuple(iterable) @@ -1361,9 +1388,6 @@ perform as purported. >>> map(sorted, powerset('ab')) [[], ['a'], ['b'], ['a', 'b']] ->>> list(compress('abcdef', [1,0,1,0,1,1])) -['a', 'c', 'e', 'f'] - >>> list(combinations_with_replacement('abc', 2)) [('a', 'a'), ('a', 'b'), ('a', 'c'), ('b', 'b'), ('b', 'c'), ('c', 'c')] |