blob: 94159b5d0e6c28b55f64b8350ddc90d4e94aaaeb [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):
Raymond Hettinger15a49502009-02-19 02:17:09 +0000313 self.assertEqual(list(compress(data='ABCDEF', selectors=[1,0,1,0,1,1])), list('ACEF'))
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000314 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
315 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
316 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
317 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
318 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
319 n = 10000
320 data = chain.from_iterable(repeat(range(6), n))
321 selectors = chain.from_iterable(repeat((0, 1)))
322 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
323 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
324 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
325 self.assertRaises(TypeError, compress, range(6)) # too few args
326 self.assertRaises(TypeError, compress, range(6), None) # too many args
327
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000328 def test_count(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000329 self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
330 self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
331 self.assertEqual(take(2, lzip('abc',count(3))), [('a', 3), ('b', 4)])
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000332 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
333 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000334 self.assertRaises(TypeError, count, 2, 3, 4)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000335 self.assertRaises(TypeError, count, 'a')
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000336 self.assertEqual(list(islice(count(maxsize-5), 10)),
337 list(range(maxsize-5, maxsize+5)))
338 self.assertEqual(list(islice(count(-maxsize-5), 10)),
339 list(range(-maxsize-5, -maxsize+5)))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000340 c = count(3)
341 self.assertEqual(repr(c), 'count(3)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000342 next(c)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000343 self.assertEqual(repr(c), 'count(4)')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000344 c = count(-9)
345 self.assertEqual(repr(c), 'count(-9)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000346 next(c)
Raymond Hettinger30729212009-02-12 06:28:27 +0000347 self.assertEqual(repr(count(10.25)), 'count(10.25)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000348 self.assertEqual(next(c), -8)
Christian Heimesa37d4c62007-12-04 23:02:19 +0000349 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 +0000350 # Test repr (ignoring the L in longs)
351 r1 = repr(count(i)).replace('L', '')
352 r2 = 'count(%r)'.__mod__(i).replace('L', '')
353 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000354
Raymond Hettinger30729212009-02-12 06:28:27 +0000355 def test_count_with_stride(self):
356 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
357 self.assertEqual(lzip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
358 self.assertEqual(lzip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000359 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000360 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
361 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
362 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettinger8a882a72009-02-12 12:08:18 +0000363 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
364 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerde09acf2009-02-12 12:53:53 +0000365 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
366 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000367 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
368 c = count(3, 5)
369 self.assertEqual(repr(c), 'count(3, 5)')
370 next(c)
371 self.assertEqual(repr(c), 'count(8, 5)')
372 c = count(-9, 0)
373 self.assertEqual(repr(c), 'count(-9, 0)')
374 next(c)
375 self.assertEqual(repr(c), 'count(-9, 0)')
376 c = count(-9, -3)
377 self.assertEqual(repr(c), 'count(-9, -3)')
378 next(c)
379 self.assertEqual(repr(c), 'count(-12, -3)')
380 self.assertEqual(repr(c), 'count(-12, -3)')
381 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
382 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
383 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
384 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
385 for j in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 1, 10, sys.maxsize-5, sys.maxsize+5):
386 # Test repr (ignoring the L in longs)
387 r1 = repr(count(i, j)).replace('L', '')
388 if j == 1:
389 r2 = ('count(%r)' % i).replace('L', '')
390 else:
391 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
392 self.assertEqual(r1, r2)
393
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000394 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000395 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000396 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000397 self.assertRaises(TypeError, cycle)
398 self.assertRaises(TypeError, cycle, 5)
399 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000400
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000401 def test_groupby(self):
402 # Check whether it accepts arguments correctly
403 self.assertEqual([], list(groupby([])))
404 self.assertEqual([], list(groupby([], key=id)))
405 self.assertRaises(TypeError, list, groupby('abc', []))
406 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000407 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000408
409 # Check normal input
410 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
411 (2,15,22), (3,16,23), (3,17,23)]
412 dup = []
413 for k, g in groupby(s, lambda r:r[0]):
414 for elem in g:
415 self.assertEqual(k, elem[0])
416 dup.append(elem)
417 self.assertEqual(s, dup)
418
419 # Check nested case
420 dup = []
421 for k, g in groupby(s, lambda r:r[0]):
422 for ik, ig in groupby(g, lambda r:r[2]):
423 for elem in ig:
424 self.assertEqual(k, elem[0])
425 self.assertEqual(ik, elem[2])
426 dup.append(elem)
427 self.assertEqual(s, dup)
428
429 # Check case where inner iterator is not used
430 keys = [k for k, g in groupby(s, lambda r:r[0])]
431 expectedkeys = set([r[0] for r in s])
432 self.assertEqual(set(keys), expectedkeys)
433 self.assertEqual(len(keys), len(expectedkeys))
434
435 # Exercise pipes and filters style
436 s = 'abracadabra'
437 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000438 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000439 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
440 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000441 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000442 self.assertEqual(r, ['a', 'b', 'r'])
443 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000444 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000445 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
446 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000447 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000448 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
449
Georg Brandla18af4e2007-04-21 15:47:16 +0000450 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000451 class ExpectedError(Exception):
452 pass
453 def delayed_raise(n=0):
454 for i in range(n):
455 yield 'yo'
456 raise ExpectedError
457 def gulp(iterable, keyp=None, func=list):
458 return [func(g) for k, g in groupby(iterable, keyp)]
459
Georg Brandla18af4e2007-04-21 15:47:16 +0000460 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000461 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000462 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000463 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
464
465 # __cmp__ failure
466 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000467 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000468 raise ExpectedError
469 s = [DummyCmp(), DummyCmp(), None]
470
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000471 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000472 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000473 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000474 self.assertRaises(ExpectedError, gulp, s)
475
476 # keyfunc failure
477 def keyfunc(obj):
478 if keyfunc.skip > 0:
479 keyfunc.skip -= 1
480 return obj
481 else:
482 raise ExpectedError
483
484 # keyfunc failure on outer object
485 keyfunc.skip = 0
486 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
487 keyfunc.skip = 1
488 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
489
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000490 def test_filter(self):
491 self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
492 self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])
493 self.assertEqual(list(filter(bool, [0,1,0,2,0])), [1,2])
494 self.assertEqual(take(4, filter(isEven, count())), [0,2,4,6])
495 self.assertRaises(TypeError, filter)
496 self.assertRaises(TypeError, filter, lambda x:x)
497 self.assertRaises(TypeError, filter, lambda x:x, range(6), 7)
498 self.assertRaises(TypeError, filter, isEven, 3)
499 self.assertRaises(TypeError, next, filter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000500
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000501 def test_filterfalse(self):
502 self.assertEqual(list(filterfalse(isEven, range(6))), [1,3,5])
503 self.assertEqual(list(filterfalse(None, [0,1,0,2,0])), [0,0,0])
504 self.assertEqual(list(filterfalse(bool, [0,1,0,2,0])), [0,0,0])
505 self.assertEqual(take(4, filterfalse(isEven, count())), [1,3,5,7])
506 self.assertRaises(TypeError, filterfalse)
507 self.assertRaises(TypeError, filterfalse, lambda x:x)
508 self.assertRaises(TypeError, filterfalse, lambda x:x, range(6), 7)
509 self.assertRaises(TypeError, filterfalse, isEven, 3)
510 self.assertRaises(TypeError, next, filterfalse(range(6), range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000511
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000512 def test_zip(self):
513 # XXX This is rather silly now that builtin zip() calls zip()...
514 ans = [(x,y) for x, y in zip('abc',count())]
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000515 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000516 self.assertEqual(list(zip('abc', range(6))), lzip('abc', range(6)))
517 self.assertEqual(list(zip('abcdef', range(3))), lzip('abcdef', range(3)))
518 self.assertEqual(take(3,zip('abcdef', count())), lzip('abcdef', range(3)))
519 self.assertEqual(list(zip('abcdef')), lzip('abcdef'))
520 self.assertEqual(list(zip()), lzip())
521 self.assertRaises(TypeError, zip, 3)
522 self.assertRaises(TypeError, zip, range(3), 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000523 # Check tuple re-use (implementation detail)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000524 self.assertEqual([tuple(list(pair)) for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000525 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000526 self.assertEqual([pair for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000527 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000528 ids = list(map(id, zip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000529 self.assertEqual(min(ids), max(ids))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000530 ids = list(map(id, list(zip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000531 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000532
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000533 def test_ziplongest(self):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000534 for args in [
535 ['abc', range(6)],
536 [range(6), 'abc'],
537 [range(1000), range(2000,2100), range(3000,3050)],
538 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
539 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
540 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000541 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
542 for i in range(max(map(len, args)))]
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000543 self.assertEqual(list(zip_longest(*args)), target)
544 self.assertEqual(list(zip_longest(*args, **{})), target)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000545 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 +0000546 self.assertEqual(list(zip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000547
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000548 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 +0000549
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000550 self.assertEqual(list(zip_longest()), list(zip()))
551 self.assertEqual(list(zip_longest([])), list(zip([])))
552 self.assertEqual(list(zip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000553
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000554 self.assertEqual(list(zip_longest('abc', 'defg', **{})),
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000555 list(zip(list('abc')+[None], 'defg'))) # empty keyword dict
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000556 self.assertRaises(TypeError, zip_longest, 3)
557 self.assertRaises(TypeError, zip_longest, range(3), 3)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000558
559 for stmt in [
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000560 "zip_longest('abc', fv=1)",
561 "zip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000562 ]:
563 try:
564 eval(stmt, globals(), locals())
565 except TypeError:
566 pass
567 else:
568 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000569
Thomas Wouterscf297e42007-02-23 15:07:44 +0000570 # Check tuple re-use (implementation detail)
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000571 self.assertEqual([tuple(list(pair)) for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000572 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000573 self.assertEqual([pair for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000574 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000575 ids = list(map(id, zip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000576 self.assertEqual(min(ids), max(ids))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000577 ids = list(map(id, list(zip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000578 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
579
Christian Heimesc3f30c42008-02-22 16:37:40 +0000580 def test_product(self):
581 for args, result in [
Christian Heimes78644762008-03-04 23:39:23 +0000582 ([], [()]), # zero iterables
Christian Heimesc3f30c42008-02-22 16:37:40 +0000583 (['ab'], [('a',), ('b',)]), # one iterable
584 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
585 ([range(0), range(2), range(3)], []), # first iterable with zero length
586 ([range(2), range(0), range(3)], []), # middle iterable with zero length
587 ([range(2), range(3), range(0)], []), # last iterable with zero length
588 ]:
589 self.assertEqual(list(product(*args)), result)
Christian Heimesf16baeb2008-02-29 14:57:44 +0000590 for r in range(4):
591 self.assertEqual(list(product(*(args*r))),
592 list(product(*args, **dict(repeat=r))))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000593 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
594 self.assertRaises(TypeError, product, range(6), None)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000595
596 def product1(*args, **kwds):
597 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
598 n = len(pools)
599 if n == 0:
600 yield ()
601 return
602 if any(len(pool) == 0 for pool in pools):
603 return
604 indices = [0] * n
605 yield tuple(pool[i] for pool, i in zip(pools, indices))
606 while 1:
607 for i in reversed(range(n)): # right to left
608 if indices[i] == len(pools[i]) - 1:
609 continue
610 indices[i] += 1
611 for j in range(i+1, n):
612 indices[j] = 0
613 yield tuple(pool[i] for pool, i in zip(pools, indices))
614 break
615 else:
616 return
617
618 def product2(*args, **kwds):
619 'Pure python version used in docs'
620 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
621 result = [[]]
622 for pool in pools:
623 result = [x+[y] for x in result for y in pool]
624 for prod in result:
625 yield tuple(prod)
626
Christian Heimesc3f30c42008-02-22 16:37:40 +0000627 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
628 set('abcdefg'), range(11), tuple(range(13))]
629 for i in range(100):
630 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Christian Heimes78644762008-03-04 23:39:23 +0000631 expected_len = prod(map(len, args))
632 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000633 self.assertEqual(list(product(*args)), list(product1(*args)))
634 self.assertEqual(list(product(*args)), list(product2(*args)))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000635 args = map(iter, args)
Christian Heimes78644762008-03-04 23:39:23 +0000636 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000637
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000638 # Test implementation detail: tuple re-use
639 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
640 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000641
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000642 def test_repeat(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000643 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000644 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000645 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000646 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000647 self.assertEqual(list(repeat('a', 0)), [])
648 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000649 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000650 self.assertRaises(TypeError, repeat, None, 3, 4)
651 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000652 r = repeat(1+0j)
653 self.assertEqual(repr(r), 'repeat((1+0j))')
654 r = repeat(1+0j, 5)
655 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
656 list(r)
657 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000658
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000659 def test_map(self):
660 self.assertEqual(list(map(operator.pow, range(3), range(1,7))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000661 [0**1, 1**2, 2**3])
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000662 def tupleize(*args):
663 return args
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000664 self.assertEqual(list(map(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000665 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000666 self.assertEqual(list(map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000667 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000668 self.assertEqual(take(2,map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000669 [('a',0),('b',1)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000670 self.assertEqual(list(map(operator.pow, [])), [])
671 self.assertRaises(TypeError, map)
672 self.assertRaises(TypeError, list, map(None, range(3), range(3)))
673 self.assertRaises(TypeError, map, operator.neg)
674 self.assertRaises(TypeError, next, map(10, range(5)))
675 self.assertRaises(ValueError, next, map(errfunc, [4], [5]))
676 self.assertRaises(TypeError, next, map(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000677
678 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000679 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
680 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000681 self.assertEqual(take(3, starmap(operator.pow, zip(count(), count(1)))),
Raymond Hettinger02420702003-06-29 20:36:23 +0000682 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000683 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +0000684 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
685 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000686 self.assertRaises(TypeError, starmap)
687 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000688 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
689 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
690 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000691
692 def test_islice(self):
693 for args in [ # islice(args) should agree with range(args)
694 (10, 20, 3),
695 (10, 3, 20),
696 (10, 20),
697 (10, 3),
698 (20,)
699 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000700 self.assertEqual(list(islice(range(100), *args)),
701 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000702
703 for args, tgtargs in [ # Stop when seqn is exhausted
704 ((10, 110, 3), ((10, 100, 3))),
705 ((10, 110), ((10, 100))),
706 ((110,), (100,))
707 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000708 self.assertEqual(list(islice(range(100), *args)),
709 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000710
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000711 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +0000712 self.assertEqual(list(islice(range(10), None)), list(range(10)))
713 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
714 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
715 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
716 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000717
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000718 # Test number of items consumed SF #1171417
719 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +0000720 self.assertEqual(list(islice(it, 3)), list(range(3)))
721 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000722
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000723 # Test invalid arguments
Guido van Rossum805365e2007-05-07 22:24:25 +0000724 self.assertRaises(TypeError, islice, range(10))
725 self.assertRaises(TypeError, islice, range(10), 1, 2, 3, 4)
726 self.assertRaises(ValueError, islice, range(10), -5, 10, 1)
727 self.assertRaises(ValueError, islice, range(10), 1, -5, -1)
728 self.assertRaises(ValueError, islice, range(10), 1, 10, -1)
729 self.assertRaises(ValueError, islice, range(10), 1, 10, 0)
730 self.assertRaises(ValueError, islice, range(10), 'a')
731 self.assertRaises(ValueError, islice, range(10), 'a', 1)
732 self.assertRaises(ValueError, islice, range(10), 1, 'a')
733 self.assertRaises(ValueError, islice, range(10), 'a', 1, 1)
734 self.assertRaises(ValueError, islice, range(10), 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +0000735 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000736
737 def test_takewhile(self):
738 data = [1, 3, 5, 20, 2, 4, 6, 8]
739 underten = lambda x: x<10
740 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000741 self.assertEqual(list(takewhile(underten, [])), [])
742 self.assertRaises(TypeError, takewhile)
743 self.assertRaises(TypeError, takewhile, operator.pow)
744 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000745 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
746 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000747 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
748 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +0000749 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000750
751 def test_dropwhile(self):
752 data = [1, 3, 5, 20, 2, 4, 6, 8]
753 underten = lambda x: x<10
754 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000755 self.assertEqual(list(dropwhile(underten, [])), [])
756 self.assertRaises(TypeError, dropwhile)
757 self.assertRaises(TypeError, dropwhile, operator.pow)
758 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000759 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
760 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000761
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000762 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000763 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000764 def irange(n):
Guido van Rossum805365e2007-05-07 22:24:25 +0000765 for i in range(n):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000766 yield i
767
768 a, b = tee([]) # test empty iterator
769 self.assertEqual(list(a), [])
770 self.assertEqual(list(b), [])
771
772 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +0000773 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000774
775 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +0000776 self.assertEqual(list(a), list(range(n)))
777 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000778
779 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000780 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000781 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000782 del a
Guido van Rossum805365e2007-05-07 22:24:25 +0000783 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000784
785 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000786 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000787 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000788 del b
Guido van Rossum805365e2007-05-07 22:24:25 +0000789 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000790
Guido van Rossum805365e2007-05-07 22:24:25 +0000791 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000792 order = [0]*n + [1]*n
793 random.shuffle(order)
794 lists = ([], [])
795 its = tee(irange(n))
796 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +0000797 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000798 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +0000799 self.assertEqual(lists[0], list(range(n)))
800 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000801
Raymond Hettingerad983e72003-11-12 14:32:26 +0000802 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000803 self.assertRaises(TypeError, tee)
804 self.assertRaises(TypeError, tee, 3)
805 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000806 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000807
Raymond Hettingerad983e72003-11-12 14:32:26 +0000808 # tee object should be instantiable
809 a, b = tee('abc')
810 c = type(a)('def')
811 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000812
Raymond Hettingerad983e72003-11-12 14:32:26 +0000813 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +0000814 a, b, c = tee(range(2000), 3)
815 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000816 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +0000817 self.assertEqual(list(b), list(range(2000)))
818 self.assertEqual([next(c), next(c)], list(range(2)))
819 self.assertEqual(list(a), list(range(100,2000)))
820 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000821
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000822 # test values of n
823 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000824 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +0000825 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000826 result = tee('abc', n)
827 self.assertEqual(type(result), tuple)
828 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000829 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000830
Raymond Hettingerad983e72003-11-12 14:32:26 +0000831 # tee pass-through to copyable iterator
832 a, b = tee('abc')
833 c, d = tee(a)
834 self.assert_(a is c)
835
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000836 # test tee_new
837 t1, t2 = tee('abc')
838 tnew = type(t1)
839 self.assertRaises(TypeError, tnew)
840 self.assertRaises(TypeError, tnew, 10)
841 t3 = tnew(t1)
842 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000843
Raymond Hettingera9f60922004-10-17 16:40:14 +0000844 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +0000845 a, b = tee(range(10))
Raymond Hettingera9f60922004-10-17 16:40:14 +0000846 p = proxy(a)
847 self.assertEqual(getattr(p, '__class__'), type(b))
848 del a
849 self.assertRaises(ReferenceError, getattr, p, '__class__')
850
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000851 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000852 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000853
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000854 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +0000855 self.assertRaises(StopIteration, next, f([]))
856 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000857
Georg Brandla18af4e2007-04-21 15:47:16 +0000858 self.assertRaises(StopIteration, next, islice([], None))
859 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000860
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000861 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +0000862 self.assertRaises(StopIteration, next, p)
863 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000864 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +0000865 self.assertRaises(StopIteration, next, p)
866 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000867
Georg Brandla18af4e2007-04-21 15:47:16 +0000868 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000869
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000870 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +0000871 self.assertRaises(StopIteration, next, f(lambda x:x, []))
872 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000873
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000874class TestExamples(unittest.TestCase):
875
876 def test_chain(self):
877 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
878
879 def test_chain_from_iterable(self):
880 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
881
882 def test_combinations(self):
883 self.assertEqual(list(combinations('ABCD', 2)),
884 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
885 self.assertEqual(list(combinations(range(4), 3)),
886 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
887
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000888 def test_combinations_with_replacement(self):
889 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
890 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
891
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000892 def test_compress(self):
893 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
894
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000895 def test_count(self):
896 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
897
898 def test_cycle(self):
899 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
900
901 def test_dropwhile(self):
902 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
903
904 def test_groupby(self):
905 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
906 list('ABCDAB'))
907 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
908 [list('AAAA'), list('BBB'), list('CC'), list('D')])
909
910 def test_filter(self):
911 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
912
913 def test_filterfalse(self):
914 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
915
916 def test_map(self):
917 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
918
919 def test_islice(self):
920 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
921 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
922 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
923 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
924
925 def test_zip(self):
926 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
927
928 def test_zip_longest(self):
929 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
930 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
931
932 def test_permutations(self):
933 self.assertEqual(list(permutations('ABCD', 2)),
934 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
935 self.assertEqual(list(permutations(range(3))),
936 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
937
938 def test_product(self):
939 self.assertEqual(list(product('ABCD', 'xy')),
940 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
941 self.assertEqual(list(product(range(2), repeat=3)),
942 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
943 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
944
945 def test_repeat(self):
946 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
947
948 def test_stapmap(self):
949 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
950 [32, 9, 1000])
951
952 def test_takewhile(self):
953 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
954
955
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000956class TestGC(unittest.TestCase):
957
958 def makecycle(self, iterator, container):
959 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +0000960 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000961 del container, iterator
962
963 def test_chain(self):
964 a = []
965 self.makecycle(chain(a), a)
966
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000967 def test_chain_from_iterable(self):
968 a = []
969 self.makecycle(chain.from_iterable([a]), a)
970
971 def test_combinations(self):
972 a = []
973 self.makecycle(combinations([1,2,a,3], 3), a)
974
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000975 def test_combinations_with_replacement(self):
976 a = []
977 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
978
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000979 def test_compress(self):
980 a = []
981 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
982
Raymond Hettingerd6280f42009-02-16 20:50:56 +0000983 def test_count(self):
984 a = []
985 Int = type('Int', (int,), dict(x=a))
986 self.makecycle(count(Int(0), Int(1)), a)
987
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000988 def test_cycle(self):
989 a = []
990 self.makecycle(cycle([a]*2), a)
991
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000992 def test_dropwhile(self):
993 a = []
994 self.makecycle(dropwhile(bool, [0, a, a]), a)
995
996 def test_groupby(self):
997 a = []
998 self.makecycle(groupby([a]*2, lambda x:x), a)
999
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001000 def test_issue2246(self):
1001 # Issue 2246 -- the _grouper iterator was not included in GC
1002 n = 10
1003 keyfunc = lambda x: x
1004 for i, j in groupby(range(n), key=keyfunc):
1005 keyfunc.__dict__.setdefault('x',[]).append(j)
1006
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001007 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001008 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001009 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001010
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001011 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001012 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001013 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001014
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001015 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001016 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001017 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001018
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001019 def test_zip_longest(self):
1020 a = []
1021 self.makecycle(zip_longest([a]*2, [a]*3), a)
1022 b = [a, None]
1023 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1024
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001025 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001026 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001027 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001028
1029 def test_islice(self):
1030 a = []
1031 self.makecycle(islice([a]*2, None), a)
1032
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001033 def test_permutations(self):
1034 a = []
1035 self.makecycle(permutations([1,2,a,3], 3), a)
1036
1037 def test_product(self):
1038 a = []
1039 self.makecycle(product([1,2,a,3], repeat=3), a)
1040
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001041 def test_repeat(self):
1042 a = []
1043 self.makecycle(repeat(a), a)
1044
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001045 def test_starmap(self):
1046 a = []
1047 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1048
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001049 def test_takewhile(self):
1050 a = []
1051 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1052
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001053def R(seqn):
1054 'Regular generator'
1055 for i in seqn:
1056 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001057
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001058class G:
1059 'Sequence using __getitem__'
1060 def __init__(self, seqn):
1061 self.seqn = seqn
1062 def __getitem__(self, i):
1063 return self.seqn[i]
1064
1065class I:
1066 'Sequence using iterator protocol'
1067 def __init__(self, seqn):
1068 self.seqn = seqn
1069 self.i = 0
1070 def __iter__(self):
1071 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001072 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001073 if self.i >= len(self.seqn): raise StopIteration
1074 v = self.seqn[self.i]
1075 self.i += 1
1076 return v
1077
1078class Ig:
1079 'Sequence using iterator protocol defined with a generator'
1080 def __init__(self, seqn):
1081 self.seqn = seqn
1082 self.i = 0
1083 def __iter__(self):
1084 for val in self.seqn:
1085 yield val
1086
1087class X:
1088 'Missing __getitem__ and __iter__'
1089 def __init__(self, seqn):
1090 self.seqn = seqn
1091 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001092 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001093 if self.i >= len(self.seqn): raise StopIteration
1094 v = self.seqn[self.i]
1095 self.i += 1
1096 return v
1097
1098class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001099 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001100 def __init__(self, seqn):
1101 self.seqn = seqn
1102 self.i = 0
1103 def __iter__(self):
1104 return self
1105
1106class E:
1107 'Test propagation of exceptions'
1108 def __init__(self, seqn):
1109 self.seqn = seqn
1110 self.i = 0
1111 def __iter__(self):
1112 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001113 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001114 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001115
1116class S:
1117 'Test immediate stop'
1118 def __init__(self, seqn):
1119 pass
1120 def __iter__(self):
1121 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001122 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001123 raise StopIteration
1124
1125def L(seqn):
1126 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001127 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001128
1129
1130class TestVariousIteratorArgs(unittest.TestCase):
1131
1132 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001133 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001134 for g in (G, I, Ig, S, L, R):
1135 self.assertEqual(list(chain(g(s))), list(g(s)))
1136 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001137 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001138 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001139 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1140
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001141 def test_compress(self):
1142 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1143 n = len(s)
1144 for g in (G, I, Ig, S, L, R):
1145 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1146 self.assertRaises(TypeError, compress, X(s), repeat(1))
1147 self.assertRaises(TypeError, compress, N(s), repeat(1))
1148 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1149
Christian Heimesc3f30c42008-02-22 16:37:40 +00001150 def test_product(self):
1151 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1152 self.assertRaises(TypeError, product, X(s))
1153 self.assertRaises(TypeError, product, N(s))
1154 self.assertRaises(ZeroDivisionError, product, E(s))
1155
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001156 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001157 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001158 for g in (G, I, Ig, S, L, R):
1159 tgtlen = len(s) * 3
1160 expected = list(g(s))*3
1161 actual = list(islice(cycle(g(s)), tgtlen))
1162 self.assertEqual(actual, expected)
1163 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001164 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001165 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1166
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001167 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001168 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001169 for g in (G, I, Ig, S, L, R):
1170 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1171 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001172 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001173 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1174
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001175 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001176 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001177 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001178 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001179 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001180 self.assertRaises(TypeError, filter, isEven, X(s))
1181 self.assertRaises(TypeError, filter, isEven, N(s))
1182 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001183
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001184 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001185 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001186 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001187 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001188 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001189 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1190 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1191 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001192
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001193 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001194 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001195 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001196 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1197 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1198 self.assertRaises(TypeError, zip, X(s))
1199 self.assertRaises(TypeError, zip, N(s))
1200 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001201
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001202 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001203 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001204 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001205 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1206 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1207 self.assertRaises(TypeError, zip_longest, X(s))
1208 self.assertRaises(TypeError, zip_longest, N(s))
1209 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001210
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001211 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001212 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001213 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001214 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001215 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001216 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001217 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001218 self.assertRaises(TypeError, map, onearg, X(s))
1219 self.assertRaises(TypeError, map, onearg, N(s))
1220 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001221
1222 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001223 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001224 for g in (G, I, Ig, S, L, R):
1225 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1226 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001227 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001228 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1229
1230 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001231 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001232 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001233 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001234 self.assertEqual(list(starmap(operator.pow, g(ss))),
1235 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001236 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001237 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001238 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1239
1240 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001241 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001242 for g in (G, I, Ig, S, L, R):
1243 tgt = []
1244 for elem in g(s):
1245 if not isEven(elem): break
1246 tgt.append(elem)
1247 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1248 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001249 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001250 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1251
1252 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001253 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001254 for g in (G, I, Ig, S, L, R):
1255 tgt = []
1256 for elem in g(s):
1257 if not tgt and isOdd(elem): continue
1258 tgt.append(elem)
1259 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1260 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001261 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001262 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1263
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001264 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001265 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001266 for g in (G, I, Ig, S, L, R):
1267 it1, it2 = tee(g(s))
1268 self.assertEqual(list(it1), list(g(s)))
1269 self.assertEqual(list(it2), list(g(s)))
1270 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001271 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001272 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1273
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001274class LengthTransparency(unittest.TestCase):
1275
1276 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001277 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001278 self.assertEqual(len(repeat(None, 50)), 50)
1279 self.assertRaises(TypeError, len, repeat(None))
1280
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001281class RegressionTests(unittest.TestCase):
1282
1283 def test_sf_793826(self):
1284 # Fix Armin Rigo's successful efforts to wreak havoc
1285
1286 def mutatingtuple(tuple1, f, tuple2):
1287 # this builds a tuple t which is a copy of tuple1,
1288 # then calls f(t), then mutates t to be equal to tuple2
1289 # (needs len(tuple1) == len(tuple2)).
1290 def g(value, first=[1]):
1291 if first:
1292 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00001293 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001294 return value
1295 items = list(tuple2)
1296 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001297 gen = map(g, items)
1298 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00001299 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001300
1301 def f(t):
1302 global T
1303 T = t
1304 first[:] = list(T)
1305
1306 first = []
1307 mutatingtuple((1,2,3), f, (4,5,6))
1308 second = list(T)
1309 self.assertEqual(first, second)
1310
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001311
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001312 def test_sf_950057(self):
1313 # Make sure that chain() and cycle() catch exceptions immediately
1314 # rather than when shifting between input sources
1315
1316 def gen1():
1317 hist.append(0)
1318 yield 1
1319 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001320 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001321 hist.append(2)
1322
1323 def gen2(x):
1324 hist.append(3)
1325 yield 2
1326 hist.append(4)
1327 if x:
1328 raise StopIteration
1329
1330 hist = []
1331 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1332 self.assertEqual(hist, [0,1])
1333
1334 hist = []
1335 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1336 self.assertEqual(hist, [0,1])
1337
1338 hist = []
1339 self.assertRaises(AssertionError, list, cycle(gen1()))
1340 self.assertEqual(hist, [0,1])
1341
Thomas Woutersb2137042007-02-01 18:02:27 +00001342class SubclassWithKwargsTest(unittest.TestCase):
1343 def test_keywords_in_subclass(self):
1344 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001345 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001346 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00001347 class Subclass(cls):
1348 def __init__(self, newarg=None, *args):
1349 cls.__init__(self, *args)
1350 try:
1351 Subclass(newarg=1)
1352 except TypeError as err:
1353 # we expect type errors because of wrong argument count
1354 self.failIf("does not take keyword arguments" in err.args[0])
1355
1356
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001357libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001358
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001359
1360>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001361>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00001362... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001363...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001364Check 1200 is for $120.15
1365Check 1201 is for $764.05
1366Check 1202 is for $823.14
1367
1368>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001369>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00001370... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00001371...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000013721
13738
137427
1375
1376>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001377>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00001378... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00001379...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001380Alex
1381Laura
1382Martin
1383Walter
1384Samuele
1385
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001386>>> from operator import itemgetter
1387>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00001388>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001389>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001390... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001391...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000013921 ['a', 'c', 'e']
13932 ['b', 'd', 'f']
13943 ['g']
1395
Raymond Hettinger734fb572004-01-20 20:04:40 +00001396# Find runs of consecutive numbers using groupby. The key to the solution
1397# is differencing with a range so that consecutive numbers all appear in
1398# same group.
1399>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00001400>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001401... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001402...
Raymond Hettinger734fb572004-01-20 20:04:40 +00001403[1]
1404[4, 5, 6]
1405[10]
1406[15, 16, 17, 18]
1407[22]
1408[25, 26, 27, 28]
1409
Georg Brandl3dbca812008-07-23 16:10:53 +00001410>>> def take(n, iterable):
1411... "Return first n items of the iterable as a list"
1412... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001413
Georg Brandl3dbca812008-07-23 16:10:53 +00001414>>> def enumerate(iterable, start=0):
1415... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001416
Georg Brandl3dbca812008-07-23 16:10:53 +00001417>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001418... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00001419... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001420
1421>>> def nth(iterable, n):
Raymond Hettingerd04fa312009-02-04 19:45:13 +00001422... "Returns the nth item or None"
1423... return next(islice(iterable, n, None), None)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001424
Georg Brandl3dbca812008-07-23 16:10:53 +00001425>>> def quantify(iterable, pred=bool):
1426... "Count how many times the predicate is true"
1427... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001428
Georg Brandl3dbca812008-07-23 16:10:53 +00001429>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001430... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00001431... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001432
Georg Brandl3dbca812008-07-23 16:10:53 +00001433>>> def ncycles(iterable, n):
1434... "Returns the seqeuence elements n times"
1435... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001436
1437>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001438... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001439
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001440>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001441... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001442
1443>>> def repeatfunc(func, times=None, *args):
1444... "Repeat calls to func with specified arguments."
1445... " Example: repeatfunc(random.random)"
1446... if times is None:
1447... return starmap(func, repeat(args))
1448... else:
1449... return starmap(func, repeat(args, times))
1450
Raymond Hettingerd591f662003-10-26 15:34:50 +00001451>>> def pairwise(iterable):
1452... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1453... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001454... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00001455... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001456... except StopIteration:
1457... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001458... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001459
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001460>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001461... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001462... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00001463... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001464
1465>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001466... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001467... # Recipe credited to George Sakkis
1468... pending = len(iterables)
1469... nexts = cycle(iter(it).__next__ for it in iterables)
1470... while pending:
1471... try:
1472... for next in nexts:
1473... yield next()
1474... except StopIteration:
1475... pending -= 1
1476... nexts = cycle(islice(nexts, pending))
1477
1478>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00001479... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1480... s = list(iterable)
1481... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001482
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001483>>> def unique_everseen(iterable, key=None):
1484... "List unique elements, preserving order. Remember all elements ever seen."
1485... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1486... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1487... seen = set()
1488... seen_add = seen.add
1489... if key is None:
1490... for element in iterable:
1491... if element not in seen:
1492... seen_add(element)
1493... yield element
1494... else:
1495... for element in iterable:
1496... k = key(element)
1497... if k not in seen:
1498... seen_add(k)
1499... yield element
1500
1501>>> def unique_justseen(iterable, key=None):
1502... "List unique elements, preserving order. Remember only the element just seen."
1503... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1504... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1505... return map(next, map(itemgetter(1), groupby(iterable, key)))
1506
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001507This is not part of the examples but it tests to make sure the definitions
1508perform as purported.
1509
Raymond Hettingera098b332003-09-08 23:58:40 +00001510>>> take(10, count())
1511[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1512
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001513>>> list(enumerate('abc'))
1514[(0, 'a'), (1, 'b'), (2, 'c')]
1515
1516>>> list(islice(tabulate(lambda x: 2*x), 4))
1517[0, 2, 4, 6]
1518
1519>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00001520'd'
1521
1522>>> nth('abcde', 9) is None
1523True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001524
Guido van Rossum805365e2007-05-07 22:24:25 +00001525>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000152650
1527
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001528>>> a = [[1, 2, 3], [4, 5, 6]]
1529>>> flatten(a)
1530[1, 2, 3, 4, 5, 6]
1531
1532>>> list(repeatfunc(pow, 5, 2, 3))
1533[8, 8, 8, 8, 8]
1534
1535>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001536>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001537[0, 0, 0, 0, 0]
1538
Raymond Hettingerd591f662003-10-26 15:34:50 +00001539>>> list(pairwise('abcd'))
1540[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001541
Raymond Hettingerd591f662003-10-26 15:34:50 +00001542>>> list(pairwise([]))
1543[]
1544
1545>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001546[]
1547
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001548>>> list(islice(padnone('abc'), 0, 6))
1549['a', 'b', 'c', None, None, None]
1550
1551>>> list(ncycles('abc', 3))
1552['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1553
1554>>> dotproduct([1,2,3], [4,5,6])
155532
1556
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001557>>> list(grouper(3, 'abcdefg', 'x'))
1558[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1559
1560>>> list(roundrobin('abc', 'd', 'ef'))
1561['a', 'd', 'e', 'b', 'f', 'c']
1562
Raymond Hettingerace67332009-01-26 02:23:50 +00001563>>> list(powerset([1,2,3]))
1564[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001565
Raymond Hettinger191e8502009-01-27 13:29:43 +00001566>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1567True
1568
1569>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1570True
1571
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001572>>> list(unique_everseen('AAAABBBCCDAABBB'))
1573['A', 'B', 'C', 'D']
1574
1575>>> list(unique_everseen('ABBCcAD', str.lower))
1576['A', 'B', 'C', 'D']
1577
1578>>> list(unique_justseen('AAAABBBCCDAABBB'))
1579['A', 'B', 'C', 'D', 'A', 'B']
1580
1581>>> list(unique_justseen('ABBCcAD', str.lower))
1582['A', 'B', 'C', 'A', 'D']
1583
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001584"""
1585
1586__test__ = {'libreftest' : libreftest}
1587
1588def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001589 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00001590 RegressionTests, LengthTransparency,
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001591 SubclassWithKwargsTest, TestExamples)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001592 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001593
1594 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001595 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001596 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001597 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00001598 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001599 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001600 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001601 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001602 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001603
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001604 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001605 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001606
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001607if __name__ == "__main__":
1608 test_main(verbose=True)