blob: 812a1efdb48089aafc13349270b7eef15395c4ff [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 Hettinger6a5b0272003-10-24 08:45:23 +0000982 def test_cycle(self):
983 a = []
984 self.makecycle(cycle([a]*2), a)
985
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000986 def test_dropwhile(self):
987 a = []
988 self.makecycle(dropwhile(bool, [0, a, a]), a)
989
990 def test_groupby(self):
991 a = []
992 self.makecycle(groupby([a]*2, lambda x:x), a)
993
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000994 def test_issue2246(self):
995 # Issue 2246 -- the _grouper iterator was not included in GC
996 n = 10
997 keyfunc = lambda x: x
998 for i, j in groupby(range(n), key=keyfunc):
999 keyfunc.__dict__.setdefault('x',[]).append(j)
1000
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001001 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001002 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001003 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001004
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001005 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001006 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001007 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001008
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001009 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001010 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001011 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001012
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001013 def test_zip_longest(self):
1014 a = []
1015 self.makecycle(zip_longest([a]*2, [a]*3), a)
1016 b = [a, None]
1017 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1018
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001019 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001020 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001021 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001022
1023 def test_islice(self):
1024 a = []
1025 self.makecycle(islice([a]*2, None), a)
1026
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001027 def test_permutations(self):
1028 a = []
1029 self.makecycle(permutations([1,2,a,3], 3), a)
1030
1031 def test_product(self):
1032 a = []
1033 self.makecycle(product([1,2,a,3], repeat=3), a)
1034
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001035 def test_repeat(self):
1036 a = []
1037 self.makecycle(repeat(a), a)
1038
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001039 def test_starmap(self):
1040 a = []
1041 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1042
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001043 def test_takewhile(self):
1044 a = []
1045 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1046
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001047def R(seqn):
1048 'Regular generator'
1049 for i in seqn:
1050 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001051
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001052class G:
1053 'Sequence using __getitem__'
1054 def __init__(self, seqn):
1055 self.seqn = seqn
1056 def __getitem__(self, i):
1057 return self.seqn[i]
1058
1059class I:
1060 'Sequence using iterator protocol'
1061 def __init__(self, seqn):
1062 self.seqn = seqn
1063 self.i = 0
1064 def __iter__(self):
1065 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001066 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001067 if self.i >= len(self.seqn): raise StopIteration
1068 v = self.seqn[self.i]
1069 self.i += 1
1070 return v
1071
1072class Ig:
1073 'Sequence using iterator protocol defined with a generator'
1074 def __init__(self, seqn):
1075 self.seqn = seqn
1076 self.i = 0
1077 def __iter__(self):
1078 for val in self.seqn:
1079 yield val
1080
1081class X:
1082 'Missing __getitem__ and __iter__'
1083 def __init__(self, seqn):
1084 self.seqn = seqn
1085 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001086 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001087 if self.i >= len(self.seqn): raise StopIteration
1088 v = self.seqn[self.i]
1089 self.i += 1
1090 return v
1091
1092class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001093 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001094 def __init__(self, seqn):
1095 self.seqn = seqn
1096 self.i = 0
1097 def __iter__(self):
1098 return self
1099
1100class E:
1101 'Test propagation of exceptions'
1102 def __init__(self, seqn):
1103 self.seqn = seqn
1104 self.i = 0
1105 def __iter__(self):
1106 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001107 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001108 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001109
1110class S:
1111 'Test immediate stop'
1112 def __init__(self, seqn):
1113 pass
1114 def __iter__(self):
1115 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001116 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001117 raise StopIteration
1118
1119def L(seqn):
1120 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001121 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001122
1123
1124class TestVariousIteratorArgs(unittest.TestCase):
1125
1126 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001127 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001128 for g in (G, I, Ig, S, L, R):
1129 self.assertEqual(list(chain(g(s))), list(g(s)))
1130 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001131 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001132 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001133 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1134
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001135 def test_compress(self):
1136 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1137 n = len(s)
1138 for g in (G, I, Ig, S, L, R):
1139 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1140 self.assertRaises(TypeError, compress, X(s), repeat(1))
1141 self.assertRaises(TypeError, compress, N(s), repeat(1))
1142 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1143
Christian Heimesc3f30c42008-02-22 16:37:40 +00001144 def test_product(self):
1145 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1146 self.assertRaises(TypeError, product, X(s))
1147 self.assertRaises(TypeError, product, N(s))
1148 self.assertRaises(ZeroDivisionError, product, E(s))
1149
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001150 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001151 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001152 for g in (G, I, Ig, S, L, R):
1153 tgtlen = len(s) * 3
1154 expected = list(g(s))*3
1155 actual = list(islice(cycle(g(s)), tgtlen))
1156 self.assertEqual(actual, expected)
1157 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001158 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001159 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1160
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001161 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001162 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001163 for g in (G, I, Ig, S, L, R):
1164 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1165 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001166 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001167 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1168
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001169 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001170 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001171 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001172 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001173 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001174 self.assertRaises(TypeError, filter, isEven, X(s))
1175 self.assertRaises(TypeError, filter, isEven, N(s))
1176 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001177
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001178 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001179 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001180 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001181 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001182 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001183 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1184 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1185 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001186
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001187 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001188 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001189 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001190 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1191 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1192 self.assertRaises(TypeError, zip, X(s))
1193 self.assertRaises(TypeError, zip, N(s))
1194 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001195
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001196 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001197 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001198 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001199 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1200 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1201 self.assertRaises(TypeError, zip_longest, X(s))
1202 self.assertRaises(TypeError, zip_longest, N(s))
1203 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001204
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001205 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001206 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001207 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001208 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001209 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001210 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001211 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001212 self.assertRaises(TypeError, map, onearg, X(s))
1213 self.assertRaises(TypeError, map, onearg, N(s))
1214 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001215
1216 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001217 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001218 for g in (G, I, Ig, S, L, R):
1219 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1220 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001221 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001222 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1223
1224 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001225 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001226 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001227 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001228 self.assertEqual(list(starmap(operator.pow, g(ss))),
1229 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001230 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001231 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001232 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1233
1234 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001235 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001236 for g in (G, I, Ig, S, L, R):
1237 tgt = []
1238 for elem in g(s):
1239 if not isEven(elem): break
1240 tgt.append(elem)
1241 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1242 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001243 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001244 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1245
1246 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001247 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001248 for g in (G, I, Ig, S, L, R):
1249 tgt = []
1250 for elem in g(s):
1251 if not tgt and isOdd(elem): continue
1252 tgt.append(elem)
1253 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1254 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001255 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001256 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1257
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001258 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001259 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001260 for g in (G, I, Ig, S, L, R):
1261 it1, it2 = tee(g(s))
1262 self.assertEqual(list(it1), list(g(s)))
1263 self.assertEqual(list(it2), list(g(s)))
1264 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001265 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001266 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1267
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001268class LengthTransparency(unittest.TestCase):
1269
1270 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001271 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001272 self.assertEqual(len(repeat(None, 50)), 50)
1273 self.assertRaises(TypeError, len, repeat(None))
1274
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001275class RegressionTests(unittest.TestCase):
1276
1277 def test_sf_793826(self):
1278 # Fix Armin Rigo's successful efforts to wreak havoc
1279
1280 def mutatingtuple(tuple1, f, tuple2):
1281 # this builds a tuple t which is a copy of tuple1,
1282 # then calls f(t), then mutates t to be equal to tuple2
1283 # (needs len(tuple1) == len(tuple2)).
1284 def g(value, first=[1]):
1285 if first:
1286 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00001287 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001288 return value
1289 items = list(tuple2)
1290 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001291 gen = map(g, items)
1292 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00001293 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001294
1295 def f(t):
1296 global T
1297 T = t
1298 first[:] = list(T)
1299
1300 first = []
1301 mutatingtuple((1,2,3), f, (4,5,6))
1302 second = list(T)
1303 self.assertEqual(first, second)
1304
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001305
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001306 def test_sf_950057(self):
1307 # Make sure that chain() and cycle() catch exceptions immediately
1308 # rather than when shifting between input sources
1309
1310 def gen1():
1311 hist.append(0)
1312 yield 1
1313 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001314 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001315 hist.append(2)
1316
1317 def gen2(x):
1318 hist.append(3)
1319 yield 2
1320 hist.append(4)
1321 if x:
1322 raise StopIteration
1323
1324 hist = []
1325 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1326 self.assertEqual(hist, [0,1])
1327
1328 hist = []
1329 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1330 self.assertEqual(hist, [0,1])
1331
1332 hist = []
1333 self.assertRaises(AssertionError, list, cycle(gen1()))
1334 self.assertEqual(hist, [0,1])
1335
Thomas Woutersb2137042007-02-01 18:02:27 +00001336class SubclassWithKwargsTest(unittest.TestCase):
1337 def test_keywords_in_subclass(self):
1338 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001339 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001340 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00001341 class Subclass(cls):
1342 def __init__(self, newarg=None, *args):
1343 cls.__init__(self, *args)
1344 try:
1345 Subclass(newarg=1)
1346 except TypeError as err:
1347 # we expect type errors because of wrong argument count
1348 self.failIf("does not take keyword arguments" in err.args[0])
1349
1350
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001351libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001352
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001353
1354>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001355>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00001356... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001357...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001358Check 1200 is for $120.15
1359Check 1201 is for $764.05
1360Check 1202 is for $823.14
1361
1362>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001363>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00001364... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00001365...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000013661
13678
136827
1369
1370>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001371>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00001372... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00001373...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001374Alex
1375Laura
1376Martin
1377Walter
1378Samuele
1379
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001380>>> from operator import itemgetter
1381>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00001382>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001383>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001384... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001385...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000013861 ['a', 'c', 'e']
13872 ['b', 'd', 'f']
13883 ['g']
1389
Raymond Hettinger734fb572004-01-20 20:04:40 +00001390# Find runs of consecutive numbers using groupby. The key to the solution
1391# is differencing with a range so that consecutive numbers all appear in
1392# same group.
1393>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00001394>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001395... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001396...
Raymond Hettinger734fb572004-01-20 20:04:40 +00001397[1]
1398[4, 5, 6]
1399[10]
1400[15, 16, 17, 18]
1401[22]
1402[25, 26, 27, 28]
1403
Georg Brandl3dbca812008-07-23 16:10:53 +00001404>>> def take(n, iterable):
1405... "Return first n items of the iterable as a list"
1406... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001407
Georg Brandl3dbca812008-07-23 16:10:53 +00001408>>> def enumerate(iterable, start=0):
1409... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001410
Georg Brandl3dbca812008-07-23 16:10:53 +00001411>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001412... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00001413... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001414
1415>>> def nth(iterable, n):
Raymond Hettingerd04fa312009-02-04 19:45:13 +00001416... "Returns the nth item or None"
1417... return next(islice(iterable, n, None), None)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001418
Georg Brandl3dbca812008-07-23 16:10:53 +00001419>>> def quantify(iterable, pred=bool):
1420... "Count how many times the predicate is true"
1421... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001422
Georg Brandl3dbca812008-07-23 16:10:53 +00001423>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001424... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00001425... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001426
Georg Brandl3dbca812008-07-23 16:10:53 +00001427>>> def ncycles(iterable, n):
1428... "Returns the seqeuence elements n times"
1429... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001430
1431>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001432... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001433
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001434>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001435... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001436
1437>>> def repeatfunc(func, times=None, *args):
1438... "Repeat calls to func with specified arguments."
1439... " Example: repeatfunc(random.random)"
1440... if times is None:
1441... return starmap(func, repeat(args))
1442... else:
1443... return starmap(func, repeat(args, times))
1444
Raymond Hettingerd591f662003-10-26 15:34:50 +00001445>>> def pairwise(iterable):
1446... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1447... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001448... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00001449... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001450... except StopIteration:
1451... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001452... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001453
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001454>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001455... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001456... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00001457... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001458
1459>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001460... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001461... # Recipe credited to George Sakkis
1462... pending = len(iterables)
1463... nexts = cycle(iter(it).__next__ for it in iterables)
1464... while pending:
1465... try:
1466... for next in nexts:
1467... yield next()
1468... except StopIteration:
1469... pending -= 1
1470... nexts = cycle(islice(nexts, pending))
1471
1472>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00001473... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1474... s = list(iterable)
1475... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001476
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001477>>> def unique_everseen(iterable, key=None):
1478... "List unique elements, preserving order. Remember all elements ever seen."
1479... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1480... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1481... seen = set()
1482... seen_add = seen.add
1483... if key is None:
1484... for element in iterable:
1485... if element not in seen:
1486... seen_add(element)
1487... yield element
1488... else:
1489... for element in iterable:
1490... k = key(element)
1491... if k not in seen:
1492... seen_add(k)
1493... yield element
1494
1495>>> def unique_justseen(iterable, key=None):
1496... "List unique elements, preserving order. Remember only the element just seen."
1497... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1498... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1499... return map(next, map(itemgetter(1), groupby(iterable, key)))
1500
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001501This is not part of the examples but it tests to make sure the definitions
1502perform as purported.
1503
Raymond Hettingera098b332003-09-08 23:58:40 +00001504>>> take(10, count())
1505[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1506
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001507>>> list(enumerate('abc'))
1508[(0, 'a'), (1, 'b'), (2, 'c')]
1509
1510>>> list(islice(tabulate(lambda x: 2*x), 4))
1511[0, 2, 4, 6]
1512
1513>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00001514'd'
1515
1516>>> nth('abcde', 9) is None
1517True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001518
Guido van Rossum805365e2007-05-07 22:24:25 +00001519>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000152050
1521
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001522>>> a = [[1, 2, 3], [4, 5, 6]]
1523>>> flatten(a)
1524[1, 2, 3, 4, 5, 6]
1525
1526>>> list(repeatfunc(pow, 5, 2, 3))
1527[8, 8, 8, 8, 8]
1528
1529>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001530>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001531[0, 0, 0, 0, 0]
1532
Raymond Hettingerd591f662003-10-26 15:34:50 +00001533>>> list(pairwise('abcd'))
1534[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001535
Raymond Hettingerd591f662003-10-26 15:34:50 +00001536>>> list(pairwise([]))
1537[]
1538
1539>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001540[]
1541
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001542>>> list(islice(padnone('abc'), 0, 6))
1543['a', 'b', 'c', None, None, None]
1544
1545>>> list(ncycles('abc', 3))
1546['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1547
1548>>> dotproduct([1,2,3], [4,5,6])
154932
1550
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001551>>> list(grouper(3, 'abcdefg', 'x'))
1552[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1553
1554>>> list(roundrobin('abc', 'd', 'ef'))
1555['a', 'd', 'e', 'b', 'f', 'c']
1556
Raymond Hettingerace67332009-01-26 02:23:50 +00001557>>> list(powerset([1,2,3]))
1558[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001559
Raymond Hettinger191e8502009-01-27 13:29:43 +00001560>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1561True
1562
1563>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1564True
1565
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001566>>> list(unique_everseen('AAAABBBCCDAABBB'))
1567['A', 'B', 'C', 'D']
1568
1569>>> list(unique_everseen('ABBCcAD', str.lower))
1570['A', 'B', 'C', 'D']
1571
1572>>> list(unique_justseen('AAAABBBCCDAABBB'))
1573['A', 'B', 'C', 'D', 'A', 'B']
1574
1575>>> list(unique_justseen('ABBCcAD', str.lower))
1576['A', 'B', 'C', 'A', 'D']
1577
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001578"""
1579
1580__test__ = {'libreftest' : libreftest}
1581
1582def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001583 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00001584 RegressionTests, LengthTransparency,
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001585 SubclassWithKwargsTest, TestExamples)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001586 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001587
1588 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001589 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001590 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001591 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00001592 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001593 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001594 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001595 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001596 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001597
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001598 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001599 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001600
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001601if __name__ == "__main__":
1602 test_main(verbose=True)