blob: bfe13ab8ed2d04d3477d34671ce4417bfd410b6d [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001import unittest
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002from test import support
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003from itertools import *
Raymond Hettingera9f60922004-10-17 16:40:14 +00004from weakref import proxy
Raymond Hettinger8a882a72009-02-12 12:08:18 +00005from decimal import Decimal
Raymond Hettinger2012f172003-02-07 05:32:58 +00006import sys
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00007import operator
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00008import random
Christian Heimesc3f30c42008-02-22 16:37:40 +00009from functools import reduce
Benjamin Petersonee8712c2008-05-20 21:35:26 +000010maxsize = support.MAX_Py_ssize_t
Guido van Rossum360e4b82007-05-14 22:51:27 +000011minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000012
Guido van Rossum801f0d72006-08-24 19:48:10 +000013def lzip(*args):
14 return list(zip(*args))
15
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000016def 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
Georg Brandla18af4e2007-04-21 15:47:16 +000041 def __next__(self):
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000042 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
Christian Heimes78644762008-03-04 23:39:23 +000048def prod(iterable):
49 return reduce(operator.mul, iterable, 1)
50
Christian Heimes380f7f22008-02-28 11:19:05 +000051def fact(n):
52 'Factorial'
Christian Heimes78644762008-03-04 23:39:23 +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):
Christian Heimesdd15f6c2008-03-16 00:07:10 +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))
Christian Heimesf16baeb2008-02-29 14:57:44 +000070
71 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]))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000077
Christian Heimes380f7f22008-02-28 11:19:05 +000078 def test_combinations(self):
Raymond Hettinger5bad41e2009-01-08 21:01:54 +000079 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Christian Heimes380f7f22008-02-28 11:19:05 +000080 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Christian Heimes78644762008-03-04 23:39:23 +000081 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Christian Heimes380f7f22008-02-28 11:19:05 +000082 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Raymond Hettinger5bad41e2009-01-08 21:01:54 +000083 self.assertEqual(list(combinations('abc', 32)), []) # r > n
Christian Heimes380f7f22008-02-28 11:19:05 +000084 self.assertEqual(list(combinations(range(4), 3)),
85 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
Christian Heimes78644762008-03-04 23:39:23 +000086
87 def combinations1(iterable, r):
88 'Pure python version shown in the docs'
89 pool = tuple(iterable)
90 n = len(pool)
Raymond Hettinger5bad41e2009-01-08 21:01:54 +000091 if r > n:
92 return
Christian Heimes78644762008-03-04 23:39:23 +000093 indices = list(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 Hettingereb508ad2009-01-27 09:35:21 +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
Christian Heimes78644762008-03-04 23:39:23 +0000122 for n in range(7):
Christian Heimes380f7f22008-02-28 11:19:05 +0000123 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000124 for r in range(n+2):
Christian Heimes380f7f22008-02-28 11:19:05 +0000125 result = list(combinations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000126 self.assertEqual(len(result), 0 if r>n else fact(n) / fact(r) / fact(n-r)) # right number of combs
Christian Heimes380f7f22008-02-28 11:19:05 +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
133 self.assert_(all(e in values for e in c)) # elements taken from input iterable
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000134 self.assertEqual(list(c),
135 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000136 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000137 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000138 self.assertEqual(result, list(combinations3(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000139
140 # Test implementation detail: tuple re-use
141 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
142 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
143
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000144 def test_combinations_with_replacement(self):
145 cwr = combinations_with_replacement
146 self.assertRaises(TypeError, cwr, 'abc') # missing r argument
147 self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
148 self.assertRaises(TypeError, cwr, None) # pool is not iterable
149 self.assertRaises(ValueError, cwr, 'abc', -2) # r is negative
150 self.assertEqual(list(cwr('ABC', 2)),
151 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
152
153 def cwr1(iterable, r):
154 'Pure python version shown in the docs'
155 # number items returned: (n+r-1)! / r! / (n-1)! when n>0
156 pool = tuple(iterable)
157 n = len(pool)
158 if not n and r:
159 return
160 indices = [0] * r
161 yield tuple(pool[i] for i in indices)
162 while 1:
163 for i in reversed(range(r)):
164 if indices[i] != n - 1:
165 break
166 else:
167 return
168 indices[i:] = [indices[i] + 1] * (r - i)
169 yield tuple(pool[i] for i in indices)
170
171 def cwr2(iterable, r):
172 'Pure python version shown in the docs'
173 pool = tuple(iterable)
174 n = len(pool)
175 for indices in product(range(n), repeat=r):
176 if sorted(indices) == list(indices):
177 yield tuple(pool[i] for i in indices)
178
179 def numcombs(n, r):
180 if not n:
181 return 0 if r else 1
182 return fact(n+r-1) / fact(r)/ fact(n-1)
183
184 for n in range(7):
185 values = [5*x-12 for x in range(n)]
186 for r in range(n+2):
187 result = list(cwr(values, r))
188
189 self.assertEqual(len(result), numcombs(n, r)) # right number of combs
190 self.assertEqual(len(result), len(set(result))) # no repeats
191 self.assertEqual(result, sorted(result)) # lexicographic order
192
193 regular_combs = list(combinations(values, r)) # compare to combs without replacement
194 if n == 0 or r <= 1:
195 self.assertEquals(result, regular_combs) # cases that should be identical
196 else:
197 self.assert_(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
198
199 for c in result:
200 self.assertEqual(len(c), r) # r-length combinations
201 noruns = [k for k,v in groupby(c)] # combo without consecutive repeats
202 self.assertEqual(len(noruns), len(set(noruns))) # no repeats other than consecutive
203 self.assertEqual(list(c), sorted(c)) # keep original ordering
204 self.assert_(all(e in values for e in c)) # elements taken from input iterable
205 self.assertEqual(noruns,
206 [e for e in values if e in c]) # comb is a subsequence of the input iterable
207 self.assertEqual(result, list(cwr1(values, r))) # matches first pure python version
208 self.assertEqual(result, list(cwr2(values, r))) # matches second pure python version
209
210 # Test implementation detail: tuple re-use
211 self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
212 self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
213
Christian Heimes78644762008-03-04 23:39:23 +0000214 def test_permutations(self):
215 self.assertRaises(TypeError, permutations) # too few arguments
216 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000217 self.assertRaises(TypeError, permutations, None) # pool is not iterable
218 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000219 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000220 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Christian Heimes78644762008-03-04 23:39:23 +0000221 self.assertEqual(list(permutations(range(3), 2)),
222 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
223
224 def permutations1(iterable, r=None):
225 'Pure python version shown in the docs'
226 pool = tuple(iterable)
227 n = len(pool)
228 r = n if r is None else r
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000229 if r > n:
230 return
Christian Heimes78644762008-03-04 23:39:23 +0000231 indices = list(range(n))
232 cycles = list(range(n-r+1, n+1))[::-1]
233 yield tuple(pool[i] for i in indices[:r])
234 while n:
235 for i in reversed(range(r)):
236 cycles[i] -= 1
237 if cycles[i] == 0:
238 indices[i:] = indices[i+1:] + indices[i:i+1]
239 cycles[i] = n - i
240 else:
241 j = cycles[i]
242 indices[i], indices[-j] = indices[-j], indices[i]
243 yield tuple(pool[i] for i in indices[:r])
244 break
245 else:
246 return
247
248 def permutations2(iterable, r=None):
249 'Pure python version shown in the docs'
250 pool = tuple(iterable)
251 n = len(pool)
252 r = n if r is None else r
253 for indices in product(range(n), repeat=r):
254 if len(set(indices)) == r:
255 yield tuple(pool[i] for i in indices)
256
257 for n in range(7):
258 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000259 for r in range(n+2):
Christian Heimes78644762008-03-04 23:39:23 +0000260 result = list(permutations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000261 self.assertEqual(len(result), 0 if r>n else fact(n) / fact(n-r)) # right number of perms
Christian Heimes78644762008-03-04 23:39:23 +0000262 self.assertEqual(len(result), len(set(result))) # no repeats
263 self.assertEqual(result, sorted(result)) # lexicographic order
264 for p in result:
265 self.assertEqual(len(p), r) # r-length permutations
266 self.assertEqual(len(set(p)), r) # no duplicate elements
267 self.assert_(all(e in values for e in p)) # elements taken from input iterable
268 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000269 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000270 if r == n:
271 self.assertEqual(result, list(permutations(values, None))) # test r as None
272 self.assertEqual(result, list(permutations(values))) # test default r
273
274 # Test implementation detail: tuple re-use
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000275 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Christian Heimes78644762008-03-04 23:39:23 +0000276 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Christian Heimes380f7f22008-02-28 11:19:05 +0000277
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000278 def test_combinatorics(self):
279 # Test relationships between product(), permutations(),
280 # combinations() and combinations_with_replacement().
281
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000282 for n in range(6):
283 s = 'ABCDEFG'[:n]
284 for r in range(8):
285 prod = list(product(s, repeat=r))
286 cwr = list(combinations_with_replacement(s, r))
287 perm = list(permutations(s, r))
288 comb = list(combinations(s, r))
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000289
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000290 # Check size
291 self.assertEquals(len(prod), n**r)
292 self.assertEquals(len(cwr), (fact(n+r-1) / fact(r)/ fact(n-1)) if n else (not r))
293 self.assertEquals(len(perm), 0 if r>n else fact(n) / fact(n-r))
294 self.assertEquals(len(comb), 0 if r>n else fact(n) / fact(r) / fact(n-r))
295
296 # Check lexicographic order without repeated tuples
297 self.assertEquals(prod, sorted(set(prod)))
298 self.assertEquals(cwr, sorted(set(cwr)))
299 self.assertEquals(perm, sorted(set(perm)))
300 self.assertEquals(comb, sorted(set(comb)))
301
302 # Check interrelationships
303 self.assertEquals(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
304 self.assertEquals(perm, [t for t in prod if len(set(t))==r]) # perm: prods with no dups
305 self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
306 self.assertEqual(comb, [t for t in cwr if len(set(t))==r]) # comb: cwrs without dups
307 self.assertEqual(comb, list(filter(set(cwr).__contains__, perm))) # comb: perm that is a cwr
308 self.assertEqual(comb, list(filter(set(perm).__contains__, cwr))) # comb: cwr that is a perm
309 self.assertEqual(comb, sorted(set(cwr) & set(perm))) # comb: both a cwr and a perm
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000310
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000311 def test_compress(self):
312 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
313 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
314 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
315 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
316 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
317 n = 10000
318 data = chain.from_iterable(repeat(range(6), n))
319 selectors = chain.from_iterable(repeat((0, 1)))
320 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
321 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
322 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
323 self.assertRaises(TypeError, compress, range(6)) # too few args
324 self.assertRaises(TypeError, compress, range(6), None) # too many args
325
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000326 def test_count(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000327 self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
328 self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
329 self.assertEqual(take(2, lzip('abc',count(3))), [('a', 3), ('b', 4)])
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000330 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
331 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000332 self.assertRaises(TypeError, count, 2, 3, 4)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000333 self.assertRaises(TypeError, count, 'a')
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000334 self.assertEqual(list(islice(count(maxsize-5), 10)),
335 list(range(maxsize-5, maxsize+5)))
336 self.assertEqual(list(islice(count(-maxsize-5), 10)),
337 list(range(-maxsize-5, -maxsize+5)))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000338 c = count(3)
339 self.assertEqual(repr(c), 'count(3)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000340 next(c)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000341 self.assertEqual(repr(c), 'count(4)')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000342 c = count(-9)
343 self.assertEqual(repr(c), 'count(-9)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000344 next(c)
Raymond Hettinger30729212009-02-12 06:28:27 +0000345 self.assertEqual(repr(count(10.25)), 'count(10.25)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000346 self.assertEqual(next(c), -8)
Christian Heimesa37d4c62007-12-04 23:02:19 +0000347 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000348 # Test repr (ignoring the L in longs)
349 r1 = repr(count(i)).replace('L', '')
350 r2 = 'count(%r)'.__mod__(i).replace('L', '')
351 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000352
Raymond Hettinger30729212009-02-12 06:28:27 +0000353 def test_count_with_stride(self):
354 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
355 self.assertEqual(lzip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
356 self.assertEqual(lzip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
357 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
358 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
359 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettinger8a882a72009-02-12 12:08:18 +0000360 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
361 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettinger30729212009-02-12 06:28:27 +0000362 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
363 c = count(3, 5)
364 self.assertEqual(repr(c), 'count(3, 5)')
365 next(c)
366 self.assertEqual(repr(c), 'count(8, 5)')
367 c = count(-9, 0)
368 self.assertEqual(repr(c), 'count(-9, 0)')
369 next(c)
370 self.assertEqual(repr(c), 'count(-9, 0)')
371 c = count(-9, -3)
372 self.assertEqual(repr(c), 'count(-9, -3)')
373 next(c)
374 self.assertEqual(repr(c), 'count(-12, -3)')
375 self.assertEqual(repr(c), 'count(-12, -3)')
376 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
377 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
378 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
379 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
380 for j in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 1, 10, sys.maxsize-5, sys.maxsize+5):
381 # Test repr (ignoring the L in longs)
382 r1 = repr(count(i, j)).replace('L', '')
383 if j == 1:
384 r2 = ('count(%r)' % i).replace('L', '')
385 else:
386 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
387 self.assertEqual(r1, r2)
388
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000389 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000390 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000391 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000392 self.assertRaises(TypeError, cycle)
393 self.assertRaises(TypeError, cycle, 5)
394 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000395
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000396 def test_groupby(self):
397 # Check whether it accepts arguments correctly
398 self.assertEqual([], list(groupby([])))
399 self.assertEqual([], list(groupby([], key=id)))
400 self.assertRaises(TypeError, list, groupby('abc', []))
401 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000402 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000403
404 # Check normal input
405 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
406 (2,15,22), (3,16,23), (3,17,23)]
407 dup = []
408 for k, g in groupby(s, lambda r:r[0]):
409 for elem in g:
410 self.assertEqual(k, elem[0])
411 dup.append(elem)
412 self.assertEqual(s, dup)
413
414 # Check nested case
415 dup = []
416 for k, g in groupby(s, lambda r:r[0]):
417 for ik, ig in groupby(g, lambda r:r[2]):
418 for elem in ig:
419 self.assertEqual(k, elem[0])
420 self.assertEqual(ik, elem[2])
421 dup.append(elem)
422 self.assertEqual(s, dup)
423
424 # Check case where inner iterator is not used
425 keys = [k for k, g in groupby(s, lambda r:r[0])]
426 expectedkeys = set([r[0] for r in s])
427 self.assertEqual(set(keys), expectedkeys)
428 self.assertEqual(len(keys), len(expectedkeys))
429
430 # Exercise pipes and filters style
431 s = 'abracadabra'
432 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000433 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000434 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
435 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000436 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000437 self.assertEqual(r, ['a', 'b', 'r'])
438 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000439 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000440 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
441 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000442 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000443 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
444
Georg Brandla18af4e2007-04-21 15:47:16 +0000445 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000446 class ExpectedError(Exception):
447 pass
448 def delayed_raise(n=0):
449 for i in range(n):
450 yield 'yo'
451 raise ExpectedError
452 def gulp(iterable, keyp=None, func=list):
453 return [func(g) for k, g in groupby(iterable, keyp)]
454
Georg Brandla18af4e2007-04-21 15:47:16 +0000455 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000456 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000457 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000458 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
459
460 # __cmp__ failure
461 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000462 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000463 raise ExpectedError
464 s = [DummyCmp(), DummyCmp(), None]
465
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000466 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000467 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000468 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000469 self.assertRaises(ExpectedError, gulp, s)
470
471 # keyfunc failure
472 def keyfunc(obj):
473 if keyfunc.skip > 0:
474 keyfunc.skip -= 1
475 return obj
476 else:
477 raise ExpectedError
478
479 # keyfunc failure on outer object
480 keyfunc.skip = 0
481 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
482 keyfunc.skip = 1
483 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
484
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000485 def test_filter(self):
486 self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
487 self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])
488 self.assertEqual(list(filter(bool, [0,1,0,2,0])), [1,2])
489 self.assertEqual(take(4, filter(isEven, count())), [0,2,4,6])
490 self.assertRaises(TypeError, filter)
491 self.assertRaises(TypeError, filter, lambda x:x)
492 self.assertRaises(TypeError, filter, lambda x:x, range(6), 7)
493 self.assertRaises(TypeError, filter, isEven, 3)
494 self.assertRaises(TypeError, next, filter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000495
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000496 def test_filterfalse(self):
497 self.assertEqual(list(filterfalse(isEven, range(6))), [1,3,5])
498 self.assertEqual(list(filterfalse(None, [0,1,0,2,0])), [0,0,0])
499 self.assertEqual(list(filterfalse(bool, [0,1,0,2,0])), [0,0,0])
500 self.assertEqual(take(4, filterfalse(isEven, count())), [1,3,5,7])
501 self.assertRaises(TypeError, filterfalse)
502 self.assertRaises(TypeError, filterfalse, lambda x:x)
503 self.assertRaises(TypeError, filterfalse, lambda x:x, range(6), 7)
504 self.assertRaises(TypeError, filterfalse, isEven, 3)
505 self.assertRaises(TypeError, next, filterfalse(range(6), range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000506
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000507 def test_zip(self):
508 # XXX This is rather silly now that builtin zip() calls zip()...
509 ans = [(x,y) for x, y in zip('abc',count())]
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000510 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000511 self.assertEqual(list(zip('abc', range(6))), lzip('abc', range(6)))
512 self.assertEqual(list(zip('abcdef', range(3))), lzip('abcdef', range(3)))
513 self.assertEqual(take(3,zip('abcdef', count())), lzip('abcdef', range(3)))
514 self.assertEqual(list(zip('abcdef')), lzip('abcdef'))
515 self.assertEqual(list(zip()), lzip())
516 self.assertRaises(TypeError, zip, 3)
517 self.assertRaises(TypeError, zip, range(3), 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000518 # Check tuple re-use (implementation detail)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000519 self.assertEqual([tuple(list(pair)) for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000520 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000521 self.assertEqual([pair for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000522 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000523 ids = list(map(id, zip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000524 self.assertEqual(min(ids), max(ids))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000525 ids = list(map(id, list(zip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000526 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000527
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000528 def test_ziplongest(self):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000529 for args in [
530 ['abc', range(6)],
531 [range(6), 'abc'],
532 [range(1000), range(2000,2100), range(3000,3050)],
533 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
534 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
535 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000536 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
537 for i in range(max(map(len, args)))]
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000538 self.assertEqual(list(zip_longest(*args)), target)
539 self.assertEqual(list(zip_longest(*args, **{})), target)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000540 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000541 self.assertEqual(list(zip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000542
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000543 self.assertEqual(take(3,zip_longest('abcdef', count())), list(zip('abcdef', range(3)))) # take 3 from infinite input
Thomas Wouterscf297e42007-02-23 15:07:44 +0000544
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000545 self.assertEqual(list(zip_longest()), list(zip()))
546 self.assertEqual(list(zip_longest([])), list(zip([])))
547 self.assertEqual(list(zip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000548
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000549 self.assertEqual(list(zip_longest('abc', 'defg', **{})),
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000550 list(zip(list('abc')+[None], 'defg'))) # empty keyword dict
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000551 self.assertRaises(TypeError, zip_longest, 3)
552 self.assertRaises(TypeError, zip_longest, range(3), 3)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000553
554 for stmt in [
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000555 "zip_longest('abc', fv=1)",
556 "zip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000557 ]:
558 try:
559 eval(stmt, globals(), locals())
560 except TypeError:
561 pass
562 else:
563 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000564
Thomas Wouterscf297e42007-02-23 15:07:44 +0000565 # Check tuple re-use (implementation detail)
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000566 self.assertEqual([tuple(list(pair)) for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000567 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000568 self.assertEqual([pair for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000569 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000570 ids = list(map(id, zip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000571 self.assertEqual(min(ids), max(ids))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000572 ids = list(map(id, list(zip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000573 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
574
Christian Heimesc3f30c42008-02-22 16:37:40 +0000575 def test_product(self):
576 for args, result in [
Christian Heimes78644762008-03-04 23:39:23 +0000577 ([], [()]), # zero iterables
Christian Heimesc3f30c42008-02-22 16:37:40 +0000578 (['ab'], [('a',), ('b',)]), # one iterable
579 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
580 ([range(0), range(2), range(3)], []), # first iterable with zero length
581 ([range(2), range(0), range(3)], []), # middle iterable with zero length
582 ([range(2), range(3), range(0)], []), # last iterable with zero length
583 ]:
584 self.assertEqual(list(product(*args)), result)
Christian Heimesf16baeb2008-02-29 14:57:44 +0000585 for r in range(4):
586 self.assertEqual(list(product(*(args*r))),
587 list(product(*args, **dict(repeat=r))))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000588 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
589 self.assertRaises(TypeError, product, range(6), None)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000590
591 def product1(*args, **kwds):
592 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
593 n = len(pools)
594 if n == 0:
595 yield ()
596 return
597 if any(len(pool) == 0 for pool in pools):
598 return
599 indices = [0] * n
600 yield tuple(pool[i] for pool, i in zip(pools, indices))
601 while 1:
602 for i in reversed(range(n)): # right to left
603 if indices[i] == len(pools[i]) - 1:
604 continue
605 indices[i] += 1
606 for j in range(i+1, n):
607 indices[j] = 0
608 yield tuple(pool[i] for pool, i in zip(pools, indices))
609 break
610 else:
611 return
612
613 def product2(*args, **kwds):
614 'Pure python version used in docs'
615 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
616 result = [[]]
617 for pool in pools:
618 result = [x+[y] for x in result for y in pool]
619 for prod in result:
620 yield tuple(prod)
621
Christian Heimesc3f30c42008-02-22 16:37:40 +0000622 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
623 set('abcdefg'), range(11), tuple(range(13))]
624 for i in range(100):
625 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Christian Heimes78644762008-03-04 23:39:23 +0000626 expected_len = prod(map(len, args))
627 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000628 self.assertEqual(list(product(*args)), list(product1(*args)))
629 self.assertEqual(list(product(*args)), list(product2(*args)))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000630 args = map(iter, args)
Christian Heimes78644762008-03-04 23:39:23 +0000631 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000632
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000633 # Test implementation detail: tuple re-use
634 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
635 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000636
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000637 def test_repeat(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000638 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000639 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000640 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000641 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000642 self.assertEqual(list(repeat('a', 0)), [])
643 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000644 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000645 self.assertRaises(TypeError, repeat, None, 3, 4)
646 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000647 r = repeat(1+0j)
648 self.assertEqual(repr(r), 'repeat((1+0j))')
649 r = repeat(1+0j, 5)
650 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
651 list(r)
652 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000653
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000654 def test_map(self):
655 self.assertEqual(list(map(operator.pow, range(3), range(1,7))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000656 [0**1, 1**2, 2**3])
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000657 def tupleize(*args):
658 return args
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000659 self.assertEqual(list(map(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000660 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000661 self.assertEqual(list(map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000662 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000663 self.assertEqual(take(2,map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000664 [('a',0),('b',1)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000665 self.assertEqual(list(map(operator.pow, [])), [])
666 self.assertRaises(TypeError, map)
667 self.assertRaises(TypeError, list, map(None, range(3), range(3)))
668 self.assertRaises(TypeError, map, operator.neg)
669 self.assertRaises(TypeError, next, map(10, range(5)))
670 self.assertRaises(ValueError, next, map(errfunc, [4], [5]))
671 self.assertRaises(TypeError, next, map(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000672
673 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000674 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
675 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000676 self.assertEqual(take(3, starmap(operator.pow, zip(count(), count(1)))),
Raymond Hettinger02420702003-06-29 20:36:23 +0000677 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000678 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +0000679 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
680 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000681 self.assertRaises(TypeError, starmap)
682 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000683 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
684 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
685 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000686
687 def test_islice(self):
688 for args in [ # islice(args) should agree with range(args)
689 (10, 20, 3),
690 (10, 3, 20),
691 (10, 20),
692 (10, 3),
693 (20,)
694 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000695 self.assertEqual(list(islice(range(100), *args)),
696 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000697
698 for args, tgtargs in [ # Stop when seqn is exhausted
699 ((10, 110, 3), ((10, 100, 3))),
700 ((10, 110), ((10, 100))),
701 ((110,), (100,))
702 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000703 self.assertEqual(list(islice(range(100), *args)),
704 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000705
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000706 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +0000707 self.assertEqual(list(islice(range(10), None)), list(range(10)))
708 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
709 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
710 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
711 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000712
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000713 # Test number of items consumed SF #1171417
714 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +0000715 self.assertEqual(list(islice(it, 3)), list(range(3)))
716 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000717
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000718 # Test invalid arguments
Guido van Rossum805365e2007-05-07 22:24:25 +0000719 self.assertRaises(TypeError, islice, range(10))
720 self.assertRaises(TypeError, islice, range(10), 1, 2, 3, 4)
721 self.assertRaises(ValueError, islice, range(10), -5, 10, 1)
722 self.assertRaises(ValueError, islice, range(10), 1, -5, -1)
723 self.assertRaises(ValueError, islice, range(10), 1, 10, -1)
724 self.assertRaises(ValueError, islice, range(10), 1, 10, 0)
725 self.assertRaises(ValueError, islice, range(10), 'a')
726 self.assertRaises(ValueError, islice, range(10), 'a', 1)
727 self.assertRaises(ValueError, islice, range(10), 1, 'a')
728 self.assertRaises(ValueError, islice, range(10), 'a', 1, 1)
729 self.assertRaises(ValueError, islice, range(10), 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +0000730 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000731
732 def test_takewhile(self):
733 data = [1, 3, 5, 20, 2, 4, 6, 8]
734 underten = lambda x: x<10
735 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000736 self.assertEqual(list(takewhile(underten, [])), [])
737 self.assertRaises(TypeError, takewhile)
738 self.assertRaises(TypeError, takewhile, operator.pow)
739 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000740 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
741 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000742 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
743 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +0000744 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000745
746 def test_dropwhile(self):
747 data = [1, 3, 5, 20, 2, 4, 6, 8]
748 underten = lambda x: x<10
749 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000750 self.assertEqual(list(dropwhile(underten, [])), [])
751 self.assertRaises(TypeError, dropwhile)
752 self.assertRaises(TypeError, dropwhile, operator.pow)
753 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000754 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
755 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000756
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000757 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000758 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000759 def irange(n):
Guido van Rossum805365e2007-05-07 22:24:25 +0000760 for i in range(n):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000761 yield i
762
763 a, b = tee([]) # test empty iterator
764 self.assertEqual(list(a), [])
765 self.assertEqual(list(b), [])
766
767 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +0000768 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000769
770 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +0000771 self.assertEqual(list(a), list(range(n)))
772 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000773
774 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000775 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000776 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000777 del a
Guido van Rossum805365e2007-05-07 22:24:25 +0000778 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000779
780 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000781 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000782 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000783 del b
Guido van Rossum805365e2007-05-07 22:24:25 +0000784 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000785
Guido van Rossum805365e2007-05-07 22:24:25 +0000786 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000787 order = [0]*n + [1]*n
788 random.shuffle(order)
789 lists = ([], [])
790 its = tee(irange(n))
791 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +0000792 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000793 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +0000794 self.assertEqual(lists[0], list(range(n)))
795 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000796
Raymond Hettingerad983e72003-11-12 14:32:26 +0000797 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000798 self.assertRaises(TypeError, tee)
799 self.assertRaises(TypeError, tee, 3)
800 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000801 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000802
Raymond Hettingerad983e72003-11-12 14:32:26 +0000803 # tee object should be instantiable
804 a, b = tee('abc')
805 c = type(a)('def')
806 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000807
Raymond Hettingerad983e72003-11-12 14:32:26 +0000808 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +0000809 a, b, c = tee(range(2000), 3)
810 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000811 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +0000812 self.assertEqual(list(b), list(range(2000)))
813 self.assertEqual([next(c), next(c)], list(range(2)))
814 self.assertEqual(list(a), list(range(100,2000)))
815 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000816
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000817 # test values of n
818 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000819 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +0000820 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000821 result = tee('abc', n)
822 self.assertEqual(type(result), tuple)
823 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000824 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000825
Raymond Hettingerad983e72003-11-12 14:32:26 +0000826 # tee pass-through to copyable iterator
827 a, b = tee('abc')
828 c, d = tee(a)
829 self.assert_(a is c)
830
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000831 # test tee_new
832 t1, t2 = tee('abc')
833 tnew = type(t1)
834 self.assertRaises(TypeError, tnew)
835 self.assertRaises(TypeError, tnew, 10)
836 t3 = tnew(t1)
837 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000838
Raymond Hettingera9f60922004-10-17 16:40:14 +0000839 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +0000840 a, b = tee(range(10))
Raymond Hettingera9f60922004-10-17 16:40:14 +0000841 p = proxy(a)
842 self.assertEqual(getattr(p, '__class__'), type(b))
843 del a
844 self.assertRaises(ReferenceError, getattr, p, '__class__')
845
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000846 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000847 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000848
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000849 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +0000850 self.assertRaises(StopIteration, next, f([]))
851 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000852
Georg Brandla18af4e2007-04-21 15:47:16 +0000853 self.assertRaises(StopIteration, next, islice([], None))
854 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000855
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000856 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +0000857 self.assertRaises(StopIteration, next, p)
858 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000859 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +0000860 self.assertRaises(StopIteration, next, p)
861 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000862
Georg Brandla18af4e2007-04-21 15:47:16 +0000863 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000864
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000865 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +0000866 self.assertRaises(StopIteration, next, f(lambda x:x, []))
867 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000868
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000869class TestExamples(unittest.TestCase):
870
871 def test_chain(self):
872 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
873
874 def test_chain_from_iterable(self):
875 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
876
877 def test_combinations(self):
878 self.assertEqual(list(combinations('ABCD', 2)),
879 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
880 self.assertEqual(list(combinations(range(4), 3)),
881 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
882
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000883 def test_combinations_with_replacement(self):
884 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
885 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
886
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000887 def test_compress(self):
888 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
889
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000890 def test_count(self):
891 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
892
893 def test_cycle(self):
894 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
895
896 def test_dropwhile(self):
897 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
898
899 def test_groupby(self):
900 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
901 list('ABCDAB'))
902 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
903 [list('AAAA'), list('BBB'), list('CC'), list('D')])
904
905 def test_filter(self):
906 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
907
908 def test_filterfalse(self):
909 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
910
911 def test_map(self):
912 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
913
914 def test_islice(self):
915 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
916 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
917 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
918 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
919
920 def test_zip(self):
921 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
922
923 def test_zip_longest(self):
924 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
925 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
926
927 def test_permutations(self):
928 self.assertEqual(list(permutations('ABCD', 2)),
929 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
930 self.assertEqual(list(permutations(range(3))),
931 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
932
933 def test_product(self):
934 self.assertEqual(list(product('ABCD', 'xy')),
935 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
936 self.assertEqual(list(product(range(2), repeat=3)),
937 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
938 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
939
940 def test_repeat(self):
941 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
942
943 def test_stapmap(self):
944 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
945 [32, 9, 1000])
946
947 def test_takewhile(self):
948 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
949
950
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000951class TestGC(unittest.TestCase):
952
953 def makecycle(self, iterator, container):
954 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +0000955 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000956 del container, iterator
957
958 def test_chain(self):
959 a = []
960 self.makecycle(chain(a), a)
961
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000962 def test_chain_from_iterable(self):
963 a = []
964 self.makecycle(chain.from_iterable([a]), a)
965
966 def test_combinations(self):
967 a = []
968 self.makecycle(combinations([1,2,a,3], 3), a)
969
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000970 def test_combinations_with_replacement(self):
971 a = []
972 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
973
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000974 def test_compress(self):
975 a = []
976 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
977
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000978 def test_cycle(self):
979 a = []
980 self.makecycle(cycle([a]*2), a)
981
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000982 def test_dropwhile(self):
983 a = []
984 self.makecycle(dropwhile(bool, [0, a, a]), a)
985
986 def test_groupby(self):
987 a = []
988 self.makecycle(groupby([a]*2, lambda x:x), a)
989
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000990 def test_issue2246(self):
991 # Issue 2246 -- the _grouper iterator was not included in GC
992 n = 10
993 keyfunc = lambda x: x
994 for i, j in groupby(range(n), key=keyfunc):
995 keyfunc.__dict__.setdefault('x',[]).append(j)
996
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000997 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000998 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000999 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001000
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001001 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001002 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001003 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001004
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001005 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001006 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001007 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001008
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001009 def test_zip_longest(self):
1010 a = []
1011 self.makecycle(zip_longest([a]*2, [a]*3), a)
1012 b = [a, None]
1013 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1014
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001015 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001016 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001017 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001018
1019 def test_islice(self):
1020 a = []
1021 self.makecycle(islice([a]*2, None), a)
1022
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001023 def test_permutations(self):
1024 a = []
1025 self.makecycle(permutations([1,2,a,3], 3), a)
1026
1027 def test_product(self):
1028 a = []
1029 self.makecycle(product([1,2,a,3], repeat=3), a)
1030
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001031 def test_repeat(self):
1032 a = []
1033 self.makecycle(repeat(a), a)
1034
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001035 def test_starmap(self):
1036 a = []
1037 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1038
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001039 def test_takewhile(self):
1040 a = []
1041 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1042
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001043def R(seqn):
1044 'Regular generator'
1045 for i in seqn:
1046 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001047
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001048class G:
1049 'Sequence using __getitem__'
1050 def __init__(self, seqn):
1051 self.seqn = seqn
1052 def __getitem__(self, i):
1053 return self.seqn[i]
1054
1055class I:
1056 'Sequence using iterator protocol'
1057 def __init__(self, seqn):
1058 self.seqn = seqn
1059 self.i = 0
1060 def __iter__(self):
1061 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001062 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001063 if self.i >= len(self.seqn): raise StopIteration
1064 v = self.seqn[self.i]
1065 self.i += 1
1066 return v
1067
1068class Ig:
1069 'Sequence using iterator protocol defined with a generator'
1070 def __init__(self, seqn):
1071 self.seqn = seqn
1072 self.i = 0
1073 def __iter__(self):
1074 for val in self.seqn:
1075 yield val
1076
1077class X:
1078 'Missing __getitem__ and __iter__'
1079 def __init__(self, seqn):
1080 self.seqn = seqn
1081 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001082 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001083 if self.i >= len(self.seqn): raise StopIteration
1084 v = self.seqn[self.i]
1085 self.i += 1
1086 return v
1087
1088class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001089 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001090 def __init__(self, seqn):
1091 self.seqn = seqn
1092 self.i = 0
1093 def __iter__(self):
1094 return self
1095
1096class E:
1097 'Test propagation of exceptions'
1098 def __init__(self, seqn):
1099 self.seqn = seqn
1100 self.i = 0
1101 def __iter__(self):
1102 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001103 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001104 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001105
1106class S:
1107 'Test immediate stop'
1108 def __init__(self, seqn):
1109 pass
1110 def __iter__(self):
1111 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001112 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001113 raise StopIteration
1114
1115def L(seqn):
1116 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001117 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001118
1119
1120class TestVariousIteratorArgs(unittest.TestCase):
1121
1122 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001123 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001124 for g in (G, I, Ig, S, L, R):
1125 self.assertEqual(list(chain(g(s))), list(g(s)))
1126 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001127 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001128 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001129 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1130
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001131 def test_compress(self):
1132 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1133 n = len(s)
1134 for g in (G, I, Ig, S, L, R):
1135 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1136 self.assertRaises(TypeError, compress, X(s), repeat(1))
1137 self.assertRaises(TypeError, compress, N(s), repeat(1))
1138 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1139
Christian Heimesc3f30c42008-02-22 16:37:40 +00001140 def test_product(self):
1141 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1142 self.assertRaises(TypeError, product, X(s))
1143 self.assertRaises(TypeError, product, N(s))
1144 self.assertRaises(ZeroDivisionError, product, E(s))
1145
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001146 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001147 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001148 for g in (G, I, Ig, S, L, R):
1149 tgtlen = len(s) * 3
1150 expected = list(g(s))*3
1151 actual = list(islice(cycle(g(s)), tgtlen))
1152 self.assertEqual(actual, expected)
1153 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001154 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001155 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1156
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001157 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001158 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001159 for g in (G, I, Ig, S, L, R):
1160 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1161 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001162 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001163 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1164
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001165 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001166 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001167 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001168 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001169 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001170 self.assertRaises(TypeError, filter, isEven, X(s))
1171 self.assertRaises(TypeError, filter, isEven, N(s))
1172 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001173
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001174 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001175 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001176 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001177 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001178 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001179 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1180 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1181 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001182
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001183 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001184 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001185 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001186 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1187 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1188 self.assertRaises(TypeError, zip, X(s))
1189 self.assertRaises(TypeError, zip, N(s))
1190 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001191
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001192 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001193 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001194 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001195 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1196 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1197 self.assertRaises(TypeError, zip_longest, X(s))
1198 self.assertRaises(TypeError, zip_longest, N(s))
1199 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001200
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001201 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001202 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001203 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001204 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001205 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001206 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001207 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001208 self.assertRaises(TypeError, map, onearg, X(s))
1209 self.assertRaises(TypeError, map, onearg, N(s))
1210 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001211
1212 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001213 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001214 for g in (G, I, Ig, S, L, R):
1215 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1216 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001217 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001218 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1219
1220 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001221 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001222 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001223 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001224 self.assertEqual(list(starmap(operator.pow, g(ss))),
1225 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001226 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001227 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001228 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1229
1230 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001231 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001232 for g in (G, I, Ig, S, L, R):
1233 tgt = []
1234 for elem in g(s):
1235 if not isEven(elem): break
1236 tgt.append(elem)
1237 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1238 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001239 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001240 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1241
1242 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001243 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001244 for g in (G, I, Ig, S, L, R):
1245 tgt = []
1246 for elem in g(s):
1247 if not tgt and isOdd(elem): continue
1248 tgt.append(elem)
1249 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1250 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001251 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001252 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1253
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001254 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001255 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001256 for g in (G, I, Ig, S, L, R):
1257 it1, it2 = tee(g(s))
1258 self.assertEqual(list(it1), list(g(s)))
1259 self.assertEqual(list(it2), list(g(s)))
1260 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001261 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001262 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1263
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001264class LengthTransparency(unittest.TestCase):
1265
1266 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001267 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001268 self.assertEqual(len(repeat(None, 50)), 50)
1269 self.assertRaises(TypeError, len, repeat(None))
1270
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001271class RegressionTests(unittest.TestCase):
1272
1273 def test_sf_793826(self):
1274 # Fix Armin Rigo's successful efforts to wreak havoc
1275
1276 def mutatingtuple(tuple1, f, tuple2):
1277 # this builds a tuple t which is a copy of tuple1,
1278 # then calls f(t), then mutates t to be equal to tuple2
1279 # (needs len(tuple1) == len(tuple2)).
1280 def g(value, first=[1]):
1281 if first:
1282 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00001283 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001284 return value
1285 items = list(tuple2)
1286 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001287 gen = map(g, items)
1288 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00001289 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001290
1291 def f(t):
1292 global T
1293 T = t
1294 first[:] = list(T)
1295
1296 first = []
1297 mutatingtuple((1,2,3), f, (4,5,6))
1298 second = list(T)
1299 self.assertEqual(first, second)
1300
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001301
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001302 def test_sf_950057(self):
1303 # Make sure that chain() and cycle() catch exceptions immediately
1304 # rather than when shifting between input sources
1305
1306 def gen1():
1307 hist.append(0)
1308 yield 1
1309 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001310 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001311 hist.append(2)
1312
1313 def gen2(x):
1314 hist.append(3)
1315 yield 2
1316 hist.append(4)
1317 if x:
1318 raise StopIteration
1319
1320 hist = []
1321 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1322 self.assertEqual(hist, [0,1])
1323
1324 hist = []
1325 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1326 self.assertEqual(hist, [0,1])
1327
1328 hist = []
1329 self.assertRaises(AssertionError, list, cycle(gen1()))
1330 self.assertEqual(hist, [0,1])
1331
Thomas Woutersb2137042007-02-01 18:02:27 +00001332class SubclassWithKwargsTest(unittest.TestCase):
1333 def test_keywords_in_subclass(self):
1334 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001335 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001336 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00001337 class Subclass(cls):
1338 def __init__(self, newarg=None, *args):
1339 cls.__init__(self, *args)
1340 try:
1341 Subclass(newarg=1)
1342 except TypeError as err:
1343 # we expect type errors because of wrong argument count
1344 self.failIf("does not take keyword arguments" in err.args[0])
1345
1346
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001347libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001348
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001349
1350>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001351>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00001352... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001353...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001354Check 1200 is for $120.15
1355Check 1201 is for $764.05
1356Check 1202 is for $823.14
1357
1358>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001359>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00001360... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00001361...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000013621
13638
136427
1365
1366>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001367>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00001368... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00001369...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001370Alex
1371Laura
1372Martin
1373Walter
1374Samuele
1375
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001376>>> from operator import itemgetter
1377>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00001378>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001379>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001380... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001381...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000013821 ['a', 'c', 'e']
13832 ['b', 'd', 'f']
13843 ['g']
1385
Raymond Hettinger734fb572004-01-20 20:04:40 +00001386# Find runs of consecutive numbers using groupby. The key to the solution
1387# is differencing with a range so that consecutive numbers all appear in
1388# same group.
1389>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00001390>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001391... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001392...
Raymond Hettinger734fb572004-01-20 20:04:40 +00001393[1]
1394[4, 5, 6]
1395[10]
1396[15, 16, 17, 18]
1397[22]
1398[25, 26, 27, 28]
1399
Georg Brandl3dbca812008-07-23 16:10:53 +00001400>>> def take(n, iterable):
1401... "Return first n items of the iterable as a list"
1402... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001403
Georg Brandl3dbca812008-07-23 16:10:53 +00001404>>> def enumerate(iterable, start=0):
1405... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001406
Georg Brandl3dbca812008-07-23 16:10:53 +00001407>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001408... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00001409... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001410
1411>>> def nth(iterable, n):
Raymond Hettingerd04fa312009-02-04 19:45:13 +00001412... "Returns the nth item or None"
1413... return next(islice(iterable, n, None), None)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001414
Georg Brandl3dbca812008-07-23 16:10:53 +00001415>>> def quantify(iterable, pred=bool):
1416... "Count how many times the predicate is true"
1417... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001418
Georg Brandl3dbca812008-07-23 16:10:53 +00001419>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001420... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00001421... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001422
Georg Brandl3dbca812008-07-23 16:10:53 +00001423>>> def ncycles(iterable, n):
1424... "Returns the seqeuence elements n times"
1425... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001426
1427>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001428... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001429
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001430>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001431... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001432
1433>>> def repeatfunc(func, times=None, *args):
1434... "Repeat calls to func with specified arguments."
1435... " Example: repeatfunc(random.random)"
1436... if times is None:
1437... return starmap(func, repeat(args))
1438... else:
1439... return starmap(func, repeat(args, times))
1440
Raymond Hettingerd591f662003-10-26 15:34:50 +00001441>>> def pairwise(iterable):
1442... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1443... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001444... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00001445... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001446... except StopIteration:
1447... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001448... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001449
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001450>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001451... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001452... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00001453... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001454
1455>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001456... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001457... # Recipe credited to George Sakkis
1458... pending = len(iterables)
1459... nexts = cycle(iter(it).__next__ for it in iterables)
1460... while pending:
1461... try:
1462... for next in nexts:
1463... yield next()
1464... except StopIteration:
1465... pending -= 1
1466... nexts = cycle(islice(nexts, pending))
1467
1468>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00001469... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1470... s = list(iterable)
1471... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001472
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001473>>> def unique_everseen(iterable, key=None):
1474... "List unique elements, preserving order. Remember all elements ever seen."
1475... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1476... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1477... seen = set()
1478... seen_add = seen.add
1479... if key is None:
1480... for element in iterable:
1481... if element not in seen:
1482... seen_add(element)
1483... yield element
1484... else:
1485... for element in iterable:
1486... k = key(element)
1487... if k not in seen:
1488... seen_add(k)
1489... yield element
1490
1491>>> def unique_justseen(iterable, key=None):
1492... "List unique elements, preserving order. Remember only the element just seen."
1493... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1494... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1495... return map(next, map(itemgetter(1), groupby(iterable, key)))
1496
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001497This is not part of the examples but it tests to make sure the definitions
1498perform as purported.
1499
Raymond Hettingera098b332003-09-08 23:58:40 +00001500>>> take(10, count())
1501[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1502
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001503>>> list(enumerate('abc'))
1504[(0, 'a'), (1, 'b'), (2, 'c')]
1505
1506>>> list(islice(tabulate(lambda x: 2*x), 4))
1507[0, 2, 4, 6]
1508
1509>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00001510'd'
1511
1512>>> nth('abcde', 9) is None
1513True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001514
Guido van Rossum805365e2007-05-07 22:24:25 +00001515>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000151650
1517
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001518>>> a = [[1, 2, 3], [4, 5, 6]]
1519>>> flatten(a)
1520[1, 2, 3, 4, 5, 6]
1521
1522>>> list(repeatfunc(pow, 5, 2, 3))
1523[8, 8, 8, 8, 8]
1524
1525>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001526>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001527[0, 0, 0, 0, 0]
1528
Raymond Hettingerd591f662003-10-26 15:34:50 +00001529>>> list(pairwise('abcd'))
1530[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001531
Raymond Hettingerd591f662003-10-26 15:34:50 +00001532>>> list(pairwise([]))
1533[]
1534
1535>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001536[]
1537
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001538>>> list(islice(padnone('abc'), 0, 6))
1539['a', 'b', 'c', None, None, None]
1540
1541>>> list(ncycles('abc', 3))
1542['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1543
1544>>> dotproduct([1,2,3], [4,5,6])
154532
1546
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001547>>> list(grouper(3, 'abcdefg', 'x'))
1548[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1549
1550>>> list(roundrobin('abc', 'd', 'ef'))
1551['a', 'd', 'e', 'b', 'f', 'c']
1552
Raymond Hettingerace67332009-01-26 02:23:50 +00001553>>> list(powerset([1,2,3]))
1554[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001555
Raymond Hettinger191e8502009-01-27 13:29:43 +00001556>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1557True
1558
1559>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1560True
1561
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001562>>> list(unique_everseen('AAAABBBCCDAABBB'))
1563['A', 'B', 'C', 'D']
1564
1565>>> list(unique_everseen('ABBCcAD', str.lower))
1566['A', 'B', 'C', 'D']
1567
1568>>> list(unique_justseen('AAAABBBCCDAABBB'))
1569['A', 'B', 'C', 'D', 'A', 'B']
1570
1571>>> list(unique_justseen('ABBCcAD', str.lower))
1572['A', 'B', 'C', 'A', 'D']
1573
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001574"""
1575
1576__test__ = {'libreftest' : libreftest}
1577
1578def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001579 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00001580 RegressionTests, LengthTransparency,
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001581 SubclassWithKwargsTest, TestExamples)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001582 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001583
1584 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001585 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001586 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001587 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00001588 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001589 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001590 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001591 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001592 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001593
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001594 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001595 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001596
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001597if __name__ == "__main__":
1598 test_main(verbose=True)