blob: c34b915cebf98e4faa80667d8a90afbe3c38a8cd [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)])
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000358 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000359 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
360 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
361 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettinger8a882a72009-02-12 12:08:18 +0000362 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
363 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerde09acf2009-02-12 12:53:53 +0000364 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
365 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000366 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
367 c = count(3, 5)
368 self.assertEqual(repr(c), 'count(3, 5)')
369 next(c)
370 self.assertEqual(repr(c), 'count(8, 5)')
371 c = count(-9, 0)
372 self.assertEqual(repr(c), 'count(-9, 0)')
373 next(c)
374 self.assertEqual(repr(c), 'count(-9, 0)')
375 c = count(-9, -3)
376 self.assertEqual(repr(c), 'count(-9, -3)')
377 next(c)
378 self.assertEqual(repr(c), 'count(-12, -3)')
379 self.assertEqual(repr(c), 'count(-12, -3)')
380 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
381 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
382 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
383 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
384 for j in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 1, 10, sys.maxsize-5, sys.maxsize+5):
385 # Test repr (ignoring the L in longs)
386 r1 = repr(count(i, j)).replace('L', '')
387 if j == 1:
388 r2 = ('count(%r)' % i).replace('L', '')
389 else:
390 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
391 self.assertEqual(r1, r2)
392
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000393 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000394 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000395 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000396 self.assertRaises(TypeError, cycle)
397 self.assertRaises(TypeError, cycle, 5)
398 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000399
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000400 def test_groupby(self):
401 # Check whether it accepts arguments correctly
402 self.assertEqual([], list(groupby([])))
403 self.assertEqual([], list(groupby([], key=id)))
404 self.assertRaises(TypeError, list, groupby('abc', []))
405 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000406 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000407
408 # Check normal input
409 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
410 (2,15,22), (3,16,23), (3,17,23)]
411 dup = []
412 for k, g in groupby(s, lambda r:r[0]):
413 for elem in g:
414 self.assertEqual(k, elem[0])
415 dup.append(elem)
416 self.assertEqual(s, dup)
417
418 # Check nested case
419 dup = []
420 for k, g in groupby(s, lambda r:r[0]):
421 for ik, ig in groupby(g, lambda r:r[2]):
422 for elem in ig:
423 self.assertEqual(k, elem[0])
424 self.assertEqual(ik, elem[2])
425 dup.append(elem)
426 self.assertEqual(s, dup)
427
428 # Check case where inner iterator is not used
429 keys = [k for k, g in groupby(s, lambda r:r[0])]
430 expectedkeys = set([r[0] for r in s])
431 self.assertEqual(set(keys), expectedkeys)
432 self.assertEqual(len(keys), len(expectedkeys))
433
434 # Exercise pipes and filters style
435 s = 'abracadabra'
436 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000437 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000438 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
439 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000440 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000441 self.assertEqual(r, ['a', 'b', 'r'])
442 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000443 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000444 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
445 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000446 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000447 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
448
Georg Brandla18af4e2007-04-21 15:47:16 +0000449 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000450 class ExpectedError(Exception):
451 pass
452 def delayed_raise(n=0):
453 for i in range(n):
454 yield 'yo'
455 raise ExpectedError
456 def gulp(iterable, keyp=None, func=list):
457 return [func(g) for k, g in groupby(iterable, keyp)]
458
Georg Brandla18af4e2007-04-21 15:47:16 +0000459 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000460 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000461 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000462 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
463
464 # __cmp__ failure
465 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000466 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000467 raise ExpectedError
468 s = [DummyCmp(), DummyCmp(), None]
469
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000470 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000471 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000472 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000473 self.assertRaises(ExpectedError, gulp, s)
474
475 # keyfunc failure
476 def keyfunc(obj):
477 if keyfunc.skip > 0:
478 keyfunc.skip -= 1
479 return obj
480 else:
481 raise ExpectedError
482
483 # keyfunc failure on outer object
484 keyfunc.skip = 0
485 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
486 keyfunc.skip = 1
487 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
488
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000489 def test_filter(self):
490 self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
491 self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])
492 self.assertEqual(list(filter(bool, [0,1,0,2,0])), [1,2])
493 self.assertEqual(take(4, filter(isEven, count())), [0,2,4,6])
494 self.assertRaises(TypeError, filter)
495 self.assertRaises(TypeError, filter, lambda x:x)
496 self.assertRaises(TypeError, filter, lambda x:x, range(6), 7)
497 self.assertRaises(TypeError, filter, isEven, 3)
498 self.assertRaises(TypeError, next, filter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000499
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000500 def test_filterfalse(self):
501 self.assertEqual(list(filterfalse(isEven, range(6))), [1,3,5])
502 self.assertEqual(list(filterfalse(None, [0,1,0,2,0])), [0,0,0])
503 self.assertEqual(list(filterfalse(bool, [0,1,0,2,0])), [0,0,0])
504 self.assertEqual(take(4, filterfalse(isEven, count())), [1,3,5,7])
505 self.assertRaises(TypeError, filterfalse)
506 self.assertRaises(TypeError, filterfalse, lambda x:x)
507 self.assertRaises(TypeError, filterfalse, lambda x:x, range(6), 7)
508 self.assertRaises(TypeError, filterfalse, isEven, 3)
509 self.assertRaises(TypeError, next, filterfalse(range(6), range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000510
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000511 def test_zip(self):
512 # XXX This is rather silly now that builtin zip() calls zip()...
513 ans = [(x,y) for x, y in zip('abc',count())]
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000514 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000515 self.assertEqual(list(zip('abc', range(6))), lzip('abc', range(6)))
516 self.assertEqual(list(zip('abcdef', range(3))), lzip('abcdef', range(3)))
517 self.assertEqual(take(3,zip('abcdef', count())), lzip('abcdef', range(3)))
518 self.assertEqual(list(zip('abcdef')), lzip('abcdef'))
519 self.assertEqual(list(zip()), lzip())
520 self.assertRaises(TypeError, zip, 3)
521 self.assertRaises(TypeError, zip, range(3), 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000522 # Check tuple re-use (implementation detail)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000523 self.assertEqual([tuple(list(pair)) for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000524 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000525 self.assertEqual([pair for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000526 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000527 ids = list(map(id, zip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000528 self.assertEqual(min(ids), max(ids))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000529 ids = list(map(id, list(zip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000530 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000531
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000532 def test_ziplongest(self):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000533 for args in [
534 ['abc', range(6)],
535 [range(6), 'abc'],
536 [range(1000), range(2000,2100), range(3000,3050)],
537 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
538 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
539 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000540 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
541 for i in range(max(map(len, args)))]
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000542 self.assertEqual(list(zip_longest(*args)), target)
543 self.assertEqual(list(zip_longest(*args, **{})), target)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000544 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 +0000545 self.assertEqual(list(zip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000546
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000547 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 +0000548
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000549 self.assertEqual(list(zip_longest()), list(zip()))
550 self.assertEqual(list(zip_longest([])), list(zip([])))
551 self.assertEqual(list(zip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000552
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000553 self.assertEqual(list(zip_longest('abc', 'defg', **{})),
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000554 list(zip(list('abc')+[None], 'defg'))) # empty keyword dict
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000555 self.assertRaises(TypeError, zip_longest, 3)
556 self.assertRaises(TypeError, zip_longest, range(3), 3)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000557
558 for stmt in [
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000559 "zip_longest('abc', fv=1)",
560 "zip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000561 ]:
562 try:
563 eval(stmt, globals(), locals())
564 except TypeError:
565 pass
566 else:
567 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000568
Thomas Wouterscf297e42007-02-23 15:07:44 +0000569 # Check tuple re-use (implementation detail)
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000570 self.assertEqual([tuple(list(pair)) for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000571 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000572 self.assertEqual([pair for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000573 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000574 ids = list(map(id, zip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000575 self.assertEqual(min(ids), max(ids))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000576 ids = list(map(id, list(zip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000577 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
578
Christian Heimesc3f30c42008-02-22 16:37:40 +0000579 def test_product(self):
580 for args, result in [
Christian Heimes78644762008-03-04 23:39:23 +0000581 ([], [()]), # zero iterables
Christian Heimesc3f30c42008-02-22 16:37:40 +0000582 (['ab'], [('a',), ('b',)]), # one iterable
583 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
584 ([range(0), range(2), range(3)], []), # first iterable with zero length
585 ([range(2), range(0), range(3)], []), # middle iterable with zero length
586 ([range(2), range(3), range(0)], []), # last iterable with zero length
587 ]:
588 self.assertEqual(list(product(*args)), result)
Christian Heimesf16baeb2008-02-29 14:57:44 +0000589 for r in range(4):
590 self.assertEqual(list(product(*(args*r))),
591 list(product(*args, **dict(repeat=r))))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000592 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
593 self.assertRaises(TypeError, product, range(6), None)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000594
595 def product1(*args, **kwds):
596 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
597 n = len(pools)
598 if n == 0:
599 yield ()
600 return
601 if any(len(pool) == 0 for pool in pools):
602 return
603 indices = [0] * n
604 yield tuple(pool[i] for pool, i in zip(pools, indices))
605 while 1:
606 for i in reversed(range(n)): # right to left
607 if indices[i] == len(pools[i]) - 1:
608 continue
609 indices[i] += 1
610 for j in range(i+1, n):
611 indices[j] = 0
612 yield tuple(pool[i] for pool, i in zip(pools, indices))
613 break
614 else:
615 return
616
617 def product2(*args, **kwds):
618 'Pure python version used in docs'
619 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
620 result = [[]]
621 for pool in pools:
622 result = [x+[y] for x in result for y in pool]
623 for prod in result:
624 yield tuple(prod)
625
Christian Heimesc3f30c42008-02-22 16:37:40 +0000626 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
627 set('abcdefg'), range(11), tuple(range(13))]
628 for i in range(100):
629 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Christian Heimes78644762008-03-04 23:39:23 +0000630 expected_len = prod(map(len, args))
631 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000632 self.assertEqual(list(product(*args)), list(product1(*args)))
633 self.assertEqual(list(product(*args)), list(product2(*args)))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000634 args = map(iter, args)
Christian Heimes78644762008-03-04 23:39:23 +0000635 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000636
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000637 # Test implementation detail: tuple re-use
638 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
639 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000640
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000641 def test_repeat(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000642 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000643 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000644 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000645 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000646 self.assertEqual(list(repeat('a', 0)), [])
647 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000648 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000649 self.assertRaises(TypeError, repeat, None, 3, 4)
650 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000651 r = repeat(1+0j)
652 self.assertEqual(repr(r), 'repeat((1+0j))')
653 r = repeat(1+0j, 5)
654 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
655 list(r)
656 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000657
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000658 def test_map(self):
659 self.assertEqual(list(map(operator.pow, range(3), range(1,7))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000660 [0**1, 1**2, 2**3])
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000661 def tupleize(*args):
662 return args
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000663 self.assertEqual(list(map(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000664 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000665 self.assertEqual(list(map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000666 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000667 self.assertEqual(take(2,map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000668 [('a',0),('b',1)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000669 self.assertEqual(list(map(operator.pow, [])), [])
670 self.assertRaises(TypeError, map)
671 self.assertRaises(TypeError, list, map(None, range(3), range(3)))
672 self.assertRaises(TypeError, map, operator.neg)
673 self.assertRaises(TypeError, next, map(10, range(5)))
674 self.assertRaises(ValueError, next, map(errfunc, [4], [5]))
675 self.assertRaises(TypeError, next, map(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000676
677 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000678 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
679 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000680 self.assertEqual(take(3, starmap(operator.pow, zip(count(), count(1)))),
Raymond Hettinger02420702003-06-29 20:36:23 +0000681 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000682 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +0000683 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
684 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000685 self.assertRaises(TypeError, starmap)
686 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000687 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
688 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
689 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000690
691 def test_islice(self):
692 for args in [ # islice(args) should agree with range(args)
693 (10, 20, 3),
694 (10, 3, 20),
695 (10, 20),
696 (10, 3),
697 (20,)
698 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000699 self.assertEqual(list(islice(range(100), *args)),
700 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000701
702 for args, tgtargs in [ # Stop when seqn is exhausted
703 ((10, 110, 3), ((10, 100, 3))),
704 ((10, 110), ((10, 100))),
705 ((110,), (100,))
706 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000707 self.assertEqual(list(islice(range(100), *args)),
708 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000709
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000710 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +0000711 self.assertEqual(list(islice(range(10), None)), list(range(10)))
712 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
713 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
714 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
715 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000716
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000717 # Test number of items consumed SF #1171417
718 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +0000719 self.assertEqual(list(islice(it, 3)), list(range(3)))
720 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000721
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000722 # Test invalid arguments
Guido van Rossum805365e2007-05-07 22:24:25 +0000723 self.assertRaises(TypeError, islice, range(10))
724 self.assertRaises(TypeError, islice, range(10), 1, 2, 3, 4)
725 self.assertRaises(ValueError, islice, range(10), -5, 10, 1)
726 self.assertRaises(ValueError, islice, range(10), 1, -5, -1)
727 self.assertRaises(ValueError, islice, range(10), 1, 10, -1)
728 self.assertRaises(ValueError, islice, range(10), 1, 10, 0)
729 self.assertRaises(ValueError, islice, range(10), 'a')
730 self.assertRaises(ValueError, islice, range(10), 'a', 1)
731 self.assertRaises(ValueError, islice, range(10), 1, 'a')
732 self.assertRaises(ValueError, islice, range(10), 'a', 1, 1)
733 self.assertRaises(ValueError, islice, range(10), 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +0000734 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000735
736 def test_takewhile(self):
737 data = [1, 3, 5, 20, 2, 4, 6, 8]
738 underten = lambda x: x<10
739 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000740 self.assertEqual(list(takewhile(underten, [])), [])
741 self.assertRaises(TypeError, takewhile)
742 self.assertRaises(TypeError, takewhile, operator.pow)
743 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000744 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
745 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000746 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
747 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +0000748 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000749
750 def test_dropwhile(self):
751 data = [1, 3, 5, 20, 2, 4, 6, 8]
752 underten = lambda x: x<10
753 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000754 self.assertEqual(list(dropwhile(underten, [])), [])
755 self.assertRaises(TypeError, dropwhile)
756 self.assertRaises(TypeError, dropwhile, operator.pow)
757 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000758 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
759 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000760
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000761 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000762 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000763 def irange(n):
Guido van Rossum805365e2007-05-07 22:24:25 +0000764 for i in range(n):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000765 yield i
766
767 a, b = tee([]) # test empty iterator
768 self.assertEqual(list(a), [])
769 self.assertEqual(list(b), [])
770
771 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +0000772 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000773
774 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +0000775 self.assertEqual(list(a), list(range(n)))
776 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000777
778 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000779 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000780 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000781 del a
Guido van Rossum805365e2007-05-07 22:24:25 +0000782 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000783
784 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000785 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000786 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000787 del b
Guido van Rossum805365e2007-05-07 22:24:25 +0000788 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000789
Guido van Rossum805365e2007-05-07 22:24:25 +0000790 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000791 order = [0]*n + [1]*n
792 random.shuffle(order)
793 lists = ([], [])
794 its = tee(irange(n))
795 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +0000796 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000797 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +0000798 self.assertEqual(lists[0], list(range(n)))
799 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000800
Raymond Hettingerad983e72003-11-12 14:32:26 +0000801 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000802 self.assertRaises(TypeError, tee)
803 self.assertRaises(TypeError, tee, 3)
804 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000805 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000806
Raymond Hettingerad983e72003-11-12 14:32:26 +0000807 # tee object should be instantiable
808 a, b = tee('abc')
809 c = type(a)('def')
810 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000811
Raymond Hettingerad983e72003-11-12 14:32:26 +0000812 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +0000813 a, b, c = tee(range(2000), 3)
814 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000815 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +0000816 self.assertEqual(list(b), list(range(2000)))
817 self.assertEqual([next(c), next(c)], list(range(2)))
818 self.assertEqual(list(a), list(range(100,2000)))
819 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000820
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000821 # test values of n
822 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000823 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +0000824 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000825 result = tee('abc', n)
826 self.assertEqual(type(result), tuple)
827 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000828 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000829
Raymond Hettingerad983e72003-11-12 14:32:26 +0000830 # tee pass-through to copyable iterator
831 a, b = tee('abc')
832 c, d = tee(a)
833 self.assert_(a is c)
834
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000835 # test tee_new
836 t1, t2 = tee('abc')
837 tnew = type(t1)
838 self.assertRaises(TypeError, tnew)
839 self.assertRaises(TypeError, tnew, 10)
840 t3 = tnew(t1)
841 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000842
Raymond Hettingera9f60922004-10-17 16:40:14 +0000843 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +0000844 a, b = tee(range(10))
Raymond Hettingera9f60922004-10-17 16:40:14 +0000845 p = proxy(a)
846 self.assertEqual(getattr(p, '__class__'), type(b))
847 del a
848 self.assertRaises(ReferenceError, getattr, p, '__class__')
849
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000850 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000851 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000852
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000853 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +0000854 self.assertRaises(StopIteration, next, f([]))
855 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000856
Georg Brandla18af4e2007-04-21 15:47:16 +0000857 self.assertRaises(StopIteration, next, islice([], None))
858 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000859
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000860 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +0000861 self.assertRaises(StopIteration, next, p)
862 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000863 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +0000864 self.assertRaises(StopIteration, next, p)
865 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000866
Georg Brandla18af4e2007-04-21 15:47:16 +0000867 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000868
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000869 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +0000870 self.assertRaises(StopIteration, next, f(lambda x:x, []))
871 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000872
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000873class TestExamples(unittest.TestCase):
874
875 def test_chain(self):
876 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
877
878 def test_chain_from_iterable(self):
879 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
880
881 def test_combinations(self):
882 self.assertEqual(list(combinations('ABCD', 2)),
883 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
884 self.assertEqual(list(combinations(range(4), 3)),
885 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
886
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000887 def test_combinations_with_replacement(self):
888 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
889 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
890
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000891 def test_compress(self):
892 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
893
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000894 def test_count(self):
895 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
896
897 def test_cycle(self):
898 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
899
900 def test_dropwhile(self):
901 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
902
903 def test_groupby(self):
904 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
905 list('ABCDAB'))
906 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
907 [list('AAAA'), list('BBB'), list('CC'), list('D')])
908
909 def test_filter(self):
910 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
911
912 def test_filterfalse(self):
913 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
914
915 def test_map(self):
916 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
917
918 def test_islice(self):
919 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
920 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
921 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
922 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
923
924 def test_zip(self):
925 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
926
927 def test_zip_longest(self):
928 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
929 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
930
931 def test_permutations(self):
932 self.assertEqual(list(permutations('ABCD', 2)),
933 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
934 self.assertEqual(list(permutations(range(3))),
935 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
936
937 def test_product(self):
938 self.assertEqual(list(product('ABCD', 'xy')),
939 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
940 self.assertEqual(list(product(range(2), repeat=3)),
941 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
942 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
943
944 def test_repeat(self):
945 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
946
947 def test_stapmap(self):
948 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
949 [32, 9, 1000])
950
951 def test_takewhile(self):
952 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
953
954
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000955class TestGC(unittest.TestCase):
956
957 def makecycle(self, iterator, container):
958 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +0000959 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000960 del container, iterator
961
962 def test_chain(self):
963 a = []
964 self.makecycle(chain(a), a)
965
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000966 def test_chain_from_iterable(self):
967 a = []
968 self.makecycle(chain.from_iterable([a]), a)
969
970 def test_combinations(self):
971 a = []
972 self.makecycle(combinations([1,2,a,3], 3), a)
973
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000974 def test_combinations_with_replacement(self):
975 a = []
976 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
977
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000978 def test_compress(self):
979 a = []
980 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
981
Raymond Hettingerd6280f42009-02-16 20:50:56 +0000982 def test_count(self):
983 a = []
984 Int = type('Int', (int,), dict(x=a))
985 self.makecycle(count(Int(0), Int(1)), a)
986
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000987 def test_cycle(self):
988 a = []
989 self.makecycle(cycle([a]*2), a)
990
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000991 def test_dropwhile(self):
992 a = []
993 self.makecycle(dropwhile(bool, [0, a, a]), a)
994
995 def test_groupby(self):
996 a = []
997 self.makecycle(groupby([a]*2, lambda x:x), a)
998
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000999 def test_issue2246(self):
1000 # Issue 2246 -- the _grouper iterator was not included in GC
1001 n = 10
1002 keyfunc = lambda x: x
1003 for i, j in groupby(range(n), key=keyfunc):
1004 keyfunc.__dict__.setdefault('x',[]).append(j)
1005
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001006 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001007 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001008 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001009
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001010 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001011 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001012 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001013
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001014 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001015 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001016 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001017
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001018 def test_zip_longest(self):
1019 a = []
1020 self.makecycle(zip_longest([a]*2, [a]*3), a)
1021 b = [a, None]
1022 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1023
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001024 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001025 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001026 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001027
1028 def test_islice(self):
1029 a = []
1030 self.makecycle(islice([a]*2, None), a)
1031
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001032 def test_permutations(self):
1033 a = []
1034 self.makecycle(permutations([1,2,a,3], 3), a)
1035
1036 def test_product(self):
1037 a = []
1038 self.makecycle(product([1,2,a,3], repeat=3), a)
1039
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001040 def test_repeat(self):
1041 a = []
1042 self.makecycle(repeat(a), a)
1043
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001044 def test_starmap(self):
1045 a = []
1046 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1047
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001048 def test_takewhile(self):
1049 a = []
1050 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1051
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001052def R(seqn):
1053 'Regular generator'
1054 for i in seqn:
1055 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001056
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001057class G:
1058 'Sequence using __getitem__'
1059 def __init__(self, seqn):
1060 self.seqn = seqn
1061 def __getitem__(self, i):
1062 return self.seqn[i]
1063
1064class I:
1065 'Sequence using iterator protocol'
1066 def __init__(self, seqn):
1067 self.seqn = seqn
1068 self.i = 0
1069 def __iter__(self):
1070 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001071 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001072 if self.i >= len(self.seqn): raise StopIteration
1073 v = self.seqn[self.i]
1074 self.i += 1
1075 return v
1076
1077class Ig:
1078 'Sequence using iterator protocol defined with a generator'
1079 def __init__(self, seqn):
1080 self.seqn = seqn
1081 self.i = 0
1082 def __iter__(self):
1083 for val in self.seqn:
1084 yield val
1085
1086class X:
1087 'Missing __getitem__ and __iter__'
1088 def __init__(self, seqn):
1089 self.seqn = seqn
1090 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001091 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001092 if self.i >= len(self.seqn): raise StopIteration
1093 v = self.seqn[self.i]
1094 self.i += 1
1095 return v
1096
1097class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001098 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001099 def __init__(self, seqn):
1100 self.seqn = seqn
1101 self.i = 0
1102 def __iter__(self):
1103 return self
1104
1105class E:
1106 'Test propagation of exceptions'
1107 def __init__(self, seqn):
1108 self.seqn = seqn
1109 self.i = 0
1110 def __iter__(self):
1111 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001112 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001113 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001114
1115class S:
1116 'Test immediate stop'
1117 def __init__(self, seqn):
1118 pass
1119 def __iter__(self):
1120 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001121 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001122 raise StopIteration
1123
1124def L(seqn):
1125 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001126 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001127
1128
1129class TestVariousIteratorArgs(unittest.TestCase):
1130
1131 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001132 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001133 for g in (G, I, Ig, S, L, R):
1134 self.assertEqual(list(chain(g(s))), list(g(s)))
1135 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001136 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001137 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001138 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1139
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001140 def test_compress(self):
1141 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1142 n = len(s)
1143 for g in (G, I, Ig, S, L, R):
1144 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1145 self.assertRaises(TypeError, compress, X(s), repeat(1))
1146 self.assertRaises(TypeError, compress, N(s), repeat(1))
1147 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1148
Christian Heimesc3f30c42008-02-22 16:37:40 +00001149 def test_product(self):
1150 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1151 self.assertRaises(TypeError, product, X(s))
1152 self.assertRaises(TypeError, product, N(s))
1153 self.assertRaises(ZeroDivisionError, product, E(s))
1154
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001155 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001156 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001157 for g in (G, I, Ig, S, L, R):
1158 tgtlen = len(s) * 3
1159 expected = list(g(s))*3
1160 actual = list(islice(cycle(g(s)), tgtlen))
1161 self.assertEqual(actual, expected)
1162 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001163 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001164 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1165
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001166 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001167 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001168 for g in (G, I, Ig, S, L, R):
1169 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1170 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001171 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001172 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1173
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001174 def test_filter(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 Hettinger736c0ab2008-03-13 02:09:15 +00001177 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001178 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001179 self.assertRaises(TypeError, filter, isEven, X(s))
1180 self.assertRaises(TypeError, filter, isEven, N(s))
1181 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001182
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001183 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001184 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001185 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001186 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001187 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001188 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1189 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1190 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001191
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001192 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001193 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001194 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001195 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1196 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1197 self.assertRaises(TypeError, zip, X(s))
1198 self.assertRaises(TypeError, zip, N(s))
1199 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001200
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001201 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001202 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001203 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001204 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1205 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1206 self.assertRaises(TypeError, zip_longest, X(s))
1207 self.assertRaises(TypeError, zip_longest, N(s))
1208 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001209
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001210 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001211 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001212 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001213 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001214 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001215 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001216 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001217 self.assertRaises(TypeError, map, onearg, X(s))
1218 self.assertRaises(TypeError, map, onearg, N(s))
1219 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001220
1221 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001222 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001223 for g in (G, I, Ig, S, L, R):
1224 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1225 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001226 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001227 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1228
1229 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001230 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001231 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001232 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001233 self.assertEqual(list(starmap(operator.pow, g(ss))),
1234 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001235 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001236 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001237 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1238
1239 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001240 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001241 for g in (G, I, Ig, S, L, R):
1242 tgt = []
1243 for elem in g(s):
1244 if not isEven(elem): break
1245 tgt.append(elem)
1246 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1247 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001248 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001249 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1250
1251 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001252 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001253 for g in (G, I, Ig, S, L, R):
1254 tgt = []
1255 for elem in g(s):
1256 if not tgt and isOdd(elem): continue
1257 tgt.append(elem)
1258 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1259 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001260 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001261 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1262
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001263 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001264 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001265 for g in (G, I, Ig, S, L, R):
1266 it1, it2 = tee(g(s))
1267 self.assertEqual(list(it1), list(g(s)))
1268 self.assertEqual(list(it2), list(g(s)))
1269 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001270 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001271 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1272
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001273class LengthTransparency(unittest.TestCase):
1274
1275 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001276 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001277 self.assertEqual(len(repeat(None, 50)), 50)
1278 self.assertRaises(TypeError, len, repeat(None))
1279
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001280class RegressionTests(unittest.TestCase):
1281
1282 def test_sf_793826(self):
1283 # Fix Armin Rigo's successful efforts to wreak havoc
1284
1285 def mutatingtuple(tuple1, f, tuple2):
1286 # this builds a tuple t which is a copy of tuple1,
1287 # then calls f(t), then mutates t to be equal to tuple2
1288 # (needs len(tuple1) == len(tuple2)).
1289 def g(value, first=[1]):
1290 if first:
1291 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00001292 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001293 return value
1294 items = list(tuple2)
1295 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001296 gen = map(g, items)
1297 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00001298 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001299
1300 def f(t):
1301 global T
1302 T = t
1303 first[:] = list(T)
1304
1305 first = []
1306 mutatingtuple((1,2,3), f, (4,5,6))
1307 second = list(T)
1308 self.assertEqual(first, second)
1309
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001310
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001311 def test_sf_950057(self):
1312 # Make sure that chain() and cycle() catch exceptions immediately
1313 # rather than when shifting between input sources
1314
1315 def gen1():
1316 hist.append(0)
1317 yield 1
1318 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001319 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001320 hist.append(2)
1321
1322 def gen2(x):
1323 hist.append(3)
1324 yield 2
1325 hist.append(4)
1326 if x:
1327 raise StopIteration
1328
1329 hist = []
1330 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1331 self.assertEqual(hist, [0,1])
1332
1333 hist = []
1334 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1335 self.assertEqual(hist, [0,1])
1336
1337 hist = []
1338 self.assertRaises(AssertionError, list, cycle(gen1()))
1339 self.assertEqual(hist, [0,1])
1340
Thomas Woutersb2137042007-02-01 18:02:27 +00001341class SubclassWithKwargsTest(unittest.TestCase):
1342 def test_keywords_in_subclass(self):
1343 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001344 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001345 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00001346 class Subclass(cls):
1347 def __init__(self, newarg=None, *args):
1348 cls.__init__(self, *args)
1349 try:
1350 Subclass(newarg=1)
1351 except TypeError as err:
1352 # we expect type errors because of wrong argument count
1353 self.failIf("does not take keyword arguments" in err.args[0])
1354
1355
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001356libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001357
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001358
1359>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001360>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00001361... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001362...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001363Check 1200 is for $120.15
1364Check 1201 is for $764.05
1365Check 1202 is for $823.14
1366
1367>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001368>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00001369... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00001370...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000013711
13728
137327
1374
1375>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001376>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00001377... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00001378...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001379Alex
1380Laura
1381Martin
1382Walter
1383Samuele
1384
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001385>>> from operator import itemgetter
1386>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00001387>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001388>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001389... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001390...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000013911 ['a', 'c', 'e']
13922 ['b', 'd', 'f']
13933 ['g']
1394
Raymond Hettinger734fb572004-01-20 20:04:40 +00001395# Find runs of consecutive numbers using groupby. The key to the solution
1396# is differencing with a range so that consecutive numbers all appear in
1397# same group.
1398>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00001399>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001400... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001401...
Raymond Hettinger734fb572004-01-20 20:04:40 +00001402[1]
1403[4, 5, 6]
1404[10]
1405[15, 16, 17, 18]
1406[22]
1407[25, 26, 27, 28]
1408
Georg Brandl3dbca812008-07-23 16:10:53 +00001409>>> def take(n, iterable):
1410... "Return first n items of the iterable as a list"
1411... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001412
Georg Brandl3dbca812008-07-23 16:10:53 +00001413>>> def enumerate(iterable, start=0):
1414... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001415
Georg Brandl3dbca812008-07-23 16:10:53 +00001416>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001417... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00001418... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001419
1420>>> def nth(iterable, n):
Raymond Hettingerd04fa312009-02-04 19:45:13 +00001421... "Returns the nth item or None"
1422... return next(islice(iterable, n, None), None)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001423
Georg Brandl3dbca812008-07-23 16:10:53 +00001424>>> def quantify(iterable, pred=bool):
1425... "Count how many times the predicate is true"
1426... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001427
Georg Brandl3dbca812008-07-23 16:10:53 +00001428>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001429... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00001430... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001431
Georg Brandl3dbca812008-07-23 16:10:53 +00001432>>> def ncycles(iterable, n):
1433... "Returns the seqeuence elements n times"
1434... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001435
1436>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001437... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001438
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001439>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001440... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001441
1442>>> def repeatfunc(func, times=None, *args):
1443... "Repeat calls to func with specified arguments."
1444... " Example: repeatfunc(random.random)"
1445... if times is None:
1446... return starmap(func, repeat(args))
1447... else:
1448... return starmap(func, repeat(args, times))
1449
Raymond Hettingerd591f662003-10-26 15:34:50 +00001450>>> def pairwise(iterable):
1451... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1452... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001453... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00001454... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001455... except StopIteration:
1456... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001457... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001458
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001459>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001460... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001461... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00001462... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001463
1464>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001465... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001466... # Recipe credited to George Sakkis
1467... pending = len(iterables)
1468... nexts = cycle(iter(it).__next__ for it in iterables)
1469... while pending:
1470... try:
1471... for next in nexts:
1472... yield next()
1473... except StopIteration:
1474... pending -= 1
1475... nexts = cycle(islice(nexts, pending))
1476
1477>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00001478... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1479... s = list(iterable)
1480... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001481
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001482>>> def unique_everseen(iterable, key=None):
1483... "List unique elements, preserving order. Remember all elements ever seen."
1484... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1485... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1486... seen = set()
1487... seen_add = seen.add
1488... if key is None:
1489... for element in iterable:
1490... if element not in seen:
1491... seen_add(element)
1492... yield element
1493... else:
1494... for element in iterable:
1495... k = key(element)
1496... if k not in seen:
1497... seen_add(k)
1498... yield element
1499
1500>>> def unique_justseen(iterable, key=None):
1501... "List unique elements, preserving order. Remember only the element just seen."
1502... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1503... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1504... return map(next, map(itemgetter(1), groupby(iterable, key)))
1505
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001506This is not part of the examples but it tests to make sure the definitions
1507perform as purported.
1508
Raymond Hettingera098b332003-09-08 23:58:40 +00001509>>> take(10, count())
1510[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1511
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001512>>> list(enumerate('abc'))
1513[(0, 'a'), (1, 'b'), (2, 'c')]
1514
1515>>> list(islice(tabulate(lambda x: 2*x), 4))
1516[0, 2, 4, 6]
1517
1518>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00001519'd'
1520
1521>>> nth('abcde', 9) is None
1522True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001523
Guido van Rossum805365e2007-05-07 22:24:25 +00001524>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000152550
1526
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001527>>> a = [[1, 2, 3], [4, 5, 6]]
1528>>> flatten(a)
1529[1, 2, 3, 4, 5, 6]
1530
1531>>> list(repeatfunc(pow, 5, 2, 3))
1532[8, 8, 8, 8, 8]
1533
1534>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001535>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001536[0, 0, 0, 0, 0]
1537
Raymond Hettingerd591f662003-10-26 15:34:50 +00001538>>> list(pairwise('abcd'))
1539[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001540
Raymond Hettingerd591f662003-10-26 15:34:50 +00001541>>> list(pairwise([]))
1542[]
1543
1544>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001545[]
1546
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001547>>> list(islice(padnone('abc'), 0, 6))
1548['a', 'b', 'c', None, None, None]
1549
1550>>> list(ncycles('abc', 3))
1551['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1552
1553>>> dotproduct([1,2,3], [4,5,6])
155432
1555
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001556>>> list(grouper(3, 'abcdefg', 'x'))
1557[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1558
1559>>> list(roundrobin('abc', 'd', 'ef'))
1560['a', 'd', 'e', 'b', 'f', 'c']
1561
Raymond Hettingerace67332009-01-26 02:23:50 +00001562>>> list(powerset([1,2,3]))
1563[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001564
Raymond Hettinger191e8502009-01-27 13:29:43 +00001565>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1566True
1567
1568>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1569True
1570
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001571>>> list(unique_everseen('AAAABBBCCDAABBB'))
1572['A', 'B', 'C', 'D']
1573
1574>>> list(unique_everseen('ABBCcAD', str.lower))
1575['A', 'B', 'C', 'D']
1576
1577>>> list(unique_justseen('AAAABBBCCDAABBB'))
1578['A', 'B', 'C', 'D', 'A', 'B']
1579
1580>>> list(unique_justseen('ABBCcAD', str.lower))
1581['A', 'B', 'C', 'A', 'D']
1582
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001583"""
1584
1585__test__ = {'libreftest' : libreftest}
1586
1587def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001588 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00001589 RegressionTests, LengthTransparency,
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001590 SubclassWithKwargsTest, TestExamples)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001591 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001592
1593 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001594 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001595 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001596 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00001597 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001598 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001599 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001600 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001601 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001602
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001603 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001604 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001605
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001606if __name__ == "__main__":
1607 test_main(verbose=True)