blob: 567210425756ee44be8f4d8028dfa1bb5c3260f2 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001import unittest
2from test import test_support
3from itertools import *
Antoine Pitrou3ec903f2014-04-29 12:13:46 +02004import weakref
Raymond Hettingeraa044612009-02-12 12:04:26 +00005from decimal import Decimal
Raymond Hettingerdbe3bfb2009-02-12 12:43:01 +00006from fractions import Fraction
Raymond Hettinger2012f172003-02-07 05:32:58 +00007import sys
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00008import operator
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00009import random
Raymond Hettingere09f45a2009-11-30 19:44:40 +000010import copy
11import pickle
Ezio Melottidde5b942010-02-03 05:37:26 +000012from functools import reduce
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +000013maxsize = test_support.MAX_Py_ssize_t
14minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000015
16def onearg(x):
17 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000018 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000019
20def errfunc(*args):
21 'Test function that raises an error'
22 raise ValueError
23
24def gen3():
25 'Non-restartable source sequence'
26 for i in (0, 1, 2):
27 yield i
28
29def isEven(x):
30 'Test predicate'
31 return x%2==0
32
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000033def isOdd(x):
34 'Test predicate'
35 return x%2==1
36
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000037class StopNow:
38 'Class emulating an empty iterable.'
39 def __iter__(self):
40 return self
41 def next(self):
42 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000043
Raymond Hettinger02420702003-06-29 20:36:23 +000044def take(n, seq):
45 'Convenience function for partially consuming a long of infinite iterable'
46 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000047
Raymond Hettingerd553d852008-03-04 04:17:08 +000048def prod(iterable):
49 return reduce(operator.mul, iterable, 1)
50
Raymond Hettinger93e804d2008-02-26 23:40:50 +000051def fact(n):
52 'Factorial'
Raymond Hettingerd553d852008-03-04 04:17:08 +000053 return prod(range(1, n+1))
54
Raymond Hettinger96ef8112003-02-01 00:10:11 +000055class TestBasicOps(unittest.TestCase):
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000056 def test_chain(self):
Raymond Hettingerad47fa12008-03-06 20:52:01 +000057
58 def chain2(*iterables):
59 'Pure python version in the docs'
60 for it in iterables:
61 for element in it:
62 yield element
63
64 for c in (chain, chain2):
65 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
66 self.assertEqual(list(c('abc')), list('abc'))
67 self.assertEqual(list(c('')), [])
68 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
69 self.assertRaises(TypeError, list,c(2, 3))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000070
Raymond Hettingerb4cbc982008-02-28 22:46:41 +000071 def test_chain_from_iterable(self):
72 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
73 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
74 self.assertEqual(list(chain.from_iterable([''])), [])
75 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
76 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
77
Raymond Hettinger93e804d2008-02-26 23:40:50 +000078 def test_combinations(self):
Raymond Hettinger5b913e32009-01-08 06:39:04 +000079 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Raymond Hettinger93e804d2008-02-26 23:40:50 +000080 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Raymond Hettingerd553d852008-03-04 04:17:08 +000081 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Raymond Hettinger93e804d2008-02-26 23:40:50 +000082 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Raymond Hettinger5b913e32009-01-08 06:39:04 +000083 self.assertEqual(list(combinations('abc', 32)), []) # r > n
Raymond Hettinger93e804d2008-02-26 23:40:50 +000084 self.assertEqual(list(combinations(range(4), 3)),
85 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
Raymond Hettingerd553d852008-03-04 04:17:08 +000086
87 def combinations1(iterable, r):
88 'Pure python version shown in the docs'
89 pool = tuple(iterable)
90 n = len(pool)
Raymond Hettinger5b913e32009-01-08 06:39:04 +000091 if r > n:
92 return
Raymond Hettingerd553d852008-03-04 04:17:08 +000093 indices = range(r)
94 yield tuple(pool[i] for i in indices)
95 while 1:
96 for i in reversed(range(r)):
97 if indices[i] != i + n - r:
98 break
99 else:
100 return
101 indices[i] += 1
102 for j in range(i+1, r):
103 indices[j] = indices[j-1] + 1
104 yield tuple(pool[i] for i in indices)
105
106 def combinations2(iterable, r):
107 'Pure python version shown in the docs'
108 pool = tuple(iterable)
109 n = len(pool)
110 for indices in permutations(range(n), r):
111 if sorted(indices) == list(indices):
112 yield tuple(pool[i] for i in indices)
113
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000114 def combinations3(iterable, r):
115 'Pure python version from cwr()'
116 pool = tuple(iterable)
117 n = len(pool)
118 for indices in combinations_with_replacement(range(n), r):
119 if len(set(indices)) == r:
120 yield tuple(pool[i] for i in indices)
121
Raymond Hettingerd553d852008-03-04 04:17:08 +0000122 for n in range(7):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000123 values = [5*x-12 for x in range(n)]
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000124 for r in range(n+2):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000125 result = list(combinations(values, r))
Ezio Melottidde5b942010-02-03 05:37:26 +0000126 self.assertEqual(len(result), 0 if r>n else fact(n) // fact(r) // fact(n-r)) # right number of combs
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000127 self.assertEqual(len(result), len(set(result))) # no repeats
128 self.assertEqual(result, sorted(result)) # lexicographic order
129 for c in result:
130 self.assertEqual(len(c), r) # r-length combinations
131 self.assertEqual(len(set(c)), r) # no duplicate elements
132 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Peterson5c8da862009-06-30 22:57:08 +0000133 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000134 self.assertEqual(list(c),
135 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Raymond Hettingerd553d852008-03-04 04:17:08 +0000136 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000137 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000138 self.assertEqual(result, list(combinations3(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000139
Benjamin Peterson021dec12015-02-01 20:59:00 -0500140 @test_support.bigaddrspacetest
141 def test_combinations_overflow(self):
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +0200142 with self.assertRaises((OverflowError, MemoryError)):
Benjamin Peterson021dec12015-02-01 20:59:00 -0500143 combinations("AA", 2**29)
144
Alex Gaynor97376482011-07-17 16:44:31 -0700145 @test_support.impl_detail("tuple reuse is specific to CPython")
146 def test_combinations_tuple_reuse(self):
Raymond Hettingerd553d852008-03-04 04:17:08 +0000147 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
148 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
149
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000150 def test_combinations_with_replacement(self):
151 cwr = combinations_with_replacement
152 self.assertRaises(TypeError, cwr, 'abc') # missing r argument
153 self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
154 self.assertRaises(TypeError, cwr, None) # pool is not iterable
155 self.assertRaises(ValueError, cwr, 'abc', -2) # r is negative
156 self.assertEqual(list(cwr('ABC', 2)),
157 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
158
159 def cwr1(iterable, r):
160 'Pure python version shown in the docs'
161 # number items returned: (n+r-1)! / r! / (n-1)! when n>0
162 pool = tuple(iterable)
163 n = len(pool)
164 if not n and r:
165 return
166 indices = [0] * r
167 yield tuple(pool[i] for i in indices)
168 while 1:
169 for i in reversed(range(r)):
170 if indices[i] != n - 1:
171 break
172 else:
173 return
174 indices[i:] = [indices[i] + 1] * (r - i)
175 yield tuple(pool[i] for i in indices)
176
177 def cwr2(iterable, r):
178 'Pure python version shown in the docs'
179 pool = tuple(iterable)
180 n = len(pool)
181 for indices in product(range(n), repeat=r):
182 if sorted(indices) == list(indices):
183 yield tuple(pool[i] for i in indices)
184
185 def numcombs(n, r):
186 if not n:
187 return 0 if r else 1
Ezio Melottidde5b942010-02-03 05:37:26 +0000188 return fact(n+r-1) // fact(r) // fact(n-1)
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000189
190 for n in range(7):
191 values = [5*x-12 for x in range(n)]
192 for r in range(n+2):
193 result = list(cwr(values, r))
194
195 self.assertEqual(len(result), numcombs(n, r)) # right number of combs
196 self.assertEqual(len(result), len(set(result))) # no repeats
197 self.assertEqual(result, sorted(result)) # lexicographic order
198
199 regular_combs = list(combinations(values, r)) # compare to combs without replacement
200 if n == 0 or r <= 1:
Ezio Melotti2623a372010-11-21 13:34:58 +0000201 self.assertEqual(result, regular_combs) # cases that should be identical
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000202 else:
Benjamin Peterson5c8da862009-06-30 22:57:08 +0000203 self.assertTrue(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000204
205 for c in result:
206 self.assertEqual(len(c), r) # r-length combinations
207 noruns = [k for k,v in groupby(c)] # combo without consecutive repeats
208 self.assertEqual(len(noruns), len(set(noruns))) # no repeats other than consecutive
209 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Peterson5c8da862009-06-30 22:57:08 +0000210 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000211 self.assertEqual(noruns,
212 [e for e in values if e in c]) # comb is a subsequence of the input iterable
213 self.assertEqual(result, list(cwr1(values, r))) # matches first pure python version
214 self.assertEqual(result, list(cwr2(values, r))) # matches second pure python version
215
Benjamin Peterson17845c12015-02-01 21:10:47 -0500216 @test_support.bigaddrspacetest
217 def test_combinations_with_replacement_overflow(self):
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +0200218 with self.assertRaises((OverflowError, MemoryError)):
Benjamin Peterson17845c12015-02-01 21:10:47 -0500219 combinations_with_replacement("AA", 2**30)
220
Alex Gaynor97376482011-07-17 16:44:31 -0700221 @test_support.impl_detail("tuple reuse is specific to CPython")
222 def test_combinations_with_replacement_tuple_reuse(self):
223 cwr = combinations_with_replacement
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000224 self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
225 self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
226
Raymond Hettingerd553d852008-03-04 04:17:08 +0000227 def test_permutations(self):
228 self.assertRaises(TypeError, permutations) # too few arguments
229 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000230 self.assertRaises(TypeError, permutations, None) # pool is not iterable
231 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000232 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000233 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Raymond Hettingerd553d852008-03-04 04:17:08 +0000234 self.assertEqual(list(permutations(range(3), 2)),
235 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
236
237 def permutations1(iterable, r=None):
238 'Pure python version shown in the docs'
239 pool = tuple(iterable)
240 n = len(pool)
241 r = n if r is None else r
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000242 if r > n:
243 return
Raymond Hettingerd553d852008-03-04 04:17:08 +0000244 indices = range(n)
Raymond Hettingere70bb8d2008-03-23 00:55:46 +0000245 cycles = range(n, n-r, -1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000246 yield tuple(pool[i] for i in indices[:r])
247 while n:
248 for i in reversed(range(r)):
249 cycles[i] -= 1
250 if cycles[i] == 0:
251 indices[i:] = indices[i+1:] + indices[i:i+1]
252 cycles[i] = n - i
253 else:
254 j = cycles[i]
255 indices[i], indices[-j] = indices[-j], indices[i]
256 yield tuple(pool[i] for i in indices[:r])
257 break
258 else:
259 return
260
261 def permutations2(iterable, r=None):
262 'Pure python version shown in the docs'
263 pool = tuple(iterable)
264 n = len(pool)
265 r = n if r is None else r
266 for indices in product(range(n), repeat=r):
267 if len(set(indices)) == r:
268 yield tuple(pool[i] for i in indices)
269
270 for n in range(7):
271 values = [5*x-12 for x in range(n)]
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000272 for r in range(n+2):
Raymond Hettingerd553d852008-03-04 04:17:08 +0000273 result = list(permutations(values, r))
Ezio Melottidde5b942010-02-03 05:37:26 +0000274 self.assertEqual(len(result), 0 if r>n else fact(n) // fact(n-r)) # right number of perms
Raymond Hettingerd553d852008-03-04 04:17:08 +0000275 self.assertEqual(len(result), len(set(result))) # no repeats
276 self.assertEqual(result, sorted(result)) # lexicographic order
277 for p in result:
278 self.assertEqual(len(p), r) # r-length permutations
279 self.assertEqual(len(set(p)), r) # no duplicate elements
Benjamin Peterson5c8da862009-06-30 22:57:08 +0000280 self.assertTrue(all(e in values for e in p)) # elements taken from input iterable
Raymond Hettingerd553d852008-03-04 04:17:08 +0000281 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000282 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000283 if r == n:
284 self.assertEqual(result, list(permutations(values, None))) # test r as None
285 self.assertEqual(result, list(permutations(values))) # test default r
286
Benjamin Petersondda91212015-02-01 21:34:07 -0500287 @test_support.bigaddrspacetest
288 def test_permutations_overflow(self):
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +0200289 with self.assertRaises((OverflowError, MemoryError)):
Benjamin Petersondda91212015-02-01 21:34:07 -0500290 permutations("A", 2**30)
Benjamin Petersondda91212015-02-01 21:34:07 -0500291
Zachary Ware4e0df172014-04-24 13:20:27 -0500292 @test_support.impl_detail("tuple reuse is specific to CPython")
Alex Gaynor97376482011-07-17 16:44:31 -0700293 def test_permutations_tuple_reuse(self):
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000294 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000295 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000296
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000297 def test_combinatorics(self):
298 # Test relationships between product(), permutations(),
299 # combinations() and combinations_with_replacement().
300
Raymond Hettinger2f6c2e02009-01-27 10:36:14 +0000301 for n in range(6):
302 s = 'ABCDEFG'[:n]
303 for r in range(8):
304 prod = list(product(s, repeat=r))
305 cwr = list(combinations_with_replacement(s, r))
306 perm = list(permutations(s, r))
307 comb = list(combinations(s, r))
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000308
Raymond Hettinger2f6c2e02009-01-27 10:36:14 +0000309 # Check size
Ezio Melotti2623a372010-11-21 13:34:58 +0000310 self.assertEqual(len(prod), n**r)
311 self.assertEqual(len(cwr), (fact(n+r-1) // fact(r) // fact(n-1)) if n else (not r))
312 self.assertEqual(len(perm), 0 if r>n else fact(n) // fact(n-r))
313 self.assertEqual(len(comb), 0 if r>n else fact(n) // fact(r) // fact(n-r))
Raymond Hettinger2f6c2e02009-01-27 10:36:14 +0000314
315 # Check lexicographic order without repeated tuples
Ezio Melotti2623a372010-11-21 13:34:58 +0000316 self.assertEqual(prod, sorted(set(prod)))
317 self.assertEqual(cwr, sorted(set(cwr)))
318 self.assertEqual(perm, sorted(set(perm)))
319 self.assertEqual(comb, sorted(set(comb)))
Raymond Hettinger2f6c2e02009-01-27 10:36:14 +0000320
321 # Check interrelationships
Ezio Melotti2623a372010-11-21 13:34:58 +0000322 self.assertEqual(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
323 self.assertEqual(perm, [t for t in prod if len(set(t))==r]) # perm: prods with no dups
Raymond Hettinger2f6c2e02009-01-27 10:36:14 +0000324 self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
325 self.assertEqual(comb, [t for t in cwr if len(set(t))==r]) # comb: cwrs without dups
326 self.assertEqual(comb, filter(set(cwr).__contains__, perm)) # comb: perm that is a cwr
327 self.assertEqual(comb, filter(set(perm).__contains__, cwr)) # comb: cwr that is a perm
328 self.assertEqual(comb, sorted(set(cwr) & set(perm))) # comb: both a cwr and a perm
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000329
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000330 def test_compress(self):
Raymond Hettinger2e2909f2009-02-19 02:15:14 +0000331 self.assertEqual(list(compress(data='ABCDEF', selectors=[1,0,1,0,1,1])), list('ACEF'))
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000332 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
333 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
334 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
335 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
336 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
337 n = 10000
338 data = chain.from_iterable(repeat(range(6), n))
339 selectors = chain.from_iterable(repeat((0, 1)))
340 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
341 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
342 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
343 self.assertRaises(TypeError, compress, range(6)) # too few args
344 self.assertRaises(TypeError, compress, range(6), None) # too many args
345
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000346 def test_count(self):
347 self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
348 self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000349 self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000350 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
351 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000352 self.assertRaises(TypeError, count, 2, 3, 4)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000353 self.assertRaises(TypeError, count, 'a')
Serhiy Storchaka7f8ce852016-09-10 09:53:29 +0300354 self.assertEqual(take(10, count(maxsize-5)), range(maxsize-5, maxsize+5))
355 self.assertEqual(take(10, count(-maxsize-5)), range(-maxsize-5, -maxsize+5))
356 self.assertEqual(take(3, count(3.25)), [3.25, 4.25, 5.25])
357 self.assertEqual(take(3, count(3.25-4j)), [3.25-4j, 4.25-4j, 5.25-4j])
358 self.assertEqual(take(3, count(Decimal('1.1'))),
359 [Decimal('1.1'), Decimal('2.1'), Decimal('3.1')])
360 self.assertEqual(take(3, count(Fraction(2, 3))),
361 [Fraction(2, 3), Fraction(5, 3), Fraction(8, 3)])
362 BIGINT = 1<<1000
363 self.assertEqual(take(3, count(BIGINT)), [BIGINT, BIGINT+1, BIGINT+2])
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000364 c = count(3)
365 self.assertEqual(repr(c), 'count(3)')
366 c.next()
367 self.assertEqual(repr(c), 'count(4)')
Jack Diederich36234e82006-09-21 17:50:26 +0000368 c = count(-9)
369 self.assertEqual(repr(c), 'count(-9)')
370 c.next()
Serhiy Storchaka7f8ce852016-09-10 09:53:29 +0300371 self.assertEqual(next(c), -8)
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000372 self.assertEqual(repr(count(10.25)), 'count(10.25)')
Serhiy Storchaka7f8ce852016-09-10 09:53:29 +0300373 self.assertEqual(repr(count(10.0)), 'count(10.0)')
374 self.assertEqual(type(next(count(10.0))), float)
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000375 for i in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 10, sys.maxint-5, sys.maxint+5):
Raymond Hettinger3a0de082007-10-12 17:53:11 +0000376 # Test repr (ignoring the L in longs)
377 r1 = repr(count(i)).replace('L', '')
378 r2 = 'count(%r)'.__mod__(i).replace('L', '')
379 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000380
Raymond Hettingere09f45a2009-11-30 19:44:40 +0000381 # check copy, deepcopy, pickle
382 for value in -3, 3, sys.maxint-5, sys.maxint+5:
383 c = count(value)
384 self.assertEqual(next(copy.copy(c)), value)
385 self.assertEqual(next(copy.deepcopy(c)), value)
Serhiy Storchaka655720e2014-12-15 14:02:43 +0200386 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
387 self.assertEqual(next(pickle.loads(pickle.dumps(c, proto))), value)
Raymond Hettingere09f45a2009-11-30 19:44:40 +0000388
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000389 def test_count_with_stride(self):
390 self.assertEqual(zip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingera4038032009-02-14 00:25:51 +0000391 self.assertEqual(zip('abc',count(start=2,step=3)),
392 [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingeraa681c72009-02-21 07:17:22 +0000393 self.assertEqual(zip('abc',count(step=-1)),
394 [('a', 0), ('b', -1), ('c', -2)])
Serhiy Storchaka7f8ce852016-09-10 09:53:29 +0300395 self.assertRaises(TypeError, count, 'a', 'b')
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000396 self.assertEqual(zip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
397 self.assertEqual(zip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
Serhiy Storchaka7f8ce852016-09-10 09:53:29 +0300398 self.assertEqual(zip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
399 self.assertEqual(zip('abc',count(2,1L)), [('a', 2), ('b', 3), ('c', 4)])
400 self.assertEqual(zip('abc',count(2,3L)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000401 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
402 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
Serhiy Storchaka7f8ce852016-09-10 09:53:29 +0300403 self.assertEqual(take(3, count(10, maxsize+5)),
404 range(10, 10+3*(maxsize+5), maxsize+5))
405 self.assertEqual(take(3, count(2, 1.25)), [2, 3.25, 4.5])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000406 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettingeraa044612009-02-12 12:04:26 +0000407 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
408 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerdbe3bfb2009-02-12 12:43:01 +0000409 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
410 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Serhiy Storchaka7f8ce852016-09-10 09:53:29 +0300411 BIGINT = 1<<1000
412 self.assertEqual(take(3, count(step=BIGINT)), [0, BIGINT, 2*BIGINT])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000413 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
414 c = count(3, 5)
415 self.assertEqual(repr(c), 'count(3, 5)')
416 c.next()
417 self.assertEqual(repr(c), 'count(8, 5)')
418 c = count(-9, 0)
419 self.assertEqual(repr(c), 'count(-9, 0)')
420 c.next()
421 self.assertEqual(repr(c), 'count(-9, 0)')
422 c = count(-9, -3)
423 self.assertEqual(repr(c), 'count(-9, -3)')
424 c.next()
425 self.assertEqual(repr(c), 'count(-12, -3)')
426 self.assertEqual(repr(c), 'count(-12, -3)')
427 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
428 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
429 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
Serhiy Storchaka7f8ce852016-09-10 09:53:29 +0300430 self.assertEqual(repr(count(10, 1.00)), 'count(10, 1.0)')
431 c = count(10, 1.0)
432 self.assertEqual(type(next(c)), int)
433 self.assertEqual(type(next(c)), float)
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000434 for i in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 10, sys.maxint-5, sys.maxint+5):
435 for j in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 1, 10, sys.maxint-5, sys.maxint+5):
436 # Test repr (ignoring the L in longs)
437 r1 = repr(count(i, j)).replace('L', '')
438 if j == 1:
439 r2 = ('count(%r)' % i).replace('L', '')
440 else:
441 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
442 self.assertEqual(r1, r2)
443
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000444 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000445 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000446 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000447 self.assertRaises(TypeError, cycle)
448 self.assertRaises(TypeError, cycle, 5)
449 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000450
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000451 def test_groupby(self):
452 # Check whether it accepts arguments correctly
453 self.assertEqual([], list(groupby([])))
454 self.assertEqual([], list(groupby([], key=id)))
455 self.assertRaises(TypeError, list, groupby('abc', []))
456 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000457 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000458
459 # Check normal input
460 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
461 (2,15,22), (3,16,23), (3,17,23)]
462 dup = []
463 for k, g in groupby(s, lambda r:r[0]):
464 for elem in g:
465 self.assertEqual(k, elem[0])
466 dup.append(elem)
467 self.assertEqual(s, dup)
468
469 # Check nested case
470 dup = []
471 for k, g in groupby(s, lambda r:r[0]):
472 for ik, ig in groupby(g, lambda r:r[2]):
473 for elem in ig:
474 self.assertEqual(k, elem[0])
475 self.assertEqual(ik, elem[2])
476 dup.append(elem)
477 self.assertEqual(s, dup)
478
479 # Check case where inner iterator is not used
480 keys = [k for k, g in groupby(s, lambda r:r[0])]
481 expectedkeys = set([r[0] for r in s])
482 self.assertEqual(set(keys), expectedkeys)
483 self.assertEqual(len(keys), len(expectedkeys))
484
485 # Exercise pipes and filters style
486 s = 'abracadabra'
487 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000488 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000489 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
490 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000491 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000492 self.assertEqual(r, ['a', 'b', 'r'])
493 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000494 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000495 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
496 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000497 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000498 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
499
500 # iter.next failure
501 class ExpectedError(Exception):
502 pass
503 def delayed_raise(n=0):
504 for i in range(n):
505 yield 'yo'
506 raise ExpectedError
507 def gulp(iterable, keyp=None, func=list):
508 return [func(g) for k, g in groupby(iterable, keyp)]
509
510 # iter.next failure on outer object
511 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
512 # iter.next failure on inner object
513 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
514
515 # __cmp__ failure
516 class DummyCmp:
517 def __cmp__(self, dst):
518 raise ExpectedError
519 s = [DummyCmp(), DummyCmp(), None]
520
521 # __cmp__ failure on outer object
522 self.assertRaises(ExpectedError, gulp, s, func=id)
523 # __cmp__ failure on inner object
524 self.assertRaises(ExpectedError, gulp, s)
525
526 # keyfunc failure
527 def keyfunc(obj):
528 if keyfunc.skip > 0:
529 keyfunc.skip -= 1
530 return obj
531 else:
532 raise ExpectedError
533
534 # keyfunc failure on outer object
535 keyfunc.skip = 0
536 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
537 keyfunc.skip = 1
538 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
539
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000540 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000541 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000542 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000543 self.assertEqual(list(ifilter(bool, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000544 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000545 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000546 self.assertRaises(TypeError, ifilter, lambda x:x)
547 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000548 self.assertRaises(TypeError, ifilter, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000549 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000550
551 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000552 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
553 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000554 self.assertEqual(list(ifilterfalse(bool, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000555 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000556 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000557 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
558 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000559 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000560 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000561
562 def test_izip(self):
563 ans = [(x,y) for x, y in izip('abc',count())]
564 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000565 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
566 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
Raymond Hettinger02420702003-06-29 20:36:23 +0000567 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000568 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000569 self.assertEqual(list(izip()), zip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000570 self.assertRaises(TypeError, izip, 3)
571 self.assertRaises(TypeError, izip, range(3), 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000572 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
573 zip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000574 self.assertEqual([pair for pair in izip('abc', 'def')],
575 zip('abc', 'def'))
Alex Gaynor97376482011-07-17 16:44:31 -0700576
577 @test_support.impl_detail("tuple reuse is specific to CPython")
Zachary Ware4e0df172014-04-24 13:20:27 -0500578 def test_izip_tuple_reuse(self):
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000579 ids = map(id, izip('abc', 'def'))
580 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000581 ids = map(id, list(izip('abc', 'def')))
582 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000583
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000584 def test_iziplongest(self):
585 for args in [
586 ['abc', range(6)],
587 [range(6), 'abc'],
588 [range(1000), range(2000,2100), range(3000,3050)],
589 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
590 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
591 ]:
Ezio Melottidde5b942010-02-03 05:37:26 +0000592 # target = map(None, *args) <- this raises a py3k warning
593 # this is the replacement:
594 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
595 for i in range(max(map(len, args)))]
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000596 self.assertEqual(list(izip_longest(*args)), target)
597 self.assertEqual(list(izip_longest(*args, **{})), target)
598 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
599 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Tim Petersea5962f2007-03-12 18:07:52 +0000600
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000601 self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
602
603 self.assertEqual(list(izip_longest()), zip())
604 self.assertEqual(list(izip_longest([])), zip([]))
605 self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
Tim Petersea5962f2007-03-12 18:07:52 +0000606
Ezio Melottidde5b942010-02-03 05:37:26 +0000607 self.assertEqual(list(izip_longest('abc', 'defg', **{})),
608 zip(list('abc') + [None], 'defg')) # empty keyword dict
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000609 self.assertRaises(TypeError, izip_longest, 3)
610 self.assertRaises(TypeError, izip_longest, range(3), 3)
611
612 for stmt in [
613 "izip_longest('abc', fv=1)",
Tim Petersea5962f2007-03-12 18:07:52 +0000614 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000615 ]:
616 try:
617 eval(stmt, globals(), locals())
618 except TypeError:
619 pass
620 else:
621 self.fail('Did not raise Type in: ' + stmt)
Tim Petersea5962f2007-03-12 18:07:52 +0000622
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000623 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
624 zip('abc', 'def'))
625 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
626 zip('abc', 'def'))
Alex Gaynor97376482011-07-17 16:44:31 -0700627
628 @test_support.impl_detail("tuple reuse is specific to CPython")
629 def test_izip_longest_tuple_reuse(self):
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000630 ids = map(id, izip_longest('abc', 'def'))
631 self.assertEqual(min(ids), max(ids))
632 ids = map(id, list(izip_longest('abc', 'def')))
633 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
634
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +0000635 def test_bug_7244(self):
636
637 class Repeater(object):
638 # this class is similar to itertools.repeat
639 def __init__(self, o, t, e):
640 self.o = o
641 self.t = int(t)
642 self.e = e
643 def __iter__(self): # its iterator is itself
644 return self
645 def next(self):
646 if self.t > 0:
647 self.t -= 1
648 return self.o
649 else:
650 raise self.e
651
652 # Formerly this code in would fail in debug mode
653 # with Undetected Error and Stop Iteration
654 r1 = Repeater(1, 3, StopIteration)
655 r2 = Repeater(2, 4, StopIteration)
656 def run(r1, r2):
657 result = []
658 for i, j in izip_longest(r1, r2, fillvalue=0):
659 with test_support.captured_output('stdout'):
660 print (i, j)
661 result.append((i, j))
662 return result
663 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
664
665 # Formerly, the RuntimeError would be lost
666 # and StopIteration would stop as expected
667 r1 = Repeater(1, 3, RuntimeError)
668 r2 = Repeater(2, 4, StopIteration)
669 it = izip_longest(r1, r2, fillvalue=0)
670 self.assertEqual(next(it), (1, 2))
671 self.assertEqual(next(it), (1, 2))
672 self.assertEqual(next(it), (1, 2))
673 self.assertRaises(RuntimeError, next, it)
674
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000675 def test_product(self):
676 for args, result in [
Raymond Hettingerd553d852008-03-04 04:17:08 +0000677 ([], [()]), # zero iterables
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000678 (['ab'], [('a',), ('b',)]), # one iterable
679 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
680 ([range(0), range(2), range(3)], []), # first iterable with zero length
681 ([range(2), range(0), range(3)], []), # middle iterable with zero length
682 ([range(2), range(3), range(0)], []), # last iterable with zero length
683 ]:
684 self.assertEqual(list(product(*args)), result)
Raymond Hettinger08ff6822008-02-29 02:21:48 +0000685 for r in range(4):
686 self.assertEqual(list(product(*(args*r))),
687 list(product(*args, **dict(repeat=r))))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000688 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
689 self.assertRaises(TypeError, product, range(6), None)
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000690
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000691 def product1(*args, **kwds):
692 pools = map(tuple, args) * kwds.get('repeat', 1)
693 n = len(pools)
694 if n == 0:
695 yield ()
696 return
697 if any(len(pool) == 0 for pool in pools):
698 return
699 indices = [0] * n
700 yield tuple(pool[i] for pool, i in zip(pools, indices))
701 while 1:
702 for i in reversed(range(n)): # right to left
703 if indices[i] == len(pools[i]) - 1:
704 continue
705 indices[i] += 1
706 for j in range(i+1, n):
707 indices[j] = 0
708 yield tuple(pool[i] for pool, i in zip(pools, indices))
709 break
710 else:
711 return
712
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000713 def product2(*args, **kwds):
714 'Pure python version used in docs'
715 pools = map(tuple, args) * kwds.get('repeat', 1)
716 result = [[]]
717 for pool in pools:
718 result = [x+[y] for x in result for y in pool]
719 for prod in result:
720 yield tuple(prod)
721
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000722 argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
723 set('abcdefg'), range(11), tuple(range(13))]
724 for i in range(100):
725 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Raymond Hettingerd553d852008-03-04 04:17:08 +0000726 expected_len = prod(map(len, args))
727 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000728 self.assertEqual(list(product(*args)), list(product1(*args)))
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000729 self.assertEqual(list(product(*args)), list(product2(*args)))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000730 args = map(iter, args)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000731 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000732
Benjamin Petersondda91212015-02-01 21:34:07 -0500733 @test_support.bigaddrspacetest
734 def test_product_overflow(self):
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +0200735 with self.assertRaises((OverflowError, MemoryError)):
736 product(*(['ab']*2**5), repeat=2**25)
Benjamin Petersondda91212015-02-01 21:34:07 -0500737
Alex Gaynor97376482011-07-17 16:44:31 -0700738 @test_support.impl_detail("tuple reuse is specific to CPython")
739 def test_product_tuple_reuse(self):
Raymond Hettinger73d79632008-02-23 02:20:41 +0000740 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
741 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000742
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000743 def test_repeat(self):
Raymond Hettinger182edae2009-02-19 02:38:25 +0000744 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
Raymond Hettinger58ad2452014-06-24 21:53:45 -0700745 self.assertEqual(list(repeat(object='a', times=0)), [])
746 self.assertEqual(list(repeat(object='a', times=-1)), [])
747 self.assertEqual(list(repeat(object='a', times=-2)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000748 self.assertEqual(zip(xrange(3),repeat('a')),
749 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000750 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000751 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000752 self.assertEqual(list(repeat('a', 0)), [])
753 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000754 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000755 self.assertRaises(TypeError, repeat, None, 3, 4)
756 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000757 r = repeat(1+0j)
758 self.assertEqual(repr(r), 'repeat((1+0j))')
759 r = repeat(1+0j, 5)
760 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
761 list(r)
762 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000763
Raymond Hettinger58ad2452014-06-24 21:53:45 -0700764 def test_repeat_with_negative_times(self):
765 self.assertEqual(repr(repeat('a', -1)), "repeat('a', 0)")
766 self.assertEqual(repr(repeat('a', -2)), "repeat('a', 0)")
767 self.assertEqual(repr(repeat('a', times=-1)), "repeat('a', 0)")
768 self.assertEqual(repr(repeat('a', times=-2)), "repeat('a', 0)")
769
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000770 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000771 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
772 [0**1, 1**2, 2**3])
773 self.assertEqual(list(imap(None, 'abc', range(5))),
774 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000775 self.assertEqual(list(imap(None, 'abc', count())),
776 [('a',0),('b',1),('c',2)])
777 self.assertEqual(take(2,imap(None, 'abc', count())),
778 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000779 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000780 self.assertRaises(TypeError, imap)
781 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000782 self.assertRaises(TypeError, imap(10, range(5)).next)
783 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
784 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000785
786 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000787 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
788 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000789 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
790 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000791 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger47317092008-01-17 03:02:14 +0000792 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
793 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000794 self.assertRaises(TypeError, starmap)
795 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
796 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
797 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
798 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000799
800 def test_islice(self):
801 for args in [ # islice(args) should agree with range(args)
802 (10, 20, 3),
803 (10, 3, 20),
804 (10, 20),
805 (10, 3),
806 (20,)
807 ]:
808 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
809
810 for args, tgtargs in [ # Stop when seqn is exhausted
811 ((10, 110, 3), ((10, 100, 3))),
812 ((10, 110), ((10, 100))),
813 ((110,), (100,))
814 ]:
815 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
816
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000817 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000818 self.assertEqual(list(islice(xrange(10), None)), range(10))
Raymond Hettingerb2594052004-12-05 09:25:51 +0000819 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
820 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000821 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
822 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
823
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000824 # Test number of items consumed SF #1171417
825 it = iter(range(10))
826 self.assertEqual(list(islice(it, 3)), range(3))
827 self.assertEqual(list(it), range(3, 10))
828
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000829 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000830 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000831 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
832 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
833 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
834 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
835 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000836 self.assertRaises(ValueError, islice, xrange(10), 'a')
837 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
838 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
839 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
840 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +0000841 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000842
Raymond Hettinger061bf7a2010-11-30 03:15:35 +0000843 # Issue #10323: Less islice in a predictable state
844 c = count()
845 self.assertEqual(list(islice(c, 1, 3, 50)), [1])
846 self.assertEqual(next(c), 3)
847
Antoine Pitrou3ec903f2014-04-29 12:13:46 +0200848 # Issue #21321: check source iterator is not referenced
849 # from islice() after the latter has been exhausted
850 it = (x for x in (1, 2))
851 wr = weakref.ref(it)
852 it = islice(it, 1)
853 self.assertIsNotNone(wr())
854 list(it) # exhaust the iterator
Benjamin Petersonec9d5472014-08-24 18:07:28 -0500855 test_support.gc_collect()
Antoine Pitrou3ec903f2014-04-29 12:13:46 +0200856 self.assertIsNone(wr())
857
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000858 def test_takewhile(self):
859 data = [1, 3, 5, 20, 2, 4, 6, 8]
860 underten = lambda x: x<10
861 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000862 self.assertEqual(list(takewhile(underten, [])), [])
863 self.assertRaises(TypeError, takewhile)
864 self.assertRaises(TypeError, takewhile, operator.pow)
865 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
866 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
867 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000868 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
869 self.assertEqual(list(t), [1, 1, 1])
870 self.assertRaises(StopIteration, t.next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000871
872 def test_dropwhile(self):
873 data = [1, 3, 5, 20, 2, 4, 6, 8]
874 underten = lambda x: x<10
875 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000876 self.assertEqual(list(dropwhile(underten, [])), [])
877 self.assertRaises(TypeError, dropwhile)
878 self.assertRaises(TypeError, dropwhile, operator.pow)
879 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
880 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
881 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000882
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000883 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000884 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000885 def irange(n):
886 for i in xrange(n):
887 yield i
888
889 a, b = tee([]) # test empty iterator
890 self.assertEqual(list(a), [])
891 self.assertEqual(list(b), [])
892
893 a, b = tee(irange(n)) # test 100% interleaved
894 self.assertEqual(zip(a,b), zip(range(n),range(n)))
895
896 a, b = tee(irange(n)) # test 0% interleaved
897 self.assertEqual(list(a), range(n))
898 self.assertEqual(list(b), range(n))
899
900 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000901 for i in xrange(100):
902 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000903 del a
904 self.assertEqual(list(b), range(n))
905
906 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000907 for i in xrange(100):
908 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000909 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000910 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000911
912 for j in xrange(5): # test randomly interleaved
913 order = [0]*n + [1]*n
914 random.shuffle(order)
915 lists = ([], [])
916 its = tee(irange(n))
917 for i in order:
918 value = its[i].next()
919 lists[i].append(value)
920 self.assertEqual(lists[0], range(n))
921 self.assertEqual(lists[1], range(n))
922
Raymond Hettingerad983e72003-11-12 14:32:26 +0000923 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000924 self.assertRaises(TypeError, tee)
925 self.assertRaises(TypeError, tee, 3)
926 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000927 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000928
Raymond Hettingerad983e72003-11-12 14:32:26 +0000929 # tee object should be instantiable
930 a, b = tee('abc')
931 c = type(a)('def')
932 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000933
Raymond Hettingerad983e72003-11-12 14:32:26 +0000934 # test long-lagged and multi-way split
935 a, b, c = tee(xrange(2000), 3)
936 for i in xrange(100):
937 self.assertEqual(a.next(), i)
938 self.assertEqual(list(b), range(2000))
939 self.assertEqual([c.next(), c.next()], range(2))
940 self.assertEqual(list(a), range(100,2000))
941 self.assertEqual(list(c), range(2,2000))
942
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000943 # test values of n
944 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Neal Norwitz69e88972006-09-02 02:58:13 +0000945 self.assertRaises(ValueError, tee, [], -1)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000946 for n in xrange(5):
947 result = tee('abc', n)
948 self.assertEqual(type(result), tuple)
949 self.assertEqual(len(result), n)
950 self.assertEqual(map(list, result), [list('abc')]*n)
951
Raymond Hettingerad983e72003-11-12 14:32:26 +0000952 # tee pass-through to copyable iterator
953 a, b = tee('abc')
954 c, d = tee(a)
Benjamin Peterson5c8da862009-06-30 22:57:08 +0000955 self.assertTrue(a is c)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000956
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000957 # test tee_new
958 t1, t2 = tee('abc')
959 tnew = type(t1)
960 self.assertRaises(TypeError, tnew)
961 self.assertRaises(TypeError, tnew, 10)
962 t3 = tnew(t1)
Benjamin Peterson5c8da862009-06-30 22:57:08 +0000963 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000964
Raymond Hettingera9f60922004-10-17 16:40:14 +0000965 # test that tee objects are weak referencable
966 a, b = tee(xrange(10))
Antoine Pitrou3ec903f2014-04-29 12:13:46 +0200967 p = weakref.proxy(a)
Raymond Hettingera9f60922004-10-17 16:40:14 +0000968 self.assertEqual(getattr(p, '__class__'), type(b))
969 del a
970 self.assertRaises(ReferenceError, getattr, p, '__class__')
971
Serhiy Storchakab09ec9b2013-01-25 13:31:05 +0200972 # Issue 13454: Crash when deleting backward iterator from tee()
973 def test_tee_del_backward(self):
Serhiy Storchaka6fef14d2013-01-26 11:51:42 +0200974 forward, backward = tee(repeat(None, 20000000))
Serhiy Storchaka53ea1622015-03-28 20:38:48 +0200975 try:
976 any(forward) # exhaust the iterator
977 del backward
978 except:
979 del forward, backward
980 raise
Serhiy Storchakab09ec9b2013-01-25 13:31:05 +0200981
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000982 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000983 self.assertRaises(StopIteration, izip().next)
984
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000985 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000986 self.assertRaises(StopIteration, f([]).next)
987 self.assertRaises(StopIteration, f(StopNow()).next)
988
989 self.assertRaises(StopIteration, islice([], None).next)
990 self.assertRaises(StopIteration, islice(StopNow(), None).next)
991
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000992 p, q = tee([])
993 self.assertRaises(StopIteration, p.next)
994 self.assertRaises(StopIteration, q.next)
995 p, q = tee(StopNow())
996 self.assertRaises(StopIteration, p.next)
997 self.assertRaises(StopIteration, q.next)
998
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000999 self.assertRaises(StopIteration, repeat(None, 0).next)
1000
1001 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
1002 self.assertRaises(StopIteration, f(lambda x:x, []).next)
1003 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
1004
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001005class TestExamples(unittest.TestCase):
1006
1007 def test_chain(self):
1008 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
1009
1010 def test_chain_from_iterable(self):
1011 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
1012
1013 def test_combinations(self):
1014 self.assertEqual(list(combinations('ABCD', 2)),
1015 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
1016 self.assertEqual(list(combinations(range(4), 3)),
1017 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
1018
Raymond Hettingerd081abc2009-01-27 02:58:49 +00001019 def test_combinations_with_replacement(self):
1020 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
1021 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
1022
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001023 def test_compress(self):
1024 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
1025
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001026 def test_count(self):
1027 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
1028
1029 def test_cycle(self):
1030 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
1031
1032 def test_dropwhile(self):
1033 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
1034
1035 def test_groupby(self):
1036 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
1037 list('ABCDAB'))
1038 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
1039 [list('AAAA'), list('BBB'), list('CC'), list('D')])
1040
1041 def test_ifilter(self):
1042 self.assertEqual(list(ifilter(lambda x: x%2, range(10))), [1,3,5,7,9])
1043
1044 def test_ifilterfalse(self):
1045 self.assertEqual(list(ifilterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
1046
1047 def test_imap(self):
1048 self.assertEqual(list(imap(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
1049
1050 def test_islice(self):
1051 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
1052 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
1053 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
1054 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1055
1056 def test_izip(self):
1057 self.assertEqual(list(izip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
1058
1059 def test_izip_longest(self):
1060 self.assertEqual(list(izip_longest('ABCD', 'xy', fillvalue='-')),
1061 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
1062
1063 def test_permutations(self):
1064 self.assertEqual(list(permutations('ABCD', 2)),
1065 map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split()))
1066 self.assertEqual(list(permutations(range(3))),
1067 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
1068
1069 def test_product(self):
1070 self.assertEqual(list(product('ABCD', 'xy')),
1071 map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split()))
1072 self.assertEqual(list(product(range(2), repeat=3)),
1073 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
1074 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
1075
1076 def test_repeat(self):
1077 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
1078
1079 def test_stapmap(self):
1080 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
1081 [32, 9, 1000])
1082
1083 def test_takewhile(self):
1084 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
1085
1086
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001087class TestGC(unittest.TestCase):
1088
1089 def makecycle(self, iterator, container):
1090 container.append(iterator)
1091 iterator.next()
1092 del container, iterator
1093
1094 def test_chain(self):
1095 a = []
1096 self.makecycle(chain(a), a)
1097
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001098 def test_chain_from_iterable(self):
1099 a = []
1100 self.makecycle(chain.from_iterable([a]), a)
1101
1102 def test_combinations(self):
1103 a = []
1104 self.makecycle(combinations([1,2,a,3], 3), a)
1105
Raymond Hettingerd081abc2009-01-27 02:58:49 +00001106 def test_combinations_with_replacement(self):
1107 a = []
1108 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1109
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001110 def test_compress(self):
1111 a = []
1112 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1113
Raymond Hettingerb21d8102009-02-16 20:39:12 +00001114 def test_count(self):
1115 a = []
1116 Int = type('Int', (int,), dict(x=a))
1117 self.makecycle(count(Int(0), Int(1)), a)
1118
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001119 def test_cycle(self):
1120 a = []
1121 self.makecycle(cycle([a]*2), a)
1122
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001123 def test_dropwhile(self):
1124 a = []
1125 self.makecycle(dropwhile(bool, [0, a, a]), a)
1126
1127 def test_groupby(self):
1128 a = []
1129 self.makecycle(groupby([a]*2, lambda x:x), a)
1130
Raymond Hettingera1ca94a2008-03-06 22:51:36 +00001131 def test_issue2246(self):
1132 # Issue 2246 -- the _grouper iterator was not included in GC
1133 n = 10
1134 keyfunc = lambda x: x
1135 for i, j in groupby(xrange(n), key=keyfunc):
1136 keyfunc.__dict__.setdefault('x',[]).append(j)
1137
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001138 def test_ifilter(self):
1139 a = []
1140 self.makecycle(ifilter(lambda x:True, [a]*2), a)
1141
1142 def test_ifilterfalse(self):
1143 a = []
1144 self.makecycle(ifilterfalse(lambda x:False, a), a)
1145
1146 def test_izip(self):
1147 a = []
1148 self.makecycle(izip([a]*2, [a]*3), a)
1149
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001150 def test_izip_longest(self):
1151 a = []
1152 self.makecycle(izip_longest([a]*2, [a]*3), a)
1153 b = [a, None]
1154 self.makecycle(izip_longest([a]*2, [a]*3, fillvalue=b), a)
1155
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001156 def test_imap(self):
1157 a = []
1158 self.makecycle(imap(lambda x:x, [a]*2), a)
1159
1160 def test_islice(self):
1161 a = []
1162 self.makecycle(islice([a]*2, None), a)
1163
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001164 def test_permutations(self):
1165 a = []
1166 self.makecycle(permutations([1,2,a,3], 3), a)
1167
1168 def test_product(self):
1169 a = []
1170 self.makecycle(product([1,2,a,3], repeat=3), a)
1171
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001172 def test_repeat(self):
1173 a = []
1174 self.makecycle(repeat(a), a)
1175
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001176 def test_starmap(self):
1177 a = []
1178 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1179
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001180 def test_takewhile(self):
1181 a = []
1182 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1183
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001184def R(seqn):
1185 'Regular generator'
1186 for i in seqn:
1187 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001188
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001189class G:
1190 'Sequence using __getitem__'
1191 def __init__(self, seqn):
1192 self.seqn = seqn
1193 def __getitem__(self, i):
1194 return self.seqn[i]
1195
1196class I:
1197 'Sequence using iterator protocol'
1198 def __init__(self, seqn):
1199 self.seqn = seqn
1200 self.i = 0
1201 def __iter__(self):
1202 return self
1203 def next(self):
1204 if self.i >= len(self.seqn): raise StopIteration
1205 v = self.seqn[self.i]
1206 self.i += 1
1207 return v
1208
1209class Ig:
1210 'Sequence using iterator protocol defined with a generator'
1211 def __init__(self, seqn):
1212 self.seqn = seqn
1213 self.i = 0
1214 def __iter__(self):
1215 for val in self.seqn:
1216 yield val
1217
1218class X:
1219 'Missing __getitem__ and __iter__'
1220 def __init__(self, seqn):
1221 self.seqn = seqn
1222 self.i = 0
1223 def next(self):
1224 if self.i >= len(self.seqn): raise StopIteration
1225 v = self.seqn[self.i]
1226 self.i += 1
1227 return v
1228
1229class N:
1230 'Iterator missing next()'
1231 def __init__(self, seqn):
1232 self.seqn = seqn
1233 self.i = 0
1234 def __iter__(self):
1235 return self
1236
1237class E:
1238 'Test propagation of exceptions'
1239 def __init__(self, seqn):
1240 self.seqn = seqn
1241 self.i = 0
1242 def __iter__(self):
1243 return self
1244 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001245 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001246
1247class S:
1248 'Test immediate stop'
1249 def __init__(self, seqn):
1250 pass
1251 def __iter__(self):
1252 return self
1253 def next(self):
1254 raise StopIteration
1255
1256def L(seqn):
1257 'Test multiple tiers of iterators'
1258 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
1259
1260
1261class TestVariousIteratorArgs(unittest.TestCase):
1262
1263 def test_chain(self):
1264 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1265 for g in (G, I, Ig, S, L, R):
1266 self.assertEqual(list(chain(g(s))), list(g(s)))
1267 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001268 self.assertRaises(TypeError, list, chain(X(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001269 self.assertRaises(TypeError, list, chain(N(s)))
1270 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1271
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001272 def test_compress(self):
1273 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1274 n = len(s)
1275 for g in (G, I, Ig, S, L, R):
1276 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1277 self.assertRaises(TypeError, compress, X(s), repeat(1))
1278 self.assertRaises(TypeError, list, compress(N(s), repeat(1)))
1279 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1280
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001281 def test_product(self):
1282 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1283 self.assertRaises(TypeError, product, X(s))
1284 self.assertRaises(TypeError, product, N(s))
1285 self.assertRaises(ZeroDivisionError, product, E(s))
1286
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001287 def test_cycle(self):
1288 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1289 for g in (G, I, Ig, S, L, R):
1290 tgtlen = len(s) * 3
1291 expected = list(g(s))*3
1292 actual = list(islice(cycle(g(s)), tgtlen))
1293 self.assertEqual(actual, expected)
1294 self.assertRaises(TypeError, cycle, X(s))
1295 self.assertRaises(TypeError, list, cycle(N(s)))
1296 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1297
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001298 def test_groupby(self):
1299 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1300 for g in (G, I, Ig, S, L, R):
1301 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1302 self.assertRaises(TypeError, groupby, X(s))
1303 self.assertRaises(TypeError, list, groupby(N(s)))
1304 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1305
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001306 def test_ifilter(self):
1307 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1308 for g in (G, I, Ig, S, L, R):
1309 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
1310 self.assertRaises(TypeError, ifilter, isEven, X(s))
1311 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
1312 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
1313
1314 def test_ifilterfalse(self):
1315 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1316 for g in (G, I, Ig, S, L, R):
1317 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
1318 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
1319 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
1320 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
1321
1322 def test_izip(self):
1323 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1324 for g in (G, I, Ig, S, L, R):
1325 self.assertEqual(list(izip(g(s))), zip(g(s)))
1326 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
1327 self.assertRaises(TypeError, izip, X(s))
1328 self.assertRaises(TypeError, list, izip(N(s)))
1329 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
1330
Raymond Hettingerd36862c2007-02-21 05:20:38 +00001331 def test_iziplongest(self):
1332 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1333 for g in (G, I, Ig, S, L, R):
1334 self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
1335 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
1336 self.assertRaises(TypeError, izip_longest, X(s))
1337 self.assertRaises(TypeError, list, izip_longest(N(s)))
1338 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
1339
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001340 def test_imap(self):
1341 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1342 for g in (G, I, Ig, S, L, R):
1343 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
1344 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
1345 self.assertRaises(TypeError, imap, onearg, X(s))
1346 self.assertRaises(TypeError, list, imap(onearg, N(s)))
1347 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
1348
1349 def test_islice(self):
1350 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1351 for g in (G, I, Ig, S, L, R):
1352 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1353 self.assertRaises(TypeError, islice, X(s), 10)
1354 self.assertRaises(TypeError, list, islice(N(s), 10))
1355 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1356
1357 def test_starmap(self):
1358 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1359 for g in (G, I, Ig, S, L, R):
1360 ss = zip(s, s)
1361 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
1362 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
1363 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
1364 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1365
1366 def test_takewhile(self):
1367 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1368 for g in (G, I, Ig, S, L, R):
1369 tgt = []
1370 for elem in g(s):
1371 if not isEven(elem): break
1372 tgt.append(elem)
1373 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1374 self.assertRaises(TypeError, takewhile, isEven, X(s))
1375 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
1376 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1377
1378 def test_dropwhile(self):
1379 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1380 for g in (G, I, Ig, S, L, R):
1381 tgt = []
1382 for elem in g(s):
1383 if not tgt and isOdd(elem): continue
1384 tgt.append(elem)
1385 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1386 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
1387 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
1388 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1389
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001390 def test_tee(self):
1391 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1392 for g in (G, I, Ig, S, L, R):
1393 it1, it2 = tee(g(s))
1394 self.assertEqual(list(it1), list(g(s)))
1395 self.assertEqual(list(it2), list(g(s)))
1396 self.assertRaises(TypeError, tee, X(s))
1397 self.assertRaises(TypeError, list, tee(N(s))[0])
1398 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1399
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001400class LengthTransparency(unittest.TestCase):
1401
1402 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001403 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001404 self.assertEqual(len(repeat(None, 50)), 50)
1405 self.assertRaises(TypeError, len, repeat(None))
1406
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001407class RegressionTests(unittest.TestCase):
1408
1409 def test_sf_793826(self):
1410 # Fix Armin Rigo's successful efforts to wreak havoc
1411
1412 def mutatingtuple(tuple1, f, tuple2):
1413 # this builds a tuple t which is a copy of tuple1,
1414 # then calls f(t), then mutates t to be equal to tuple2
1415 # (needs len(tuple1) == len(tuple2)).
1416 def g(value, first=[1]):
1417 if first:
1418 del first[:]
1419 f(z.next())
1420 return value
1421 items = list(tuple2)
1422 items[1:1] = list(tuple1)
1423 gen = imap(g, items)
1424 z = izip(*[gen]*len(tuple1))
1425 z.next()
1426
1427 def f(t):
1428 global T
1429 T = t
1430 first[:] = list(T)
1431
1432 first = []
1433 mutatingtuple((1,2,3), f, (4,5,6))
1434 second = list(T)
1435 self.assertEqual(first, second)
1436
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001437
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001438 def test_sf_950057(self):
1439 # Make sure that chain() and cycle() catch exceptions immediately
1440 # rather than when shifting between input sources
1441
1442 def gen1():
1443 hist.append(0)
1444 yield 1
1445 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001446 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001447 hist.append(2)
1448
1449 def gen2(x):
1450 hist.append(3)
1451 yield 2
1452 hist.append(4)
1453 if x:
1454 raise StopIteration
1455
1456 hist = []
1457 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1458 self.assertEqual(hist, [0,1])
1459
1460 hist = []
1461 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1462 self.assertEqual(hist, [0,1])
1463
1464 hist = []
1465 self.assertRaises(AssertionError, list, cycle(gen1()))
1466 self.assertEqual(hist, [0,1])
1467
Georg Brandlb84c1372007-01-21 10:28:43 +00001468class SubclassWithKwargsTest(unittest.TestCase):
1469 def test_keywords_in_subclass(self):
1470 # count is not subclassable...
1471 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001472 starmap, islice, takewhile, dropwhile, cycle, compress):
Georg Brandlb84c1372007-01-21 10:28:43 +00001473 class Subclass(cls):
1474 def __init__(self, newarg=None, *args):
1475 cls.__init__(self, *args)
1476 try:
1477 Subclass(newarg=1)
1478 except TypeError, err:
1479 # we expect type errors because of wrong argument count
Ezio Melottiaa980582010-01-23 23:04:36 +00001480 self.assertNotIn("does not take keyword arguments", err.args[0])
Georg Brandlb84c1372007-01-21 10:28:43 +00001481
1482
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001483libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001484
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001485
1486>>> amounts = [120.15, 764.05, 823.14]
1487>>> for checknum, amount in izip(count(1200), amounts):
1488... print 'Check %d is for $%.2f' % (checknum, amount)
1489...
1490Check 1200 is for $120.15
1491Check 1201 is for $764.05
1492Check 1202 is for $823.14
1493
1494>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001495>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1496... print cube
1497...
14981
14998
150027
1501
1502>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001503>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001504... print name.title()
1505...
1506Alex
1507Laura
1508Martin
1509Walter
1510Samuele
1511
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001512>>> from operator import itemgetter
1513>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Armin Rigoa3f09272006-05-28 19:13:17 +00001514>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001515>>> for k, g in groupby(di, itemgetter(1)):
1516... print k, map(itemgetter(0), g)
1517...
15181 ['a', 'c', 'e']
15192 ['b', 'd', 'f']
15203 ['g']
1521
Raymond Hettinger734fb572004-01-20 20:04:40 +00001522# Find runs of consecutive numbers using groupby. The key to the solution
1523# is differencing with a range so that consecutive numbers all appear in
1524# same group.
1525>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Ezio Melottidde5b942010-02-03 05:37:26 +00001526>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Raymond Hettinger734fb572004-01-20 20:04:40 +00001527... print map(operator.itemgetter(1), g)
1528...
1529[1]
1530[4, 5, 6]
1531[10]
1532[15, 16, 17, 18]
1533[22]
1534[25, 26, 27, 28]
1535
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001536>>> def take(n, iterable):
1537... "Return first n items of the iterable as a list"
1538... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001539
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001540>>> def enumerate(iterable, start=0):
1541... return izip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001542
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001543>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001544... "Return function(0), function(1), ..."
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001545... return imap(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001546
Raymond Hettingerf9bce832009-02-19 05:34:35 +00001547>>> def nth(iterable, n, default=None):
1548... "Returns the nth item or a default value"
1549... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001550
Raymond Hettinger2ec5fd52016-03-06 18:06:29 -08001551>>> def all_equal(iterable):
1552... "Returns True if all the elements are equal to each other"
1553... g = groupby(iterable)
1554... return next(g, True) and not next(g, False)
1555
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001556>>> def quantify(iterable, pred=bool):
1557... "Count how many times the predicate is true"
1558... return sum(imap(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001559
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001560>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001561... "Returns the sequence elements and then returns None indefinitely"
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001562... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001563
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001564>>> def ncycles(iterable, n):
Ezio Melottic2077b02011-03-16 12:34:31 +02001565... "Returns the sequence elements n times"
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001566... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001567
1568>>> def dotproduct(vec1, vec2):
1569... return sum(imap(operator.mul, vec1, vec2))
1570
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001571>>> def flatten(listOfLists):
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001572... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001573
1574>>> def repeatfunc(func, times=None, *args):
1575... "Repeat calls to func with specified arguments."
1576... " Example: repeatfunc(random.random)"
1577... if times is None:
1578... return starmap(func, repeat(args))
1579... else:
1580... return starmap(func, repeat(args, times))
1581
Raymond Hettingerd591f662003-10-26 15:34:50 +00001582>>> def pairwise(iterable):
1583... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1584... a, b = tee(iterable)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001585... for elem in b:
1586... break
Raymond Hettingerad983e72003-11-12 14:32:26 +00001587... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001588
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001589>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001590... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001591... args = [iter(iterable)] * n
Raymond Hettingerf080e6d2008-07-31 01:19:50 +00001592... return izip_longest(fillvalue=fillvalue, *args)
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001593
1594>>> def roundrobin(*iterables):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001595... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001596... # Recipe credited to George Sakkis
1597... pending = len(iterables)
1598... nexts = cycle(iter(it).next for it in iterables)
1599... while pending:
1600... try:
1601... for next in nexts:
1602... yield next()
1603... except StopIteration:
1604... pending -= 1
1605... nexts = cycle(islice(nexts, pending))
1606
1607>>> def powerset(iterable):
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001608... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1609... s = list(iterable)
1610... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001611
Raymond Hettinger44e15812009-01-02 21:26:45 +00001612>>> def unique_everseen(iterable, key=None):
1613... "List unique elements, preserving order. Remember all elements ever seen."
1614... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1615... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1616... seen = set()
1617... seen_add = seen.add
1618... if key is None:
1619... for element in iterable:
1620... if element not in seen:
1621... seen_add(element)
1622... yield element
1623... else:
1624... for element in iterable:
1625... k = key(element)
1626... if k not in seen:
1627... seen_add(k)
1628... yield element
1629
1630>>> def unique_justseen(iterable, key=None):
1631... "List unique elements, preserving order. Remember only the element just seen."
1632... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1633... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1634... return imap(next, imap(itemgetter(1), groupby(iterable, key)))
1635
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001636This is not part of the examples but it tests to make sure the definitions
1637perform as purported.
1638
Raymond Hettingera098b332003-09-08 23:58:40 +00001639>>> take(10, count())
1640[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1641
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001642>>> list(enumerate('abc'))
1643[(0, 'a'), (1, 'b'), (2, 'c')]
1644
1645>>> list(islice(tabulate(lambda x: 2*x), 4))
1646[0, 2, 4, 6]
1647
1648>>> nth('abcde', 3)
Raymond Hettingerd507afd2009-02-04 10:52:32 +00001649'd'
1650
1651>>> nth('abcde', 9) is None
1652True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001653
Raymond Hettinger2ec5fd52016-03-06 18:06:29 -08001654>>> [all_equal(s) for s in ('', 'A', 'AAAA', 'AAAB', 'AAABA')]
1655[True, True, True, False, False]
1656
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001657>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000165850
1659
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001660>>> a = [[1, 2, 3], [4, 5, 6]]
1661>>> flatten(a)
1662[1, 2, 3, 4, 5, 6]
1663
1664>>> list(repeatfunc(pow, 5, 2, 3))
1665[8, 8, 8, 8, 8]
1666
1667>>> import random
1668>>> take(5, imap(int, repeatfunc(random.random)))
1669[0, 0, 0, 0, 0]
1670
Raymond Hettingerd591f662003-10-26 15:34:50 +00001671>>> list(pairwise('abcd'))
1672[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001673
Raymond Hettingerd591f662003-10-26 15:34:50 +00001674>>> list(pairwise([]))
1675[]
1676
1677>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001678[]
1679
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001680>>> list(islice(padnone('abc'), 0, 6))
1681['a', 'b', 'c', None, None, None]
1682
1683>>> list(ncycles('abc', 3))
1684['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1685
1686>>> dotproduct([1,2,3], [4,5,6])
168732
1688
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001689>>> list(grouper(3, 'abcdefg', 'x'))
1690[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1691
1692>>> list(roundrobin('abc', 'd', 'ef'))
1693['a', 'd', 'e', 'b', 'f', 'c']
1694
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001695>>> list(powerset([1,2,3]))
1696[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001697
Raymond Hettinger560f9a82009-01-27 13:26:35 +00001698>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1699True
1700
1701>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1702True
1703
Raymond Hettinger44e15812009-01-02 21:26:45 +00001704>>> list(unique_everseen('AAAABBBCCDAABBB'))
1705['A', 'B', 'C', 'D']
1706
1707>>> list(unique_everseen('ABBCcAD', str.lower))
1708['A', 'B', 'C', 'D']
1709
1710>>> list(unique_justseen('AAAABBBCCDAABBB'))
1711['A', 'B', 'C', 'D', 'A', 'B']
1712
1713>>> list(unique_justseen('ABBCcAD', str.lower))
1714['A', 'B', 'C', 'A', 'D']
1715
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001716"""
1717
1718__test__ = {'libreftest' : libreftest}
1719
1720def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001721 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Georg Brandlb84c1372007-01-21 10:28:43 +00001722 RegressionTests, LengthTransparency,
Serhiy Storchaka6c467a42013-04-06 22:51:29 +03001723 SubclassWithKwargsTest, TestExamples)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001724 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001725
1726 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001727 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001728 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001729 counts = [None] * 5
1730 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001731 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001732 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001733 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001734 print counts
1735
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001736 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001737 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001738
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001739if __name__ == "__main__":
1740 test_main(verbose=True)