blob: 31ef2e1578d320e5ac156c220418358d32378d48 [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 Hettingerde09acf2009-02-12 12:53:53 +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
Christian Heimesc3f30c42008-02-22 16:37:40 +000010from functools import reduce
Benjamin Petersonee8712c2008-05-20 21:35:26 +000011maxsize = support.MAX_Py_ssize_t
Guido van Rossum360e4b82007-05-14 22:51:27 +000012minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000013
Guido van Rossum801f0d72006-08-24 19:48:10 +000014def lzip(*args):
15 return list(zip(*args))
16
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000017def onearg(x):
18 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000019 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000020
21def errfunc(*args):
22 'Test function that raises an error'
23 raise ValueError
24
25def gen3():
26 'Non-restartable source sequence'
27 for i in (0, 1, 2):
28 yield i
29
30def isEven(x):
31 'Test predicate'
32 return x%2==0
33
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000034def isOdd(x):
35 'Test predicate'
36 return x%2==1
37
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000038class StopNow:
39 'Class emulating an empty iterable.'
40 def __iter__(self):
41 return self
Georg Brandla18af4e2007-04-21 15:47:16 +000042 def __next__(self):
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000043 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000044
Raymond Hettinger02420702003-06-29 20:36:23 +000045def take(n, seq):
46 'Convenience function for partially consuming a long of infinite iterable'
47 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000048
Christian Heimes78644762008-03-04 23:39:23 +000049def prod(iterable):
50 return reduce(operator.mul, iterable, 1)
51
Christian Heimes380f7f22008-02-28 11:19:05 +000052def fact(n):
53 'Factorial'
Christian Heimes78644762008-03-04 23:39:23 +000054 return prod(range(1, n+1))
55
Raymond Hettinger96ef8112003-02-01 00:10:11 +000056class TestBasicOps(unittest.TestCase):
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000057 def test_chain(self):
Christian Heimesdd15f6c2008-03-16 00:07:10 +000058
59 def chain2(*iterables):
60 'Pure python version in the docs'
61 for it in iterables:
62 for element in it:
63 yield element
64
65 for c in (chain, chain2):
66 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
67 self.assertEqual(list(c('abc')), list('abc'))
68 self.assertEqual(list(c('')), [])
69 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
70 self.assertRaises(TypeError, list,c(2, 3))
Christian Heimesf16baeb2008-02-29 14:57:44 +000071
72 def test_chain_from_iterable(self):
73 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
74 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
75 self.assertEqual(list(chain.from_iterable([''])), [])
76 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
77 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000078
Christian Heimes380f7f22008-02-28 11:19:05 +000079 def test_combinations(self):
Raymond Hettinger5bad41e2009-01-08 21:01:54 +000080 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Christian Heimes380f7f22008-02-28 11:19:05 +000081 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Christian Heimes78644762008-03-04 23:39:23 +000082 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Christian Heimes380f7f22008-02-28 11:19:05 +000083 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Raymond Hettinger5bad41e2009-01-08 21:01:54 +000084 self.assertEqual(list(combinations('abc', 32)), []) # r > n
Christian Heimes380f7f22008-02-28 11:19:05 +000085 self.assertEqual(list(combinations(range(4), 3)),
86 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
Christian Heimes78644762008-03-04 23:39:23 +000087
88 def combinations1(iterable, r):
89 'Pure python version shown in the docs'
90 pool = tuple(iterable)
91 n = len(pool)
Raymond Hettinger5bad41e2009-01-08 21:01:54 +000092 if r > n:
93 return
Christian Heimes78644762008-03-04 23:39:23 +000094 indices = list(range(r))
95 yield tuple(pool[i] for i in indices)
96 while 1:
97 for i in reversed(range(r)):
98 if indices[i] != i + n - r:
99 break
100 else:
101 return
102 indices[i] += 1
103 for j in range(i+1, r):
104 indices[j] = indices[j-1] + 1
105 yield tuple(pool[i] for i in indices)
106
107 def combinations2(iterable, r):
108 'Pure python version shown in the docs'
109 pool = tuple(iterable)
110 n = len(pool)
111 for indices in permutations(range(n), r):
112 if sorted(indices) == list(indices):
113 yield tuple(pool[i] for i in indices)
114
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000115 def combinations3(iterable, r):
116 'Pure python version from cwr()'
117 pool = tuple(iterable)
118 n = len(pool)
119 for indices in combinations_with_replacement(range(n), r):
120 if len(set(indices)) == r:
121 yield tuple(pool[i] for i in indices)
122
Christian Heimes78644762008-03-04 23:39:23 +0000123 for n in range(7):
Christian Heimes380f7f22008-02-28 11:19:05 +0000124 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000125 for r in range(n+2):
Christian Heimes380f7f22008-02-28 11:19:05 +0000126 result = list(combinations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000127 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 +0000128 self.assertEqual(len(result), len(set(result))) # no repeats
129 self.assertEqual(result, sorted(result)) # lexicographic order
130 for c in result:
131 self.assertEqual(len(c), r) # r-length combinations
132 self.assertEqual(len(set(c)), r) # no duplicate elements
133 self.assertEqual(list(c), sorted(c)) # keep original ordering
134 self.assert_(all(e in values for e in c)) # elements taken from input iterable
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000135 self.assertEqual(list(c),
136 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000137 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000138 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000139 self.assertEqual(result, list(combinations3(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000140
141 # Test implementation detail: tuple re-use
142 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
143 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
144
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000145 def test_combinations_with_replacement(self):
146 cwr = combinations_with_replacement
147 self.assertRaises(TypeError, cwr, 'abc') # missing r argument
148 self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
149 self.assertRaises(TypeError, cwr, None) # pool is not iterable
150 self.assertRaises(ValueError, cwr, 'abc', -2) # r is negative
151 self.assertEqual(list(cwr('ABC', 2)),
152 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
153
154 def cwr1(iterable, r):
155 'Pure python version shown in the docs'
156 # number items returned: (n+r-1)! / r! / (n-1)! when n>0
157 pool = tuple(iterable)
158 n = len(pool)
159 if not n and r:
160 return
161 indices = [0] * r
162 yield tuple(pool[i] for i in indices)
163 while 1:
164 for i in reversed(range(r)):
165 if indices[i] != n - 1:
166 break
167 else:
168 return
169 indices[i:] = [indices[i] + 1] * (r - i)
170 yield tuple(pool[i] for i in indices)
171
172 def cwr2(iterable, r):
173 'Pure python version shown in the docs'
174 pool = tuple(iterable)
175 n = len(pool)
176 for indices in product(range(n), repeat=r):
177 if sorted(indices) == list(indices):
178 yield tuple(pool[i] for i in indices)
179
180 def numcombs(n, r):
181 if not n:
182 return 0 if r else 1
183 return fact(n+r-1) / fact(r)/ fact(n-1)
184
185 for n in range(7):
186 values = [5*x-12 for x in range(n)]
187 for r in range(n+2):
188 result = list(cwr(values, r))
189
190 self.assertEqual(len(result), numcombs(n, r)) # right number of combs
191 self.assertEqual(len(result), len(set(result))) # no repeats
192 self.assertEqual(result, sorted(result)) # lexicographic order
193
194 regular_combs = list(combinations(values, r)) # compare to combs without replacement
195 if n == 0 or r <= 1:
196 self.assertEquals(result, regular_combs) # cases that should be identical
197 else:
198 self.assert_(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
199
200 for c in result:
201 self.assertEqual(len(c), r) # r-length combinations
202 noruns = [k for k,v in groupby(c)] # combo without consecutive repeats
203 self.assertEqual(len(noruns), len(set(noruns))) # no repeats other than consecutive
204 self.assertEqual(list(c), sorted(c)) # keep original ordering
205 self.assert_(all(e in values for e in c)) # elements taken from input iterable
206 self.assertEqual(noruns,
207 [e for e in values if e in c]) # comb is a subsequence of the input iterable
208 self.assertEqual(result, list(cwr1(values, r))) # matches first pure python version
209 self.assertEqual(result, list(cwr2(values, r))) # matches second pure python version
210
211 # Test implementation detail: tuple re-use
212 self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
213 self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
214
Christian Heimes78644762008-03-04 23:39:23 +0000215 def test_permutations(self):
216 self.assertRaises(TypeError, permutations) # too few arguments
217 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000218 self.assertRaises(TypeError, permutations, None) # pool is not iterable
219 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000220 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000221 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Christian Heimes78644762008-03-04 23:39:23 +0000222 self.assertEqual(list(permutations(range(3), 2)),
223 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
224
225 def permutations1(iterable, r=None):
226 'Pure python version shown in the docs'
227 pool = tuple(iterable)
228 n = len(pool)
229 r = n if r is None else r
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000230 if r > n:
231 return
Christian Heimes78644762008-03-04 23:39:23 +0000232 indices = list(range(n))
233 cycles = list(range(n-r+1, n+1))[::-1]
234 yield tuple(pool[i] for i in indices[:r])
235 while n:
236 for i in reversed(range(r)):
237 cycles[i] -= 1
238 if cycles[i] == 0:
239 indices[i:] = indices[i+1:] + indices[i:i+1]
240 cycles[i] = n - i
241 else:
242 j = cycles[i]
243 indices[i], indices[-j] = indices[-j], indices[i]
244 yield tuple(pool[i] for i in indices[:r])
245 break
246 else:
247 return
248
249 def permutations2(iterable, r=None):
250 'Pure python version shown in the docs'
251 pool = tuple(iterable)
252 n = len(pool)
253 r = n if r is None else r
254 for indices in product(range(n), repeat=r):
255 if len(set(indices)) == r:
256 yield tuple(pool[i] for i in indices)
257
258 for n in range(7):
259 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000260 for r in range(n+2):
Christian Heimes78644762008-03-04 23:39:23 +0000261 result = list(permutations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000262 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 +0000263 self.assertEqual(len(result), len(set(result))) # no repeats
264 self.assertEqual(result, sorted(result)) # lexicographic order
265 for p in result:
266 self.assertEqual(len(p), r) # r-length permutations
267 self.assertEqual(len(set(p)), r) # no duplicate elements
268 self.assert_(all(e in values for e in p)) # elements taken from input iterable
269 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000270 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000271 if r == n:
272 self.assertEqual(result, list(permutations(values, None))) # test r as None
273 self.assertEqual(result, list(permutations(values))) # test default r
274
275 # Test implementation detail: tuple re-use
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000276 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Christian Heimes78644762008-03-04 23:39:23 +0000277 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Christian Heimes380f7f22008-02-28 11:19:05 +0000278
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000279 def test_combinatorics(self):
280 # Test relationships between product(), permutations(),
281 # combinations() and combinations_with_replacement().
282
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000283 for n in range(6):
284 s = 'ABCDEFG'[:n]
285 for r in range(8):
286 prod = list(product(s, repeat=r))
287 cwr = list(combinations_with_replacement(s, r))
288 perm = list(permutations(s, r))
289 comb = list(combinations(s, r))
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000290
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000291 # Check size
292 self.assertEquals(len(prod), n**r)
293 self.assertEquals(len(cwr), (fact(n+r-1) / fact(r)/ fact(n-1)) if n else (not r))
294 self.assertEquals(len(perm), 0 if r>n else fact(n) / fact(n-r))
295 self.assertEquals(len(comb), 0 if r>n else fact(n) / fact(r) / fact(n-r))
296
297 # Check lexicographic order without repeated tuples
298 self.assertEquals(prod, sorted(set(prod)))
299 self.assertEquals(cwr, sorted(set(cwr)))
300 self.assertEquals(perm, sorted(set(perm)))
301 self.assertEquals(comb, sorted(set(comb)))
302
303 # Check interrelationships
304 self.assertEquals(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
305 self.assertEquals(perm, [t for t in prod if len(set(t))==r]) # perm: prods with no dups
306 self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
307 self.assertEqual(comb, [t for t in cwr if len(set(t))==r]) # comb: cwrs without dups
308 self.assertEqual(comb, list(filter(set(cwr).__contains__, perm))) # comb: perm that is a cwr
309 self.assertEqual(comb, list(filter(set(perm).__contains__, cwr))) # comb: cwr that is a perm
310 self.assertEqual(comb, sorted(set(cwr) & set(perm))) # comb: both a cwr and a perm
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000311
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000312 def test_compress(self):
313 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
314 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
315 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
316 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
317 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
318 n = 10000
319 data = chain.from_iterable(repeat(range(6), n))
320 selectors = chain.from_iterable(repeat((0, 1)))
321 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
322 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
323 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
324 self.assertRaises(TypeError, compress, range(6)) # too few args
325 self.assertRaises(TypeError, compress, range(6), None) # too many args
326
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000327 def test_count(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000328 self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
329 self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
330 self.assertEqual(take(2, lzip('abc',count(3))), [('a', 3), ('b', 4)])
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000331 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
332 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000333 self.assertRaises(TypeError, count, 2, 3, 4)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000334 self.assertRaises(TypeError, count, 'a')
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000335 self.assertEqual(list(islice(count(maxsize-5), 10)),
336 list(range(maxsize-5, maxsize+5)))
337 self.assertEqual(list(islice(count(-maxsize-5), 10)),
338 list(range(-maxsize-5, -maxsize+5)))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000339 c = count(3)
340 self.assertEqual(repr(c), 'count(3)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000341 next(c)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000342 self.assertEqual(repr(c), 'count(4)')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000343 c = count(-9)
344 self.assertEqual(repr(c), 'count(-9)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000345 next(c)
Raymond Hettinger30729212009-02-12 06:28:27 +0000346 self.assertEqual(repr(count(10.25)), 'count(10.25)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000347 self.assertEqual(next(c), -8)
Christian Heimesa37d4c62007-12-04 23:02:19 +0000348 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 +0000349 # Test repr (ignoring the L in longs)
350 r1 = repr(count(i)).replace('L', '')
351 r2 = 'count(%r)'.__mod__(i).replace('L', '')
352 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000353
Raymond Hettinger30729212009-02-12 06:28:27 +0000354 def test_count_with_stride(self):
355 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
356 self.assertEqual(lzip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
357 self.assertEqual(lzip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
358 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
359 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
360 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettinger8a882a72009-02-12 12:08:18 +0000361 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
362 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerde09acf2009-02-12 12:53:53 +0000363 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
364 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000365 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
366 c = count(3, 5)
367 self.assertEqual(repr(c), 'count(3, 5)')
368 next(c)
369 self.assertEqual(repr(c), 'count(8, 5)')
370 c = count(-9, 0)
371 self.assertEqual(repr(c), 'count(-9, 0)')
372 next(c)
373 self.assertEqual(repr(c), 'count(-9, 0)')
374 c = count(-9, -3)
375 self.assertEqual(repr(c), 'count(-9, -3)')
376 next(c)
377 self.assertEqual(repr(c), 'count(-12, -3)')
378 self.assertEqual(repr(c), 'count(-12, -3)')
379 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
380 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
381 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
382 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
383 for j in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 1, 10, sys.maxsize-5, sys.maxsize+5):
384 # Test repr (ignoring the L in longs)
385 r1 = repr(count(i, j)).replace('L', '')
386 if j == 1:
387 r2 = ('count(%r)' % i).replace('L', '')
388 else:
389 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
390 self.assertEqual(r1, r2)
391
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000392 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000393 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000394 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000395 self.assertRaises(TypeError, cycle)
396 self.assertRaises(TypeError, cycle, 5)
397 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000398
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000399 def test_groupby(self):
400 # Check whether it accepts arguments correctly
401 self.assertEqual([], list(groupby([])))
402 self.assertEqual([], list(groupby([], key=id)))
403 self.assertRaises(TypeError, list, groupby('abc', []))
404 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000405 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000406
407 # Check normal input
408 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
409 (2,15,22), (3,16,23), (3,17,23)]
410 dup = []
411 for k, g in groupby(s, lambda r:r[0]):
412 for elem in g:
413 self.assertEqual(k, elem[0])
414 dup.append(elem)
415 self.assertEqual(s, dup)
416
417 # Check nested case
418 dup = []
419 for k, g in groupby(s, lambda r:r[0]):
420 for ik, ig in groupby(g, lambda r:r[2]):
421 for elem in ig:
422 self.assertEqual(k, elem[0])
423 self.assertEqual(ik, elem[2])
424 dup.append(elem)
425 self.assertEqual(s, dup)
426
427 # Check case where inner iterator is not used
428 keys = [k for k, g in groupby(s, lambda r:r[0])]
429 expectedkeys = set([r[0] for r in s])
430 self.assertEqual(set(keys), expectedkeys)
431 self.assertEqual(len(keys), len(expectedkeys))
432
433 # Exercise pipes and filters style
434 s = 'abracadabra'
435 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000436 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000437 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
438 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000439 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000440 self.assertEqual(r, ['a', 'b', 'r'])
441 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000442 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000443 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
444 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000445 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000446 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
447
Georg Brandla18af4e2007-04-21 15:47:16 +0000448 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000449 class ExpectedError(Exception):
450 pass
451 def delayed_raise(n=0):
452 for i in range(n):
453 yield 'yo'
454 raise ExpectedError
455 def gulp(iterable, keyp=None, func=list):
456 return [func(g) for k, g in groupby(iterable, keyp)]
457
Georg Brandla18af4e2007-04-21 15:47:16 +0000458 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000459 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000460 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000461 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
462
463 # __cmp__ failure
464 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000465 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000466 raise ExpectedError
467 s = [DummyCmp(), DummyCmp(), None]
468
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000469 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000470 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000471 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000472 self.assertRaises(ExpectedError, gulp, s)
473
474 # keyfunc failure
475 def keyfunc(obj):
476 if keyfunc.skip > 0:
477 keyfunc.skip -= 1
478 return obj
479 else:
480 raise ExpectedError
481
482 # keyfunc failure on outer object
483 keyfunc.skip = 0
484 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
485 keyfunc.skip = 1
486 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
487
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000488 def test_filter(self):
489 self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
490 self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])
491 self.assertEqual(list(filter(bool, [0,1,0,2,0])), [1,2])
492 self.assertEqual(take(4, filter(isEven, count())), [0,2,4,6])
493 self.assertRaises(TypeError, filter)
494 self.assertRaises(TypeError, filter, lambda x:x)
495 self.assertRaises(TypeError, filter, lambda x:x, range(6), 7)
496 self.assertRaises(TypeError, filter, isEven, 3)
497 self.assertRaises(TypeError, next, filter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000498
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000499 def test_filterfalse(self):
500 self.assertEqual(list(filterfalse(isEven, range(6))), [1,3,5])
501 self.assertEqual(list(filterfalse(None, [0,1,0,2,0])), [0,0,0])
502 self.assertEqual(list(filterfalse(bool, [0,1,0,2,0])), [0,0,0])
503 self.assertEqual(take(4, filterfalse(isEven, count())), [1,3,5,7])
504 self.assertRaises(TypeError, filterfalse)
505 self.assertRaises(TypeError, filterfalse, lambda x:x)
506 self.assertRaises(TypeError, filterfalse, lambda x:x, range(6), 7)
507 self.assertRaises(TypeError, filterfalse, isEven, 3)
508 self.assertRaises(TypeError, next, filterfalse(range(6), range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000509
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000510 def test_zip(self):
511 # XXX This is rather silly now that builtin zip() calls zip()...
512 ans = [(x,y) for x, y in zip('abc',count())]
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000513 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000514 self.assertEqual(list(zip('abc', range(6))), lzip('abc', range(6)))
515 self.assertEqual(list(zip('abcdef', range(3))), lzip('abcdef', range(3)))
516 self.assertEqual(take(3,zip('abcdef', count())), lzip('abcdef', range(3)))
517 self.assertEqual(list(zip('abcdef')), lzip('abcdef'))
518 self.assertEqual(list(zip()), lzip())
519 self.assertRaises(TypeError, zip, 3)
520 self.assertRaises(TypeError, zip, range(3), 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000521 # Check tuple re-use (implementation detail)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000522 self.assertEqual([tuple(list(pair)) for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000523 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000524 self.assertEqual([pair for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000525 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000526 ids = list(map(id, zip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000527 self.assertEqual(min(ids), max(ids))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000528 ids = list(map(id, list(zip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000529 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000530
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000531 def test_ziplongest(self):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000532 for args in [
533 ['abc', range(6)],
534 [range(6), 'abc'],
535 [range(1000), range(2000,2100), range(3000,3050)],
536 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
537 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
538 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000539 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
540 for i in range(max(map(len, args)))]
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000541 self.assertEqual(list(zip_longest(*args)), target)
542 self.assertEqual(list(zip_longest(*args, **{})), target)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000543 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 +0000544 self.assertEqual(list(zip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000545
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000546 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 +0000547
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000548 self.assertEqual(list(zip_longest()), list(zip()))
549 self.assertEqual(list(zip_longest([])), list(zip([])))
550 self.assertEqual(list(zip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000551
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000552 self.assertEqual(list(zip_longest('abc', 'defg', **{})),
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000553 list(zip(list('abc')+[None], 'defg'))) # empty keyword dict
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000554 self.assertRaises(TypeError, zip_longest, 3)
555 self.assertRaises(TypeError, zip_longest, range(3), 3)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000556
557 for stmt in [
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000558 "zip_longest('abc', fv=1)",
559 "zip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000560 ]:
561 try:
562 eval(stmt, globals(), locals())
563 except TypeError:
564 pass
565 else:
566 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000567
Thomas Wouterscf297e42007-02-23 15:07:44 +0000568 # Check tuple re-use (implementation detail)
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000569 self.assertEqual([tuple(list(pair)) for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000570 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000571 self.assertEqual([pair for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000572 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000573 ids = list(map(id, zip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000574 self.assertEqual(min(ids), max(ids))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000575 ids = list(map(id, list(zip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000576 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
577
Christian Heimesc3f30c42008-02-22 16:37:40 +0000578 def test_product(self):
579 for args, result in [
Christian Heimes78644762008-03-04 23:39:23 +0000580 ([], [()]), # zero iterables
Christian Heimesc3f30c42008-02-22 16:37:40 +0000581 (['ab'], [('a',), ('b',)]), # one iterable
582 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
583 ([range(0), range(2), range(3)], []), # first iterable with zero length
584 ([range(2), range(0), range(3)], []), # middle iterable with zero length
585 ([range(2), range(3), range(0)], []), # last iterable with zero length
586 ]:
587 self.assertEqual(list(product(*args)), result)
Christian Heimesf16baeb2008-02-29 14:57:44 +0000588 for r in range(4):
589 self.assertEqual(list(product(*(args*r))),
590 list(product(*args, **dict(repeat=r))))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000591 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
592 self.assertRaises(TypeError, product, range(6), None)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000593
594 def product1(*args, **kwds):
595 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
596 n = len(pools)
597 if n == 0:
598 yield ()
599 return
600 if any(len(pool) == 0 for pool in pools):
601 return
602 indices = [0] * n
603 yield tuple(pool[i] for pool, i in zip(pools, indices))
604 while 1:
605 for i in reversed(range(n)): # right to left
606 if indices[i] == len(pools[i]) - 1:
607 continue
608 indices[i] += 1
609 for j in range(i+1, n):
610 indices[j] = 0
611 yield tuple(pool[i] for pool, i in zip(pools, indices))
612 break
613 else:
614 return
615
616 def product2(*args, **kwds):
617 'Pure python version used in docs'
618 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
619 result = [[]]
620 for pool in pools:
621 result = [x+[y] for x in result for y in pool]
622 for prod in result:
623 yield tuple(prod)
624
Christian Heimesc3f30c42008-02-22 16:37:40 +0000625 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
626 set('abcdefg'), range(11), tuple(range(13))]
627 for i in range(100):
628 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Christian Heimes78644762008-03-04 23:39:23 +0000629 expected_len = prod(map(len, args))
630 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000631 self.assertEqual(list(product(*args)), list(product1(*args)))
632 self.assertEqual(list(product(*args)), list(product2(*args)))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000633 args = map(iter, args)
Christian Heimes78644762008-03-04 23:39:23 +0000634 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000635
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000636 # Test implementation detail: tuple re-use
637 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
638 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000639
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000640 def test_repeat(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000641 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000642 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000643 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000644 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000645 self.assertEqual(list(repeat('a', 0)), [])
646 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000647 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000648 self.assertRaises(TypeError, repeat, None, 3, 4)
649 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000650 r = repeat(1+0j)
651 self.assertEqual(repr(r), 'repeat((1+0j))')
652 r = repeat(1+0j, 5)
653 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
654 list(r)
655 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000656
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000657 def test_map(self):
658 self.assertEqual(list(map(operator.pow, range(3), range(1,7))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000659 [0**1, 1**2, 2**3])
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000660 def tupleize(*args):
661 return args
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000662 self.assertEqual(list(map(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000663 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000664 self.assertEqual(list(map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000665 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000666 self.assertEqual(take(2,map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000667 [('a',0),('b',1)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000668 self.assertEqual(list(map(operator.pow, [])), [])
669 self.assertRaises(TypeError, map)
670 self.assertRaises(TypeError, list, map(None, range(3), range(3)))
671 self.assertRaises(TypeError, map, operator.neg)
672 self.assertRaises(TypeError, next, map(10, range(5)))
673 self.assertRaises(ValueError, next, map(errfunc, [4], [5]))
674 self.assertRaises(TypeError, next, map(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000675
676 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000677 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
678 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000679 self.assertEqual(take(3, starmap(operator.pow, zip(count(), count(1)))),
Raymond Hettinger02420702003-06-29 20:36:23 +0000680 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000681 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +0000682 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
683 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000684 self.assertRaises(TypeError, starmap)
685 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000686 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
687 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
688 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000689
690 def test_islice(self):
691 for args in [ # islice(args) should agree with range(args)
692 (10, 20, 3),
693 (10, 3, 20),
694 (10, 20),
695 (10, 3),
696 (20,)
697 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000698 self.assertEqual(list(islice(range(100), *args)),
699 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000700
701 for args, tgtargs in [ # Stop when seqn is exhausted
702 ((10, 110, 3), ((10, 100, 3))),
703 ((10, 110), ((10, 100))),
704 ((110,), (100,))
705 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000706 self.assertEqual(list(islice(range(100), *args)),
707 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000708
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000709 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +0000710 self.assertEqual(list(islice(range(10), None)), list(range(10)))
711 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
712 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
713 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
714 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000715
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000716 # Test number of items consumed SF #1171417
717 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +0000718 self.assertEqual(list(islice(it, 3)), list(range(3)))
719 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000720
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000721 # Test invalid arguments
Guido van Rossum805365e2007-05-07 22:24:25 +0000722 self.assertRaises(TypeError, islice, range(10))
723 self.assertRaises(TypeError, islice, range(10), 1, 2, 3, 4)
724 self.assertRaises(ValueError, islice, range(10), -5, 10, 1)
725 self.assertRaises(ValueError, islice, range(10), 1, -5, -1)
726 self.assertRaises(ValueError, islice, range(10), 1, 10, -1)
727 self.assertRaises(ValueError, islice, range(10), 1, 10, 0)
728 self.assertRaises(ValueError, islice, range(10), 'a')
729 self.assertRaises(ValueError, islice, range(10), 'a', 1)
730 self.assertRaises(ValueError, islice, range(10), 1, 'a')
731 self.assertRaises(ValueError, islice, range(10), 'a', 1, 1)
732 self.assertRaises(ValueError, islice, range(10), 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +0000733 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000734
735 def test_takewhile(self):
736 data = [1, 3, 5, 20, 2, 4, 6, 8]
737 underten = lambda x: x<10
738 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000739 self.assertEqual(list(takewhile(underten, [])), [])
740 self.assertRaises(TypeError, takewhile)
741 self.assertRaises(TypeError, takewhile, operator.pow)
742 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000743 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
744 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000745 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
746 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +0000747 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000748
749 def test_dropwhile(self):
750 data = [1, 3, 5, 20, 2, 4, 6, 8]
751 underten = lambda x: x<10
752 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000753 self.assertEqual(list(dropwhile(underten, [])), [])
754 self.assertRaises(TypeError, dropwhile)
755 self.assertRaises(TypeError, dropwhile, operator.pow)
756 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000757 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
758 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000759
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000760 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000761 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000762 def irange(n):
Guido van Rossum805365e2007-05-07 22:24:25 +0000763 for i in range(n):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000764 yield i
765
766 a, b = tee([]) # test empty iterator
767 self.assertEqual(list(a), [])
768 self.assertEqual(list(b), [])
769
770 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +0000771 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000772
773 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +0000774 self.assertEqual(list(a), list(range(n)))
775 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000776
777 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000778 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000779 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000780 del a
Guido van Rossum805365e2007-05-07 22:24:25 +0000781 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000782
783 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000784 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000785 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000786 del b
Guido van Rossum805365e2007-05-07 22:24:25 +0000787 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000788
Guido van Rossum805365e2007-05-07 22:24:25 +0000789 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000790 order = [0]*n + [1]*n
791 random.shuffle(order)
792 lists = ([], [])
793 its = tee(irange(n))
794 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +0000795 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000796 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +0000797 self.assertEqual(lists[0], list(range(n)))
798 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000799
Raymond Hettingerad983e72003-11-12 14:32:26 +0000800 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000801 self.assertRaises(TypeError, tee)
802 self.assertRaises(TypeError, tee, 3)
803 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000804 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000805
Raymond Hettingerad983e72003-11-12 14:32:26 +0000806 # tee object should be instantiable
807 a, b = tee('abc')
808 c = type(a)('def')
809 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000810
Raymond Hettingerad983e72003-11-12 14:32:26 +0000811 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +0000812 a, b, c = tee(range(2000), 3)
813 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000814 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +0000815 self.assertEqual(list(b), list(range(2000)))
816 self.assertEqual([next(c), next(c)], list(range(2)))
817 self.assertEqual(list(a), list(range(100,2000)))
818 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000819
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000820 # test values of n
821 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000822 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +0000823 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000824 result = tee('abc', n)
825 self.assertEqual(type(result), tuple)
826 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000827 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000828
Raymond Hettingerad983e72003-11-12 14:32:26 +0000829 # tee pass-through to copyable iterator
830 a, b = tee('abc')
831 c, d = tee(a)
832 self.assert_(a is c)
833
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000834 # test tee_new
835 t1, t2 = tee('abc')
836 tnew = type(t1)
837 self.assertRaises(TypeError, tnew)
838 self.assertRaises(TypeError, tnew, 10)
839 t3 = tnew(t1)
840 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000841
Raymond Hettingera9f60922004-10-17 16:40:14 +0000842 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +0000843 a, b = tee(range(10))
Raymond Hettingera9f60922004-10-17 16:40:14 +0000844 p = proxy(a)
845 self.assertEqual(getattr(p, '__class__'), type(b))
846 del a
847 self.assertRaises(ReferenceError, getattr, p, '__class__')
848
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000849 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000850 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000851
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000852 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +0000853 self.assertRaises(StopIteration, next, f([]))
854 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000855
Georg Brandla18af4e2007-04-21 15:47:16 +0000856 self.assertRaises(StopIteration, next, islice([], None))
857 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000858
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000859 p, q = tee([])
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 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +0000863 self.assertRaises(StopIteration, next, p)
864 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000865
Georg Brandla18af4e2007-04-21 15:47:16 +0000866 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000867
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000868 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +0000869 self.assertRaises(StopIteration, next, f(lambda x:x, []))
870 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000871
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000872class TestExamples(unittest.TestCase):
873
874 def test_chain(self):
875 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
876
877 def test_chain_from_iterable(self):
878 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
879
880 def test_combinations(self):
881 self.assertEqual(list(combinations('ABCD', 2)),
882 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
883 self.assertEqual(list(combinations(range(4), 3)),
884 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
885
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000886 def test_combinations_with_replacement(self):
887 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
888 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
889
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000890 def test_compress(self):
891 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
892
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000893 def test_count(self):
894 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
895
896 def test_cycle(self):
897 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
898
899 def test_dropwhile(self):
900 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
901
902 def test_groupby(self):
903 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
904 list('ABCDAB'))
905 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
906 [list('AAAA'), list('BBB'), list('CC'), list('D')])
907
908 def test_filter(self):
909 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
910
911 def test_filterfalse(self):
912 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
913
914 def test_map(self):
915 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
916
917 def test_islice(self):
918 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
919 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
920 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
921 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
922
923 def test_zip(self):
924 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
925
926 def test_zip_longest(self):
927 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
928 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
929
930 def test_permutations(self):
931 self.assertEqual(list(permutations('ABCD', 2)),
932 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
933 self.assertEqual(list(permutations(range(3))),
934 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
935
936 def test_product(self):
937 self.assertEqual(list(product('ABCD', 'xy')),
938 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
939 self.assertEqual(list(product(range(2), repeat=3)),
940 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
941 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
942
943 def test_repeat(self):
944 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
945
946 def test_stapmap(self):
947 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
948 [32, 9, 1000])
949
950 def test_takewhile(self):
951 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
952
953
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000954class TestGC(unittest.TestCase):
955
956 def makecycle(self, iterator, container):
957 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +0000958 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000959 del container, iterator
960
961 def test_chain(self):
962 a = []
963 self.makecycle(chain(a), a)
964
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000965 def test_chain_from_iterable(self):
966 a = []
967 self.makecycle(chain.from_iterable([a]), a)
968
969 def test_combinations(self):
970 a = []
971 self.makecycle(combinations([1,2,a,3], 3), a)
972
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000973 def test_combinations_with_replacement(self):
974 a = []
975 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
976
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000977 def test_compress(self):
978 a = []
979 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
980
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000981 def test_cycle(self):
982 a = []
983 self.makecycle(cycle([a]*2), a)
984
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000985 def test_dropwhile(self):
986 a = []
987 self.makecycle(dropwhile(bool, [0, a, a]), a)
988
989 def test_groupby(self):
990 a = []
991 self.makecycle(groupby([a]*2, lambda x:x), a)
992
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000993 def test_issue2246(self):
994 # Issue 2246 -- the _grouper iterator was not included in GC
995 n = 10
996 keyfunc = lambda x: x
997 for i, j in groupby(range(n), key=keyfunc):
998 keyfunc.__dict__.setdefault('x',[]).append(j)
999
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001000 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001001 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001002 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001003
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001004 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001005 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001006 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001007
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001008 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001009 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001010 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001011
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001012 def test_zip_longest(self):
1013 a = []
1014 self.makecycle(zip_longest([a]*2, [a]*3), a)
1015 b = [a, None]
1016 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1017
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001018 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001019 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001020 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001021
1022 def test_islice(self):
1023 a = []
1024 self.makecycle(islice([a]*2, None), a)
1025
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001026 def test_permutations(self):
1027 a = []
1028 self.makecycle(permutations([1,2,a,3], 3), a)
1029
1030 def test_product(self):
1031 a = []
1032 self.makecycle(product([1,2,a,3], repeat=3), a)
1033
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001034 def test_repeat(self):
1035 a = []
1036 self.makecycle(repeat(a), a)
1037
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001038 def test_starmap(self):
1039 a = []
1040 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1041
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001042 def test_takewhile(self):
1043 a = []
1044 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1045
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001046def R(seqn):
1047 'Regular generator'
1048 for i in seqn:
1049 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001050
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001051class G:
1052 'Sequence using __getitem__'
1053 def __init__(self, seqn):
1054 self.seqn = seqn
1055 def __getitem__(self, i):
1056 return self.seqn[i]
1057
1058class I:
1059 'Sequence using iterator protocol'
1060 def __init__(self, seqn):
1061 self.seqn = seqn
1062 self.i = 0
1063 def __iter__(self):
1064 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001065 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001066 if self.i >= len(self.seqn): raise StopIteration
1067 v = self.seqn[self.i]
1068 self.i += 1
1069 return v
1070
1071class Ig:
1072 'Sequence using iterator protocol defined with a generator'
1073 def __init__(self, seqn):
1074 self.seqn = seqn
1075 self.i = 0
1076 def __iter__(self):
1077 for val in self.seqn:
1078 yield val
1079
1080class X:
1081 'Missing __getitem__ and __iter__'
1082 def __init__(self, seqn):
1083 self.seqn = seqn
1084 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001085 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001086 if self.i >= len(self.seqn): raise StopIteration
1087 v = self.seqn[self.i]
1088 self.i += 1
1089 return v
1090
1091class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001092 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001093 def __init__(self, seqn):
1094 self.seqn = seqn
1095 self.i = 0
1096 def __iter__(self):
1097 return self
1098
1099class E:
1100 'Test propagation of exceptions'
1101 def __init__(self, seqn):
1102 self.seqn = seqn
1103 self.i = 0
1104 def __iter__(self):
1105 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001106 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001107 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001108
1109class S:
1110 'Test immediate stop'
1111 def __init__(self, seqn):
1112 pass
1113 def __iter__(self):
1114 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001115 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001116 raise StopIteration
1117
1118def L(seqn):
1119 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001120 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001121
1122
1123class TestVariousIteratorArgs(unittest.TestCase):
1124
1125 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001126 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001127 for g in (G, I, Ig, S, L, R):
1128 self.assertEqual(list(chain(g(s))), list(g(s)))
1129 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001130 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001131 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001132 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1133
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001134 def test_compress(self):
1135 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1136 n = len(s)
1137 for g in (G, I, Ig, S, L, R):
1138 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1139 self.assertRaises(TypeError, compress, X(s), repeat(1))
1140 self.assertRaises(TypeError, compress, N(s), repeat(1))
1141 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1142
Christian Heimesc3f30c42008-02-22 16:37:40 +00001143 def test_product(self):
1144 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1145 self.assertRaises(TypeError, product, X(s))
1146 self.assertRaises(TypeError, product, N(s))
1147 self.assertRaises(ZeroDivisionError, product, E(s))
1148
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001149 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001150 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001151 for g in (G, I, Ig, S, L, R):
1152 tgtlen = len(s) * 3
1153 expected = list(g(s))*3
1154 actual = list(islice(cycle(g(s)), tgtlen))
1155 self.assertEqual(actual, expected)
1156 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001157 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001158 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1159
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001160 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001161 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001162 for g in (G, I, Ig, S, L, R):
1163 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1164 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001165 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001166 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1167
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001168 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001169 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001170 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001171 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001172 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001173 self.assertRaises(TypeError, filter, isEven, X(s))
1174 self.assertRaises(TypeError, filter, isEven, N(s))
1175 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001176
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001177 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001178 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001179 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001180 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001181 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001182 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1183 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1184 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001185
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001186 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001187 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001188 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001189 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1190 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1191 self.assertRaises(TypeError, zip, X(s))
1192 self.assertRaises(TypeError, zip, N(s))
1193 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001194
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001195 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001196 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001197 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001198 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1199 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1200 self.assertRaises(TypeError, zip_longest, X(s))
1201 self.assertRaises(TypeError, zip_longest, N(s))
1202 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001203
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001204 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001205 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001206 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001207 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001208 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001209 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001210 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001211 self.assertRaises(TypeError, map, onearg, X(s))
1212 self.assertRaises(TypeError, map, onearg, N(s))
1213 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001214
1215 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001216 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001217 for g in (G, I, Ig, S, L, R):
1218 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1219 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001220 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001221 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1222
1223 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001224 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001225 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001226 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001227 self.assertEqual(list(starmap(operator.pow, g(ss))),
1228 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001229 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001230 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001231 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1232
1233 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001234 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001235 for g in (G, I, Ig, S, L, R):
1236 tgt = []
1237 for elem in g(s):
1238 if not isEven(elem): break
1239 tgt.append(elem)
1240 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1241 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001242 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001243 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1244
1245 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001246 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001247 for g in (G, I, Ig, S, L, R):
1248 tgt = []
1249 for elem in g(s):
1250 if not tgt and isOdd(elem): continue
1251 tgt.append(elem)
1252 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1253 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001254 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001255 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1256
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001257 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001258 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001259 for g in (G, I, Ig, S, L, R):
1260 it1, it2 = tee(g(s))
1261 self.assertEqual(list(it1), list(g(s)))
1262 self.assertEqual(list(it2), list(g(s)))
1263 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001264 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001265 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1266
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001267class LengthTransparency(unittest.TestCase):
1268
1269 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001270 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001271 self.assertEqual(len(repeat(None, 50)), 50)
1272 self.assertRaises(TypeError, len, repeat(None))
1273
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001274class RegressionTests(unittest.TestCase):
1275
1276 def test_sf_793826(self):
1277 # Fix Armin Rigo's successful efforts to wreak havoc
1278
1279 def mutatingtuple(tuple1, f, tuple2):
1280 # this builds a tuple t which is a copy of tuple1,
1281 # then calls f(t), then mutates t to be equal to tuple2
1282 # (needs len(tuple1) == len(tuple2)).
1283 def g(value, first=[1]):
1284 if first:
1285 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00001286 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001287 return value
1288 items = list(tuple2)
1289 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001290 gen = map(g, items)
1291 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00001292 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001293
1294 def f(t):
1295 global T
1296 T = t
1297 first[:] = list(T)
1298
1299 first = []
1300 mutatingtuple((1,2,3), f, (4,5,6))
1301 second = list(T)
1302 self.assertEqual(first, second)
1303
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001304
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001305 def test_sf_950057(self):
1306 # Make sure that chain() and cycle() catch exceptions immediately
1307 # rather than when shifting between input sources
1308
1309 def gen1():
1310 hist.append(0)
1311 yield 1
1312 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001313 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001314 hist.append(2)
1315
1316 def gen2(x):
1317 hist.append(3)
1318 yield 2
1319 hist.append(4)
1320 if x:
1321 raise StopIteration
1322
1323 hist = []
1324 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1325 self.assertEqual(hist, [0,1])
1326
1327 hist = []
1328 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1329 self.assertEqual(hist, [0,1])
1330
1331 hist = []
1332 self.assertRaises(AssertionError, list, cycle(gen1()))
1333 self.assertEqual(hist, [0,1])
1334
Thomas Woutersb2137042007-02-01 18:02:27 +00001335class SubclassWithKwargsTest(unittest.TestCase):
1336 def test_keywords_in_subclass(self):
1337 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001338 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001339 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00001340 class Subclass(cls):
1341 def __init__(self, newarg=None, *args):
1342 cls.__init__(self, *args)
1343 try:
1344 Subclass(newarg=1)
1345 except TypeError as err:
1346 # we expect type errors because of wrong argument count
1347 self.failIf("does not take keyword arguments" in err.args[0])
1348
1349
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001350libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001351
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001352
1353>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001354>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00001355... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001356...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001357Check 1200 is for $120.15
1358Check 1201 is for $764.05
1359Check 1202 is for $823.14
1360
1361>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001362>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00001363... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00001364...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000013651
13668
136727
1368
1369>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001370>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00001371... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00001372...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001373Alex
1374Laura
1375Martin
1376Walter
1377Samuele
1378
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001379>>> from operator import itemgetter
1380>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00001381>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001382>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001383... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001384...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000013851 ['a', 'c', 'e']
13862 ['b', 'd', 'f']
13873 ['g']
1388
Raymond Hettinger734fb572004-01-20 20:04:40 +00001389# Find runs of consecutive numbers using groupby. The key to the solution
1390# is differencing with a range so that consecutive numbers all appear in
1391# same group.
1392>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00001393>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001394... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001395...
Raymond Hettinger734fb572004-01-20 20:04:40 +00001396[1]
1397[4, 5, 6]
1398[10]
1399[15, 16, 17, 18]
1400[22]
1401[25, 26, 27, 28]
1402
Georg Brandl3dbca812008-07-23 16:10:53 +00001403>>> def take(n, iterable):
1404... "Return first n items of the iterable as a list"
1405... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001406
Georg Brandl3dbca812008-07-23 16:10:53 +00001407>>> def enumerate(iterable, start=0):
1408... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001409
Georg Brandl3dbca812008-07-23 16:10:53 +00001410>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001411... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00001412... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001413
1414>>> def nth(iterable, n):
Raymond Hettingerd04fa312009-02-04 19:45:13 +00001415... "Returns the nth item or None"
1416... return next(islice(iterable, n, None), None)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001417
Georg Brandl3dbca812008-07-23 16:10:53 +00001418>>> def quantify(iterable, pred=bool):
1419... "Count how many times the predicate is true"
1420... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001421
Georg Brandl3dbca812008-07-23 16:10:53 +00001422>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001423... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00001424... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001425
Georg Brandl3dbca812008-07-23 16:10:53 +00001426>>> def ncycles(iterable, n):
1427... "Returns the seqeuence elements n times"
1428... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001429
1430>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001431... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001432
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001433>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001434... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001435
1436>>> def repeatfunc(func, times=None, *args):
1437... "Repeat calls to func with specified arguments."
1438... " Example: repeatfunc(random.random)"
1439... if times is None:
1440... return starmap(func, repeat(args))
1441... else:
1442... return starmap(func, repeat(args, times))
1443
Raymond Hettingerd591f662003-10-26 15:34:50 +00001444>>> def pairwise(iterable):
1445... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1446... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001447... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00001448... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001449... except StopIteration:
1450... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001451... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001452
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001453>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001454... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001455... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00001456... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001457
1458>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001459... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001460... # Recipe credited to George Sakkis
1461... pending = len(iterables)
1462... nexts = cycle(iter(it).__next__ for it in iterables)
1463... while pending:
1464... try:
1465... for next in nexts:
1466... yield next()
1467... except StopIteration:
1468... pending -= 1
1469... nexts = cycle(islice(nexts, pending))
1470
1471>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00001472... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1473... s = list(iterable)
1474... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001475
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001476>>> def unique_everseen(iterable, key=None):
1477... "List unique elements, preserving order. Remember all elements ever seen."
1478... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1479... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1480... seen = set()
1481... seen_add = seen.add
1482... if key is None:
1483... for element in iterable:
1484... if element not in seen:
1485... seen_add(element)
1486... yield element
1487... else:
1488... for element in iterable:
1489... k = key(element)
1490... if k not in seen:
1491... seen_add(k)
1492... yield element
1493
1494>>> def unique_justseen(iterable, key=None):
1495... "List unique elements, preserving order. Remember only the element just seen."
1496... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1497... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1498... return map(next, map(itemgetter(1), groupby(iterable, key)))
1499
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001500This is not part of the examples but it tests to make sure the definitions
1501perform as purported.
1502
Raymond Hettingera098b332003-09-08 23:58:40 +00001503>>> take(10, count())
1504[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1505
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001506>>> list(enumerate('abc'))
1507[(0, 'a'), (1, 'b'), (2, 'c')]
1508
1509>>> list(islice(tabulate(lambda x: 2*x), 4))
1510[0, 2, 4, 6]
1511
1512>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00001513'd'
1514
1515>>> nth('abcde', 9) is None
1516True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001517
Guido van Rossum805365e2007-05-07 22:24:25 +00001518>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000151950
1520
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001521>>> a = [[1, 2, 3], [4, 5, 6]]
1522>>> flatten(a)
1523[1, 2, 3, 4, 5, 6]
1524
1525>>> list(repeatfunc(pow, 5, 2, 3))
1526[8, 8, 8, 8, 8]
1527
1528>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001529>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001530[0, 0, 0, 0, 0]
1531
Raymond Hettingerd591f662003-10-26 15:34:50 +00001532>>> list(pairwise('abcd'))
1533[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001534
Raymond Hettingerd591f662003-10-26 15:34:50 +00001535>>> list(pairwise([]))
1536[]
1537
1538>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001539[]
1540
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001541>>> list(islice(padnone('abc'), 0, 6))
1542['a', 'b', 'c', None, None, None]
1543
1544>>> list(ncycles('abc', 3))
1545['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1546
1547>>> dotproduct([1,2,3], [4,5,6])
154832
1549
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001550>>> list(grouper(3, 'abcdefg', 'x'))
1551[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1552
1553>>> list(roundrobin('abc', 'd', 'ef'))
1554['a', 'd', 'e', 'b', 'f', 'c']
1555
Raymond Hettingerace67332009-01-26 02:23:50 +00001556>>> list(powerset([1,2,3]))
1557[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001558
Raymond Hettinger191e8502009-01-27 13:29:43 +00001559>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1560True
1561
1562>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1563True
1564
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001565>>> list(unique_everseen('AAAABBBCCDAABBB'))
1566['A', 'B', 'C', 'D']
1567
1568>>> list(unique_everseen('ABBCcAD', str.lower))
1569['A', 'B', 'C', 'D']
1570
1571>>> list(unique_justseen('AAAABBBCCDAABBB'))
1572['A', 'B', 'C', 'D', 'A', 'B']
1573
1574>>> list(unique_justseen('ABBCcAD', str.lower))
1575['A', 'B', 'C', 'A', 'D']
1576
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001577"""
1578
1579__test__ = {'libreftest' : libreftest}
1580
1581def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001582 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00001583 RegressionTests, LengthTransparency,
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001584 SubclassWithKwargsTest, TestExamples)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001585 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001586
1587 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001588 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001589 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001590 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00001591 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001592 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001593 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001594 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001595 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001596
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001597 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001598 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001599
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001600if __name__ == "__main__":
1601 test_main(verbose=True)