blob: a1875ffda4ab04160e03afb161476971b80cba53 [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
Raymond Hettinger6c8ee7a2009-11-30 21:55:17 +000010import copy
11import pickle
Christian Heimesc3f30c42008-02-22 16:37:40 +000012from functools import reduce
Benjamin Petersonee8712c2008-05-20 21:35:26 +000013maxsize = support.MAX_Py_ssize_t
Guido van Rossum360e4b82007-05-14 22:51:27 +000014minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000015
Guido van Rossum801f0d72006-08-24 19:48:10 +000016def lzip(*args):
17 return list(zip(*args))
18
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000019def onearg(x):
20 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000021 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000022
23def errfunc(*args):
24 'Test function that raises an error'
25 raise ValueError
26
27def gen3():
28 'Non-restartable source sequence'
29 for i in (0, 1, 2):
30 yield i
31
32def isEven(x):
33 'Test predicate'
34 return x%2==0
35
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000036def isOdd(x):
37 'Test predicate'
38 return x%2==1
39
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000040class StopNow:
41 'Class emulating an empty iterable.'
42 def __iter__(self):
43 return self
Georg Brandla18af4e2007-04-21 15:47:16 +000044 def __next__(self):
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000045 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000046
Raymond Hettinger02420702003-06-29 20:36:23 +000047def take(n, seq):
48 'Convenience function for partially consuming a long of infinite iterable'
49 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000050
Christian Heimes78644762008-03-04 23:39:23 +000051def prod(iterable):
52 return reduce(operator.mul, iterable, 1)
53
Christian Heimes380f7f22008-02-28 11:19:05 +000054def fact(n):
55 'Factorial'
Christian Heimes78644762008-03-04 23:39:23 +000056 return prod(range(1, n+1))
57
Raymond Hettinger96ef8112003-02-01 00:10:11 +000058class TestBasicOps(unittest.TestCase):
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000059 def test_chain(self):
Christian Heimesdd15f6c2008-03-16 00:07:10 +000060
61 def chain2(*iterables):
62 'Pure python version in the docs'
63 for it in iterables:
64 for element in it:
65 yield element
66
67 for c in (chain, chain2):
68 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
69 self.assertEqual(list(c('abc')), list('abc'))
70 self.assertEqual(list(c('')), [])
71 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
72 self.assertRaises(TypeError, list,c(2, 3))
Christian Heimesf16baeb2008-02-29 14:57:44 +000073
74 def test_chain_from_iterable(self):
75 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
76 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
77 self.assertEqual(list(chain.from_iterable([''])), [])
78 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
79 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000080
Christian Heimes380f7f22008-02-28 11:19:05 +000081 def test_combinations(self):
Raymond Hettinger5bad41e2009-01-08 21:01:54 +000082 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Christian Heimes380f7f22008-02-28 11:19:05 +000083 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Christian Heimes78644762008-03-04 23:39:23 +000084 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Christian Heimes380f7f22008-02-28 11:19:05 +000085 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Raymond Hettinger5bad41e2009-01-08 21:01:54 +000086 self.assertEqual(list(combinations('abc', 32)), []) # r > n
Christian Heimes380f7f22008-02-28 11:19:05 +000087 self.assertEqual(list(combinations(range(4), 3)),
88 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
Christian Heimes78644762008-03-04 23:39:23 +000089
90 def combinations1(iterable, r):
91 'Pure python version shown in the docs'
92 pool = tuple(iterable)
93 n = len(pool)
Raymond Hettinger5bad41e2009-01-08 21:01:54 +000094 if r > n:
95 return
Christian Heimes78644762008-03-04 23:39:23 +000096 indices = list(range(r))
97 yield tuple(pool[i] for i in indices)
98 while 1:
99 for i in reversed(range(r)):
100 if indices[i] != i + n - r:
101 break
102 else:
103 return
104 indices[i] += 1
105 for j in range(i+1, r):
106 indices[j] = indices[j-1] + 1
107 yield tuple(pool[i] for i in indices)
108
109 def combinations2(iterable, r):
110 'Pure python version shown in the docs'
111 pool = tuple(iterable)
112 n = len(pool)
113 for indices in permutations(range(n), r):
114 if sorted(indices) == list(indices):
115 yield tuple(pool[i] for i in indices)
116
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000117 def combinations3(iterable, r):
118 'Pure python version from cwr()'
119 pool = tuple(iterable)
120 n = len(pool)
121 for indices in combinations_with_replacement(range(n), r):
122 if len(set(indices)) == r:
123 yield tuple(pool[i] for i in indices)
124
Christian Heimes78644762008-03-04 23:39:23 +0000125 for n in range(7):
Christian Heimes380f7f22008-02-28 11:19:05 +0000126 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000127 for r in range(n+2):
Christian Heimes380f7f22008-02-28 11:19:05 +0000128 result = list(combinations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000129 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 +0000130 self.assertEqual(len(result), len(set(result))) # no repeats
131 self.assertEqual(result, sorted(result)) # lexicographic order
132 for c in result:
133 self.assertEqual(len(c), r) # r-length combinations
134 self.assertEqual(len(set(c)), r) # no duplicate elements
135 self.assertEqual(list(c), sorted(c)) # keep original ordering
Georg Brandlab91fde2009-08-13 08:51:18 +0000136 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000137 self.assertEqual(list(c),
138 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000139 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000140 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000141 self.assertEqual(result, list(combinations3(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000142
143 # Test implementation detail: tuple re-use
144 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
145 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
146
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000147 def test_combinations_with_replacement(self):
148 cwr = combinations_with_replacement
149 self.assertRaises(TypeError, cwr, 'abc') # missing r argument
150 self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
151 self.assertRaises(TypeError, cwr, None) # pool is not iterable
152 self.assertRaises(ValueError, cwr, 'abc', -2) # r is negative
153 self.assertEqual(list(cwr('ABC', 2)),
154 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
155
156 def cwr1(iterable, r):
157 'Pure python version shown in the docs'
158 # number items returned: (n+r-1)! / r! / (n-1)! when n>0
159 pool = tuple(iterable)
160 n = len(pool)
161 if not n and r:
162 return
163 indices = [0] * r
164 yield tuple(pool[i] for i in indices)
165 while 1:
166 for i in reversed(range(r)):
167 if indices[i] != n - 1:
168 break
169 else:
170 return
171 indices[i:] = [indices[i] + 1] * (r - i)
172 yield tuple(pool[i] for i in indices)
173
174 def cwr2(iterable, r):
175 'Pure python version shown in the docs'
176 pool = tuple(iterable)
177 n = len(pool)
178 for indices in product(range(n), repeat=r):
179 if sorted(indices) == list(indices):
180 yield tuple(pool[i] for i in indices)
181
182 def numcombs(n, r):
183 if not n:
184 return 0 if r else 1
185 return fact(n+r-1) / fact(r)/ fact(n-1)
186
187 for n in range(7):
188 values = [5*x-12 for x in range(n)]
189 for r in range(n+2):
190 result = list(cwr(values, r))
191
192 self.assertEqual(len(result), numcombs(n, r)) # right number of combs
193 self.assertEqual(len(result), len(set(result))) # no repeats
194 self.assertEqual(result, sorted(result)) # lexicographic order
195
196 regular_combs = list(combinations(values, r)) # compare to combs without replacement
197 if n == 0 or r <= 1:
Ezio Melotti19f2aeb2010-11-21 01:30:29 +0000198 self.assertEqual(result, regular_combs) # cases that should be identical
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000199 else:
Georg Brandlab91fde2009-08-13 08:51:18 +0000200 self.assertTrue(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000201
202 for c in result:
203 self.assertEqual(len(c), r) # r-length combinations
204 noruns = [k for k,v in groupby(c)] # combo without consecutive repeats
205 self.assertEqual(len(noruns), len(set(noruns))) # no repeats other than consecutive
206 self.assertEqual(list(c), sorted(c)) # keep original ordering
Georg Brandlab91fde2009-08-13 08:51:18 +0000207 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000208 self.assertEqual(noruns,
209 [e for e in values if e in c]) # comb is a subsequence of the input iterable
210 self.assertEqual(result, list(cwr1(values, r))) # matches first pure python version
211 self.assertEqual(result, list(cwr2(values, r))) # matches second pure python version
212
213 # Test implementation detail: tuple re-use
214 self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
215 self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
216
Christian Heimes78644762008-03-04 23:39:23 +0000217 def test_permutations(self):
218 self.assertRaises(TypeError, permutations) # too few arguments
219 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000220 self.assertRaises(TypeError, permutations, None) # pool is not iterable
221 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000222 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000223 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Christian Heimes78644762008-03-04 23:39:23 +0000224 self.assertEqual(list(permutations(range(3), 2)),
225 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
226
227 def permutations1(iterable, r=None):
228 'Pure python version shown in the docs'
229 pool = tuple(iterable)
230 n = len(pool)
231 r = n if r is None else r
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000232 if r > n:
233 return
Christian Heimes78644762008-03-04 23:39:23 +0000234 indices = list(range(n))
235 cycles = list(range(n-r+1, n+1))[::-1]
236 yield tuple(pool[i] for i in indices[:r])
237 while n:
238 for i in reversed(range(r)):
239 cycles[i] -= 1
240 if cycles[i] == 0:
241 indices[i:] = indices[i+1:] + indices[i:i+1]
242 cycles[i] = n - i
243 else:
244 j = cycles[i]
245 indices[i], indices[-j] = indices[-j], indices[i]
246 yield tuple(pool[i] for i in indices[:r])
247 break
248 else:
249 return
250
251 def permutations2(iterable, r=None):
252 'Pure python version shown in the docs'
253 pool = tuple(iterable)
254 n = len(pool)
255 r = n if r is None else r
256 for indices in product(range(n), repeat=r):
257 if len(set(indices)) == r:
258 yield tuple(pool[i] for i in indices)
259
260 for n in range(7):
261 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000262 for r in range(n+2):
Christian Heimes78644762008-03-04 23:39:23 +0000263 result = list(permutations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000264 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 +0000265 self.assertEqual(len(result), len(set(result))) # no repeats
266 self.assertEqual(result, sorted(result)) # lexicographic order
267 for p in result:
268 self.assertEqual(len(p), r) # r-length permutations
269 self.assertEqual(len(set(p)), r) # no duplicate elements
Georg Brandlab91fde2009-08-13 08:51:18 +0000270 self.assertTrue(all(e in values for e in p)) # elements taken from input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000271 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000272 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000273 if r == n:
274 self.assertEqual(result, list(permutations(values, None))) # test r as None
275 self.assertEqual(result, list(permutations(values))) # test default r
276
277 # Test implementation detail: tuple re-use
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000278 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Christian Heimes78644762008-03-04 23:39:23 +0000279 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Christian Heimes380f7f22008-02-28 11:19:05 +0000280
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000281 def test_combinatorics(self):
282 # Test relationships between product(), permutations(),
283 # combinations() and combinations_with_replacement().
284
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000285 for n in range(6):
286 s = 'ABCDEFG'[:n]
287 for r in range(8):
288 prod = list(product(s, repeat=r))
289 cwr = list(combinations_with_replacement(s, r))
290 perm = list(permutations(s, r))
291 comb = list(combinations(s, r))
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000292
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000293 # Check size
Ezio Melotti19f2aeb2010-11-21 01:30:29 +0000294 self.assertEqual(len(prod), n**r)
295 self.assertEqual(len(cwr), (fact(n+r-1) / fact(r)/ fact(n-1)) if n else (not r))
296 self.assertEqual(len(perm), 0 if r>n else fact(n) / fact(n-r))
297 self.assertEqual(len(comb), 0 if r>n else fact(n) / fact(r) / fact(n-r))
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000298
299 # Check lexicographic order without repeated tuples
Ezio Melotti19f2aeb2010-11-21 01:30:29 +0000300 self.assertEqual(prod, sorted(set(prod)))
301 self.assertEqual(cwr, sorted(set(cwr)))
302 self.assertEqual(perm, sorted(set(perm)))
303 self.assertEqual(comb, sorted(set(comb)))
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000304
305 # Check interrelationships
Ezio Melotti19f2aeb2010-11-21 01:30:29 +0000306 self.assertEqual(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
307 self.assertEqual(perm, [t for t in prod if len(set(t))==r]) # perm: prods with no dups
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000308 self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
309 self.assertEqual(comb, [t for t in cwr if len(set(t))==r]) # comb: cwrs without dups
310 self.assertEqual(comb, list(filter(set(cwr).__contains__, perm))) # comb: perm that is a cwr
311 self.assertEqual(comb, list(filter(set(perm).__contains__, cwr))) # comb: cwr that is a perm
312 self.assertEqual(comb, sorted(set(cwr) & set(perm))) # comb: both a cwr and a perm
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000313
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000314 def test_compress(self):
Raymond Hettinger15a49502009-02-19 02:17:09 +0000315 self.assertEqual(list(compress(data='ABCDEF', selectors=[1,0,1,0,1,1])), list('ACEF'))
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000316 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
317 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
318 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
319 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
320 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
321 n = 10000
322 data = chain.from_iterable(repeat(range(6), n))
323 selectors = chain.from_iterable(repeat((0, 1)))
324 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
325 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
326 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
327 self.assertRaises(TypeError, compress, range(6)) # too few args
328 self.assertRaises(TypeError, compress, range(6), None) # too many args
329
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000330 def test_count(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000331 self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
332 self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
333 self.assertEqual(take(2, lzip('abc',count(3))), [('a', 3), ('b', 4)])
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000334 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
335 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000336 self.assertRaises(TypeError, count, 2, 3, 4)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000337 self.assertRaises(TypeError, count, 'a')
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000338 self.assertEqual(list(islice(count(maxsize-5), 10)),
339 list(range(maxsize-5, maxsize+5)))
340 self.assertEqual(list(islice(count(-maxsize-5), 10)),
341 list(range(-maxsize-5, -maxsize+5)))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000342 c = count(3)
343 self.assertEqual(repr(c), 'count(3)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000344 next(c)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000345 self.assertEqual(repr(c), 'count(4)')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000346 c = count(-9)
347 self.assertEqual(repr(c), 'count(-9)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000348 next(c)
Raymond Hettinger30729212009-02-12 06:28:27 +0000349 self.assertEqual(repr(count(10.25)), 'count(10.25)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000350 self.assertEqual(next(c), -8)
Christian Heimesa37d4c62007-12-04 23:02:19 +0000351 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 +0000352 # Test repr (ignoring the L in longs)
353 r1 = repr(count(i)).replace('L', '')
354 r2 = 'count(%r)'.__mod__(i).replace('L', '')
355 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000356
Raymond Hettinger6c8ee7a2009-11-30 21:55:17 +0000357 # check copy, deepcopy, pickle
358 for value in -3, 3, maxsize-5, maxsize+5:
359 c = count(value)
360 self.assertEqual(next(copy.copy(c)), value)
361 self.assertEqual(next(copy.deepcopy(c)), value)
362 self.assertEqual(next(pickle.loads(pickle.dumps(c))), value)
363
Raymond Hettinger30729212009-02-12 06:28:27 +0000364 def test_count_with_stride(self):
365 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingereb13fdd2009-02-21 22:30:12 +0000366 self.assertEqual(lzip('abc',count(start=2,step=3)),
367 [('a', 2), ('b', 5), ('c', 8)])
368 self.assertEqual(lzip('abc',count(step=-1)),
369 [('a', 0), ('b', -1), ('c', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000370 self.assertEqual(lzip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
371 self.assertEqual(lzip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000372 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000373 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
374 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
375 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettinger8a882a72009-02-12 12:08:18 +0000376 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
377 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerde09acf2009-02-12 12:53:53 +0000378 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
379 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000380 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
381 c = count(3, 5)
382 self.assertEqual(repr(c), 'count(3, 5)')
383 next(c)
384 self.assertEqual(repr(c), 'count(8, 5)')
385 c = count(-9, 0)
386 self.assertEqual(repr(c), 'count(-9, 0)')
387 next(c)
388 self.assertEqual(repr(c), 'count(-9, 0)')
389 c = count(-9, -3)
390 self.assertEqual(repr(c), 'count(-9, -3)')
391 next(c)
392 self.assertEqual(repr(c), 'count(-12, -3)')
393 self.assertEqual(repr(c), 'count(-12, -3)')
394 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
395 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
396 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
397 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
398 for j in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 1, 10, sys.maxsize-5, sys.maxsize+5):
399 # Test repr (ignoring the L in longs)
400 r1 = repr(count(i, j)).replace('L', '')
401 if j == 1:
402 r2 = ('count(%r)' % i).replace('L', '')
403 else:
404 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
405 self.assertEqual(r1, r2)
406
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000407 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000408 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000409 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000410 self.assertRaises(TypeError, cycle)
411 self.assertRaises(TypeError, cycle, 5)
412 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000413
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000414 def test_groupby(self):
415 # Check whether it accepts arguments correctly
416 self.assertEqual([], list(groupby([])))
417 self.assertEqual([], list(groupby([], key=id)))
418 self.assertRaises(TypeError, list, groupby('abc', []))
419 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000420 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000421
422 # Check normal input
423 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
424 (2,15,22), (3,16,23), (3,17,23)]
425 dup = []
426 for k, g in groupby(s, lambda r:r[0]):
427 for elem in g:
428 self.assertEqual(k, elem[0])
429 dup.append(elem)
430 self.assertEqual(s, dup)
431
432 # Check nested case
433 dup = []
434 for k, g in groupby(s, lambda r:r[0]):
435 for ik, ig in groupby(g, lambda r:r[2]):
436 for elem in ig:
437 self.assertEqual(k, elem[0])
438 self.assertEqual(ik, elem[2])
439 dup.append(elem)
440 self.assertEqual(s, dup)
441
442 # Check case where inner iterator is not used
443 keys = [k for k, g in groupby(s, lambda r:r[0])]
444 expectedkeys = set([r[0] for r in s])
445 self.assertEqual(set(keys), expectedkeys)
446 self.assertEqual(len(keys), len(expectedkeys))
447
448 # Exercise pipes and filters style
449 s = 'abracadabra'
450 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000451 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000452 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
453 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000454 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000455 self.assertEqual(r, ['a', 'b', 'r'])
456 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000457 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000458 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
459 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000460 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000461 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
462
Georg Brandla18af4e2007-04-21 15:47:16 +0000463 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000464 class ExpectedError(Exception):
465 pass
466 def delayed_raise(n=0):
467 for i in range(n):
468 yield 'yo'
469 raise ExpectedError
470 def gulp(iterable, keyp=None, func=list):
471 return [func(g) for k, g in groupby(iterable, keyp)]
472
Georg Brandla18af4e2007-04-21 15:47:16 +0000473 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000474 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000475 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000476 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
477
478 # __cmp__ failure
479 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000480 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000481 raise ExpectedError
482 s = [DummyCmp(), DummyCmp(), None]
483
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000484 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000485 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000486 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000487 self.assertRaises(ExpectedError, gulp, s)
488
489 # keyfunc failure
490 def keyfunc(obj):
491 if keyfunc.skip > 0:
492 keyfunc.skip -= 1
493 return obj
494 else:
495 raise ExpectedError
496
497 # keyfunc failure on outer object
498 keyfunc.skip = 0
499 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
500 keyfunc.skip = 1
501 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
502
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000503 def test_filter(self):
504 self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
505 self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])
506 self.assertEqual(list(filter(bool, [0,1,0,2,0])), [1,2])
507 self.assertEqual(take(4, filter(isEven, count())), [0,2,4,6])
508 self.assertRaises(TypeError, filter)
509 self.assertRaises(TypeError, filter, lambda x:x)
510 self.assertRaises(TypeError, filter, lambda x:x, range(6), 7)
511 self.assertRaises(TypeError, filter, isEven, 3)
512 self.assertRaises(TypeError, next, filter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000513
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000514 def test_filterfalse(self):
515 self.assertEqual(list(filterfalse(isEven, range(6))), [1,3,5])
516 self.assertEqual(list(filterfalse(None, [0,1,0,2,0])), [0,0,0])
517 self.assertEqual(list(filterfalse(bool, [0,1,0,2,0])), [0,0,0])
518 self.assertEqual(take(4, filterfalse(isEven, count())), [1,3,5,7])
519 self.assertRaises(TypeError, filterfalse)
520 self.assertRaises(TypeError, filterfalse, lambda x:x)
521 self.assertRaises(TypeError, filterfalse, lambda x:x, range(6), 7)
522 self.assertRaises(TypeError, filterfalse, isEven, 3)
523 self.assertRaises(TypeError, next, filterfalse(range(6), range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000524
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000525 def test_zip(self):
526 # XXX This is rather silly now that builtin zip() calls zip()...
527 ans = [(x,y) for x, y in zip('abc',count())]
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000528 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000529 self.assertEqual(list(zip('abc', range(6))), lzip('abc', range(6)))
530 self.assertEqual(list(zip('abcdef', range(3))), lzip('abcdef', range(3)))
531 self.assertEqual(take(3,zip('abcdef', count())), lzip('abcdef', range(3)))
532 self.assertEqual(list(zip('abcdef')), lzip('abcdef'))
533 self.assertEqual(list(zip()), lzip())
534 self.assertRaises(TypeError, zip, 3)
535 self.assertRaises(TypeError, zip, range(3), 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000536 # Check tuple re-use (implementation detail)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000537 self.assertEqual([tuple(list(pair)) for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000538 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000539 self.assertEqual([pair for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000540 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000541 ids = list(map(id, zip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000542 self.assertEqual(min(ids), max(ids))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000543 ids = list(map(id, list(zip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000544 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000545
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000546 def test_ziplongest(self):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000547 for args in [
548 ['abc', range(6)],
549 [range(6), 'abc'],
550 [range(1000), range(2000,2100), range(3000,3050)],
551 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
552 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
553 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000554 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
555 for i in range(max(map(len, args)))]
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000556 self.assertEqual(list(zip_longest(*args)), target)
557 self.assertEqual(list(zip_longest(*args, **{})), target)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000558 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 +0000559 self.assertEqual(list(zip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000560
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000561 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 +0000562
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000563 self.assertEqual(list(zip_longest()), list(zip()))
564 self.assertEqual(list(zip_longest([])), list(zip([])))
565 self.assertEqual(list(zip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000566
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000567 self.assertEqual(list(zip_longest('abc', 'defg', **{})),
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000568 list(zip(list('abc')+[None], 'defg'))) # empty keyword dict
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000569 self.assertRaises(TypeError, zip_longest, 3)
570 self.assertRaises(TypeError, zip_longest, range(3), 3)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000571
572 for stmt in [
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000573 "zip_longest('abc', fv=1)",
574 "zip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000575 ]:
576 try:
577 eval(stmt, globals(), locals())
578 except TypeError:
579 pass
580 else:
581 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000582
Thomas Wouterscf297e42007-02-23 15:07:44 +0000583 # Check tuple re-use (implementation detail)
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000584 self.assertEqual([tuple(list(pair)) for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000585 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000586 self.assertEqual([pair for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000587 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000588 ids = list(map(id, zip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000589 self.assertEqual(min(ids), max(ids))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000590 ids = list(map(id, list(zip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000591 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
592
Raymond Hettingera9311a32009-11-01 21:02:38 +0000593 def test_bug_7244(self):
594
595 class Repeater:
596 # this class is similar to itertools.repeat
597 def __init__(self, o, t, e):
598 self.o = o
599 self.t = int(t)
600 self.e = e
601 def __iter__(self): # its iterator is itself
602 return self
603 def __next__(self):
604 if self.t > 0:
605 self.t -= 1
606 return self.o
607 else:
608 raise self.e
609
610 # Formerly this code in would fail in debug mode
611 # with Undetected Error and Stop Iteration
612 r1 = Repeater(1, 3, StopIteration)
613 r2 = Repeater(2, 4, StopIteration)
614 def run(r1, r2):
615 result = []
616 for i, j in zip_longest(r1, r2, fillvalue=0):
617 with support.captured_output('stdout'):
618 print((i, j))
619 result.append((i, j))
620 return result
621 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
622
623 # Formerly, the RuntimeError would be lost
624 # and StopIteration would stop as expected
625 r1 = Repeater(1, 3, RuntimeError)
626 r2 = Repeater(2, 4, StopIteration)
627 it = zip_longest(r1, r2, fillvalue=0)
628 self.assertEqual(next(it), (1, 2))
629 self.assertEqual(next(it), (1, 2))
630 self.assertEqual(next(it), (1, 2))
631 self.assertRaises(RuntimeError, next, it)
632
Christian Heimesc3f30c42008-02-22 16:37:40 +0000633 def test_product(self):
634 for args, result in [
Christian Heimes78644762008-03-04 23:39:23 +0000635 ([], [()]), # zero iterables
Christian Heimesc3f30c42008-02-22 16:37:40 +0000636 (['ab'], [('a',), ('b',)]), # one iterable
637 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
638 ([range(0), range(2), range(3)], []), # first iterable with zero length
639 ([range(2), range(0), range(3)], []), # middle iterable with zero length
640 ([range(2), range(3), range(0)], []), # last iterable with zero length
641 ]:
642 self.assertEqual(list(product(*args)), result)
Christian Heimesf16baeb2008-02-29 14:57:44 +0000643 for r in range(4):
644 self.assertEqual(list(product(*(args*r))),
645 list(product(*args, **dict(repeat=r))))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000646 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
647 self.assertRaises(TypeError, product, range(6), None)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000648
649 def product1(*args, **kwds):
650 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
651 n = len(pools)
652 if n == 0:
653 yield ()
654 return
655 if any(len(pool) == 0 for pool in pools):
656 return
657 indices = [0] * n
658 yield tuple(pool[i] for pool, i in zip(pools, indices))
659 while 1:
660 for i in reversed(range(n)): # right to left
661 if indices[i] == len(pools[i]) - 1:
662 continue
663 indices[i] += 1
664 for j in range(i+1, n):
665 indices[j] = 0
666 yield tuple(pool[i] for pool, i in zip(pools, indices))
667 break
668 else:
669 return
670
671 def product2(*args, **kwds):
672 'Pure python version used in docs'
673 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
674 result = [[]]
675 for pool in pools:
676 result = [x+[y] for x in result for y in pool]
677 for prod in result:
678 yield tuple(prod)
679
Christian Heimesc3f30c42008-02-22 16:37:40 +0000680 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
681 set('abcdefg'), range(11), tuple(range(13))]
682 for i in range(100):
683 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Christian Heimes78644762008-03-04 23:39:23 +0000684 expected_len = prod(map(len, args))
685 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000686 self.assertEqual(list(product(*args)), list(product1(*args)))
687 self.assertEqual(list(product(*args)), list(product2(*args)))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000688 args = map(iter, args)
Christian Heimes78644762008-03-04 23:39:23 +0000689 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000690
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000691 # Test implementation detail: tuple re-use
692 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
693 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000694
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000695 def test_repeat(self):
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +0000696 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
Guido van Rossum805365e2007-05-07 22:24:25 +0000697 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000698 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000699 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000700 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000701 self.assertEqual(list(repeat('a', 0)), [])
702 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000703 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000704 self.assertRaises(TypeError, repeat, None, 3, 4)
705 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000706 r = repeat(1+0j)
707 self.assertEqual(repr(r), 'repeat((1+0j))')
708 r = repeat(1+0j, 5)
709 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
710 list(r)
711 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000712
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000713 def test_map(self):
714 self.assertEqual(list(map(operator.pow, range(3), range(1,7))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000715 [0**1, 1**2, 2**3])
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000716 def tupleize(*args):
717 return args
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000718 self.assertEqual(list(map(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000719 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000720 self.assertEqual(list(map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000721 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000722 self.assertEqual(take(2,map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000723 [('a',0),('b',1)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000724 self.assertEqual(list(map(operator.pow, [])), [])
725 self.assertRaises(TypeError, map)
726 self.assertRaises(TypeError, list, map(None, range(3), range(3)))
727 self.assertRaises(TypeError, map, operator.neg)
728 self.assertRaises(TypeError, next, map(10, range(5)))
729 self.assertRaises(ValueError, next, map(errfunc, [4], [5]))
730 self.assertRaises(TypeError, next, map(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000731
732 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000733 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
734 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000735 self.assertEqual(take(3, starmap(operator.pow, zip(count(), count(1)))),
Raymond Hettinger02420702003-06-29 20:36:23 +0000736 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000737 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +0000738 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
739 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000740 self.assertRaises(TypeError, starmap)
741 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000742 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
743 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
744 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000745
746 def test_islice(self):
747 for args in [ # islice(args) should agree with range(args)
748 (10, 20, 3),
749 (10, 3, 20),
750 (10, 20),
751 (10, 3),
752 (20,)
753 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000754 self.assertEqual(list(islice(range(100), *args)),
755 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000756
757 for args, tgtargs in [ # Stop when seqn is exhausted
758 ((10, 110, 3), ((10, 100, 3))),
759 ((10, 110), ((10, 100))),
760 ((110,), (100,))
761 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000762 self.assertEqual(list(islice(range(100), *args)),
763 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000764
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000765 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +0000766 self.assertEqual(list(islice(range(10), None)), list(range(10)))
767 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
768 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
769 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
770 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000771
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000772 # Test number of items consumed SF #1171417
773 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +0000774 self.assertEqual(list(islice(it, 3)), list(range(3)))
775 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000776
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000777 # Test invalid arguments
Guido van Rossum805365e2007-05-07 22:24:25 +0000778 self.assertRaises(TypeError, islice, range(10))
779 self.assertRaises(TypeError, islice, range(10), 1, 2, 3, 4)
780 self.assertRaises(ValueError, islice, range(10), -5, 10, 1)
781 self.assertRaises(ValueError, islice, range(10), 1, -5, -1)
782 self.assertRaises(ValueError, islice, range(10), 1, 10, -1)
783 self.assertRaises(ValueError, islice, range(10), 1, 10, 0)
784 self.assertRaises(ValueError, islice, range(10), 'a')
785 self.assertRaises(ValueError, islice, range(10), 'a', 1)
786 self.assertRaises(ValueError, islice, range(10), 1, 'a')
787 self.assertRaises(ValueError, islice, range(10), 'a', 1, 1)
788 self.assertRaises(ValueError, islice, range(10), 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +0000789 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000790
Raymond Hettinger101f09e2010-11-30 03:09:05 +0000791 # Issue #10323: Less islice in a predictable state
792 c = count()
793 self.assertEqual(list(islice(c, 1, 3, 50)), [1])
794 self.assertEqual(next(c), 3)
795
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000796 def test_takewhile(self):
797 data = [1, 3, 5, 20, 2, 4, 6, 8]
798 underten = lambda x: x<10
799 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000800 self.assertEqual(list(takewhile(underten, [])), [])
801 self.assertRaises(TypeError, takewhile)
802 self.assertRaises(TypeError, takewhile, operator.pow)
803 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000804 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
805 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000806 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
807 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +0000808 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000809
810 def test_dropwhile(self):
811 data = [1, 3, 5, 20, 2, 4, 6, 8]
812 underten = lambda x: x<10
813 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000814 self.assertEqual(list(dropwhile(underten, [])), [])
815 self.assertRaises(TypeError, dropwhile)
816 self.assertRaises(TypeError, dropwhile, operator.pow)
817 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000818 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
819 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000820
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000821 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000822 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000823 def irange(n):
Guido van Rossum805365e2007-05-07 22:24:25 +0000824 for i in range(n):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000825 yield i
826
827 a, b = tee([]) # test empty iterator
828 self.assertEqual(list(a), [])
829 self.assertEqual(list(b), [])
830
831 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +0000832 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000833
834 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +0000835 self.assertEqual(list(a), list(range(n)))
836 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000837
838 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000839 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000840 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000841 del a
Guido van Rossum805365e2007-05-07 22:24:25 +0000842 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000843
844 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000845 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000846 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000847 del b
Guido van Rossum805365e2007-05-07 22:24:25 +0000848 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000849
Guido van Rossum805365e2007-05-07 22:24:25 +0000850 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000851 order = [0]*n + [1]*n
852 random.shuffle(order)
853 lists = ([], [])
854 its = tee(irange(n))
855 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +0000856 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000857 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +0000858 self.assertEqual(lists[0], list(range(n)))
859 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000860
Raymond Hettingerad983e72003-11-12 14:32:26 +0000861 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000862 self.assertRaises(TypeError, tee)
863 self.assertRaises(TypeError, tee, 3)
864 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000865 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000866
Raymond Hettingerad983e72003-11-12 14:32:26 +0000867 # tee object should be instantiable
868 a, b = tee('abc')
869 c = type(a)('def')
870 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000871
Raymond Hettingerad983e72003-11-12 14:32:26 +0000872 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +0000873 a, b, c = tee(range(2000), 3)
874 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000875 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +0000876 self.assertEqual(list(b), list(range(2000)))
877 self.assertEqual([next(c), next(c)], list(range(2)))
878 self.assertEqual(list(a), list(range(100,2000)))
879 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000880
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000881 # test values of n
882 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000883 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +0000884 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000885 result = tee('abc', n)
886 self.assertEqual(type(result), tuple)
887 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000888 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000889
Raymond Hettingerad983e72003-11-12 14:32:26 +0000890 # tee pass-through to copyable iterator
891 a, b = tee('abc')
892 c, d = tee(a)
Georg Brandlab91fde2009-08-13 08:51:18 +0000893 self.assertTrue(a is c)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000894
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000895 # test tee_new
896 t1, t2 = tee('abc')
897 tnew = type(t1)
898 self.assertRaises(TypeError, tnew)
899 self.assertRaises(TypeError, tnew, 10)
900 t3 = tnew(t1)
Georg Brandlab91fde2009-08-13 08:51:18 +0000901 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000902
Raymond Hettingera9f60922004-10-17 16:40:14 +0000903 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +0000904 a, b = tee(range(10))
Raymond Hettingera9f60922004-10-17 16:40:14 +0000905 p = proxy(a)
906 self.assertEqual(getattr(p, '__class__'), type(b))
907 del a
908 self.assertRaises(ReferenceError, getattr, p, '__class__')
909
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000910 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000911 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000912
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000913 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +0000914 self.assertRaises(StopIteration, next, f([]))
915 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000916
Georg Brandla18af4e2007-04-21 15:47:16 +0000917 self.assertRaises(StopIteration, next, islice([], None))
918 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000919
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000920 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +0000921 self.assertRaises(StopIteration, next, p)
922 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000923 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +0000924 self.assertRaises(StopIteration, next, p)
925 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000926
Georg Brandla18af4e2007-04-21 15:47:16 +0000927 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000928
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000929 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +0000930 self.assertRaises(StopIteration, next, f(lambda x:x, []))
931 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000932
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000933class TestExamples(unittest.TestCase):
934
935 def test_chain(self):
936 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
937
938 def test_chain_from_iterable(self):
939 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
940
941 def test_combinations(self):
942 self.assertEqual(list(combinations('ABCD', 2)),
943 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
944 self.assertEqual(list(combinations(range(4), 3)),
945 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
946
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000947 def test_combinations_with_replacement(self):
948 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
949 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
950
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000951 def test_compress(self):
952 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
953
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000954 def test_count(self):
955 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
956
957 def test_cycle(self):
958 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
959
960 def test_dropwhile(self):
961 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
962
963 def test_groupby(self):
964 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
965 list('ABCDAB'))
966 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
967 [list('AAAA'), list('BBB'), list('CC'), list('D')])
968
969 def test_filter(self):
970 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
971
972 def test_filterfalse(self):
973 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
974
975 def test_map(self):
976 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
977
978 def test_islice(self):
979 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
980 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
981 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
982 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
983
984 def test_zip(self):
985 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
986
987 def test_zip_longest(self):
988 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
989 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
990
991 def test_permutations(self):
992 self.assertEqual(list(permutations('ABCD', 2)),
993 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
994 self.assertEqual(list(permutations(range(3))),
995 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
996
997 def test_product(self):
998 self.assertEqual(list(product('ABCD', 'xy')),
999 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
1000 self.assertEqual(list(product(range(2), repeat=3)),
1001 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
1002 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
1003
1004 def test_repeat(self):
1005 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
1006
1007 def test_stapmap(self):
1008 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
1009 [32, 9, 1000])
1010
1011 def test_takewhile(self):
1012 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
1013
1014
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001015class TestGC(unittest.TestCase):
1016
1017 def makecycle(self, iterator, container):
1018 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +00001019 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001020 del container, iterator
1021
1022 def test_chain(self):
1023 a = []
1024 self.makecycle(chain(a), a)
1025
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001026 def test_chain_from_iterable(self):
1027 a = []
1028 self.makecycle(chain.from_iterable([a]), a)
1029
1030 def test_combinations(self):
1031 a = []
1032 self.makecycle(combinations([1,2,a,3], 3), a)
1033
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001034 def test_combinations_with_replacement(self):
1035 a = []
1036 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1037
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001038 def test_compress(self):
1039 a = []
1040 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1041
Raymond Hettingerd6280f42009-02-16 20:50:56 +00001042 def test_count(self):
1043 a = []
1044 Int = type('Int', (int,), dict(x=a))
1045 self.makecycle(count(Int(0), Int(1)), a)
1046
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001047 def test_cycle(self):
1048 a = []
1049 self.makecycle(cycle([a]*2), a)
1050
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001051 def test_dropwhile(self):
1052 a = []
1053 self.makecycle(dropwhile(bool, [0, a, a]), a)
1054
1055 def test_groupby(self):
1056 a = []
1057 self.makecycle(groupby([a]*2, lambda x:x), a)
1058
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001059 def test_issue2246(self):
1060 # Issue 2246 -- the _grouper iterator was not included in GC
1061 n = 10
1062 keyfunc = lambda x: x
1063 for i, j in groupby(range(n), key=keyfunc):
1064 keyfunc.__dict__.setdefault('x',[]).append(j)
1065
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001066 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001067 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001068 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001069
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001070 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001071 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001072 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001073
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001074 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001075 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001076 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001077
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001078 def test_zip_longest(self):
1079 a = []
1080 self.makecycle(zip_longest([a]*2, [a]*3), a)
1081 b = [a, None]
1082 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1083
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001084 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001085 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001086 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001087
1088 def test_islice(self):
1089 a = []
1090 self.makecycle(islice([a]*2, None), a)
1091
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001092 def test_permutations(self):
1093 a = []
1094 self.makecycle(permutations([1,2,a,3], 3), a)
1095
1096 def test_product(self):
1097 a = []
1098 self.makecycle(product([1,2,a,3], repeat=3), a)
1099
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001100 def test_repeat(self):
1101 a = []
1102 self.makecycle(repeat(a), a)
1103
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001104 def test_starmap(self):
1105 a = []
1106 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1107
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001108 def test_takewhile(self):
1109 a = []
1110 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1111
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001112def R(seqn):
1113 'Regular generator'
1114 for i in seqn:
1115 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001116
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001117class G:
1118 'Sequence using __getitem__'
1119 def __init__(self, seqn):
1120 self.seqn = seqn
1121 def __getitem__(self, i):
1122 return self.seqn[i]
1123
1124class I:
1125 'Sequence using iterator protocol'
1126 def __init__(self, seqn):
1127 self.seqn = seqn
1128 self.i = 0
1129 def __iter__(self):
1130 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001131 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001132 if self.i >= len(self.seqn): raise StopIteration
1133 v = self.seqn[self.i]
1134 self.i += 1
1135 return v
1136
1137class Ig:
1138 'Sequence using iterator protocol defined with a generator'
1139 def __init__(self, seqn):
1140 self.seqn = seqn
1141 self.i = 0
1142 def __iter__(self):
1143 for val in self.seqn:
1144 yield val
1145
1146class X:
1147 'Missing __getitem__ and __iter__'
1148 def __init__(self, seqn):
1149 self.seqn = seqn
1150 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001151 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001152 if self.i >= len(self.seqn): raise StopIteration
1153 v = self.seqn[self.i]
1154 self.i += 1
1155 return v
1156
1157class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001158 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001159 def __init__(self, seqn):
1160 self.seqn = seqn
1161 self.i = 0
1162 def __iter__(self):
1163 return self
1164
1165class E:
1166 'Test propagation of exceptions'
1167 def __init__(self, seqn):
1168 self.seqn = seqn
1169 self.i = 0
1170 def __iter__(self):
1171 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001172 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001173 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001174
1175class S:
1176 'Test immediate stop'
1177 def __init__(self, seqn):
1178 pass
1179 def __iter__(self):
1180 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001181 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001182 raise StopIteration
1183
1184def L(seqn):
1185 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001186 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001187
1188
1189class TestVariousIteratorArgs(unittest.TestCase):
1190
1191 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001192 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001193 for g in (G, I, Ig, S, L, R):
1194 self.assertEqual(list(chain(g(s))), list(g(s)))
1195 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001196 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001197 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001198 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1199
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001200 def test_compress(self):
1201 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1202 n = len(s)
1203 for g in (G, I, Ig, S, L, R):
1204 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1205 self.assertRaises(TypeError, compress, X(s), repeat(1))
1206 self.assertRaises(TypeError, compress, N(s), repeat(1))
1207 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1208
Christian Heimesc3f30c42008-02-22 16:37:40 +00001209 def test_product(self):
1210 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1211 self.assertRaises(TypeError, product, X(s))
1212 self.assertRaises(TypeError, product, N(s))
1213 self.assertRaises(ZeroDivisionError, product, E(s))
1214
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001215 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001216 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001217 for g in (G, I, Ig, S, L, R):
1218 tgtlen = len(s) * 3
1219 expected = list(g(s))*3
1220 actual = list(islice(cycle(g(s)), tgtlen))
1221 self.assertEqual(actual, expected)
1222 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001223 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001224 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1225
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001226 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001227 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001228 for g in (G, I, Ig, S, L, R):
1229 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1230 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001231 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001232 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1233
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001234 def test_filter(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):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001237 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001238 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001239 self.assertRaises(TypeError, filter, isEven, X(s))
1240 self.assertRaises(TypeError, filter, isEven, N(s))
1241 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001242
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001243 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001244 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001245 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001246 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001247 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001248 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1249 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1250 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001251
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001252 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001253 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001254 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001255 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1256 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1257 self.assertRaises(TypeError, zip, X(s))
1258 self.assertRaises(TypeError, zip, N(s))
1259 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001260
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001261 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001262 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001263 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001264 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1265 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1266 self.assertRaises(TypeError, zip_longest, X(s))
1267 self.assertRaises(TypeError, zip_longest, N(s))
1268 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001269
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001270 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001271 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001272 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001273 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001274 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001275 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001276 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001277 self.assertRaises(TypeError, map, onearg, X(s))
1278 self.assertRaises(TypeError, map, onearg, N(s))
1279 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001280
1281 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001282 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001283 for g in (G, I, Ig, S, L, R):
1284 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1285 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001286 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001287 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1288
1289 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001290 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001291 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001292 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001293 self.assertEqual(list(starmap(operator.pow, g(ss))),
1294 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001295 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001296 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001297 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1298
1299 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001300 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001301 for g in (G, I, Ig, S, L, R):
1302 tgt = []
1303 for elem in g(s):
1304 if not isEven(elem): break
1305 tgt.append(elem)
1306 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1307 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001308 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001309 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1310
1311 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001312 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001313 for g in (G, I, Ig, S, L, R):
1314 tgt = []
1315 for elem in g(s):
1316 if not tgt and isOdd(elem): continue
1317 tgt.append(elem)
1318 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1319 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001320 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001321 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1322
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001323 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001324 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001325 for g in (G, I, Ig, S, L, R):
1326 it1, it2 = tee(g(s))
1327 self.assertEqual(list(it1), list(g(s)))
1328 self.assertEqual(list(it2), list(g(s)))
1329 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001330 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001331 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1332
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001333class LengthTransparency(unittest.TestCase):
1334
1335 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001336 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001337 self.assertEqual(len(repeat(None, 50)), 50)
1338 self.assertRaises(TypeError, len, repeat(None))
1339
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001340class RegressionTests(unittest.TestCase):
1341
1342 def test_sf_793826(self):
1343 # Fix Armin Rigo's successful efforts to wreak havoc
1344
1345 def mutatingtuple(tuple1, f, tuple2):
1346 # this builds a tuple t which is a copy of tuple1,
1347 # then calls f(t), then mutates t to be equal to tuple2
1348 # (needs len(tuple1) == len(tuple2)).
1349 def g(value, first=[1]):
1350 if first:
1351 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00001352 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001353 return value
1354 items = list(tuple2)
1355 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001356 gen = map(g, items)
1357 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00001358 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001359
1360 def f(t):
1361 global T
1362 T = t
1363 first[:] = list(T)
1364
1365 first = []
1366 mutatingtuple((1,2,3), f, (4,5,6))
1367 second = list(T)
1368 self.assertEqual(first, second)
1369
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001370
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001371 def test_sf_950057(self):
1372 # Make sure that chain() and cycle() catch exceptions immediately
1373 # rather than when shifting between input sources
1374
1375 def gen1():
1376 hist.append(0)
1377 yield 1
1378 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001379 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001380 hist.append(2)
1381
1382 def gen2(x):
1383 hist.append(3)
1384 yield 2
1385 hist.append(4)
1386 if x:
1387 raise StopIteration
1388
1389 hist = []
1390 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1391 self.assertEqual(hist, [0,1])
1392
1393 hist = []
1394 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1395 self.assertEqual(hist, [0,1])
1396
1397 hist = []
1398 self.assertRaises(AssertionError, list, cycle(gen1()))
1399 self.assertEqual(hist, [0,1])
1400
Thomas Woutersb2137042007-02-01 18:02:27 +00001401class SubclassWithKwargsTest(unittest.TestCase):
1402 def test_keywords_in_subclass(self):
1403 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001404 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001405 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00001406 class Subclass(cls):
1407 def __init__(self, newarg=None, *args):
1408 cls.__init__(self, *args)
1409 try:
1410 Subclass(newarg=1)
1411 except TypeError as err:
1412 # we expect type errors because of wrong argument count
Georg Brandlab91fde2009-08-13 08:51:18 +00001413 self.assertFalse("does not take keyword arguments" in err.args[0])
Thomas Woutersb2137042007-02-01 18:02:27 +00001414
1415
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001416libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001417
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001418
1419>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001420>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00001421... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001422...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001423Check 1200 is for $120.15
1424Check 1201 is for $764.05
1425Check 1202 is for $823.14
1426
1427>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001428>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00001429... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00001430...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000014311
14328
143327
1434
1435>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001436>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00001437... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00001438...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001439Alex
1440Laura
1441Martin
1442Walter
1443Samuele
1444
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001445>>> from operator import itemgetter
1446>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00001447>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001448>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001449... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001450...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000014511 ['a', 'c', 'e']
14522 ['b', 'd', 'f']
14533 ['g']
1454
Raymond Hettinger734fb572004-01-20 20:04:40 +00001455# Find runs of consecutive numbers using groupby. The key to the solution
1456# is differencing with a range so that consecutive numbers all appear in
1457# same group.
1458>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00001459>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001460... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001461...
Raymond Hettinger734fb572004-01-20 20:04:40 +00001462[1]
1463[4, 5, 6]
1464[10]
1465[15, 16, 17, 18]
1466[22]
1467[25, 26, 27, 28]
1468
Georg Brandl3dbca812008-07-23 16:10:53 +00001469>>> def take(n, iterable):
1470... "Return first n items of the iterable as a list"
1471... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001472
Georg Brandl3dbca812008-07-23 16:10:53 +00001473>>> def enumerate(iterable, start=0):
1474... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001475
Georg Brandl3dbca812008-07-23 16:10:53 +00001476>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001477... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00001478... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001479
Raymond Hettingercdf8ba32009-02-19 04:45:07 +00001480>>> def nth(iterable, n, default=None):
1481... "Returns the nth item or a default value"
1482... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001483
Georg Brandl3dbca812008-07-23 16:10:53 +00001484>>> def quantify(iterable, pred=bool):
1485... "Count how many times the predicate is true"
1486... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001487
Georg Brandl3dbca812008-07-23 16:10:53 +00001488>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001489... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00001490... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001491
Georg Brandl3dbca812008-07-23 16:10:53 +00001492>>> def ncycles(iterable, n):
1493... "Returns the seqeuence elements n times"
1494... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001495
1496>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001497... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001498
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001499>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001500... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001501
1502>>> def repeatfunc(func, times=None, *args):
1503... "Repeat calls to func with specified arguments."
1504... " Example: repeatfunc(random.random)"
1505... if times is None:
1506... return starmap(func, repeat(args))
1507... else:
1508... return starmap(func, repeat(args, times))
1509
Raymond Hettingerd591f662003-10-26 15:34:50 +00001510>>> def pairwise(iterable):
1511... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1512... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001513... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00001514... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001515... except StopIteration:
1516... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001517... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001518
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001519>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001520... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001521... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00001522... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001523
1524>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001525... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001526... # Recipe credited to George Sakkis
1527... pending = len(iterables)
1528... nexts = cycle(iter(it).__next__ for it in iterables)
1529... while pending:
1530... try:
1531... for next in nexts:
1532... yield next()
1533... except StopIteration:
1534... pending -= 1
1535... nexts = cycle(islice(nexts, pending))
1536
1537>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00001538... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1539... s = list(iterable)
1540... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001541
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001542>>> def unique_everseen(iterable, key=None):
1543... "List unique elements, preserving order. Remember all elements ever seen."
1544... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1545... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1546... seen = set()
1547... seen_add = seen.add
1548... if key is None:
1549... for element in iterable:
1550... if element not in seen:
1551... seen_add(element)
1552... yield element
1553... else:
1554... for element in iterable:
1555... k = key(element)
1556... if k not in seen:
1557... seen_add(k)
1558... yield element
1559
1560>>> def unique_justseen(iterable, key=None):
1561... "List unique elements, preserving order. Remember only the element just seen."
1562... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1563... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1564... return map(next, map(itemgetter(1), groupby(iterable, key)))
1565
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001566This is not part of the examples but it tests to make sure the definitions
1567perform as purported.
1568
Raymond Hettingera098b332003-09-08 23:58:40 +00001569>>> take(10, count())
1570[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1571
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001572>>> list(enumerate('abc'))
1573[(0, 'a'), (1, 'b'), (2, 'c')]
1574
1575>>> list(islice(tabulate(lambda x: 2*x), 4))
1576[0, 2, 4, 6]
1577
1578>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00001579'd'
1580
1581>>> nth('abcde', 9) is None
1582True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001583
Guido van Rossum805365e2007-05-07 22:24:25 +00001584>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000158550
1586
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001587>>> a = [[1, 2, 3], [4, 5, 6]]
1588>>> flatten(a)
1589[1, 2, 3, 4, 5, 6]
1590
1591>>> list(repeatfunc(pow, 5, 2, 3))
1592[8, 8, 8, 8, 8]
1593
1594>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001595>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001596[0, 0, 0, 0, 0]
1597
Raymond Hettingerd591f662003-10-26 15:34:50 +00001598>>> list(pairwise('abcd'))
1599[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001600
Raymond Hettingerd591f662003-10-26 15:34:50 +00001601>>> list(pairwise([]))
1602[]
1603
1604>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001605[]
1606
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001607>>> list(islice(padnone('abc'), 0, 6))
1608['a', 'b', 'c', None, None, None]
1609
1610>>> list(ncycles('abc', 3))
1611['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1612
1613>>> dotproduct([1,2,3], [4,5,6])
161432
1615
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001616>>> list(grouper(3, 'abcdefg', 'x'))
1617[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1618
1619>>> list(roundrobin('abc', 'd', 'ef'))
1620['a', 'd', 'e', 'b', 'f', 'c']
1621
Raymond Hettingerace67332009-01-26 02:23:50 +00001622>>> list(powerset([1,2,3]))
1623[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001624
Raymond Hettinger191e8502009-01-27 13:29:43 +00001625>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1626True
1627
1628>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1629True
1630
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001631>>> list(unique_everseen('AAAABBBCCDAABBB'))
1632['A', 'B', 'C', 'D']
1633
1634>>> list(unique_everseen('ABBCcAD', str.lower))
1635['A', 'B', 'C', 'D']
1636
1637>>> list(unique_justseen('AAAABBBCCDAABBB'))
1638['A', 'B', 'C', 'D', 'A', 'B']
1639
1640>>> list(unique_justseen('ABBCcAD', str.lower))
1641['A', 'B', 'C', 'A', 'D']
1642
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001643"""
1644
1645__test__ = {'libreftest' : libreftest}
1646
1647def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001648 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00001649 RegressionTests, LengthTransparency,
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001650 SubclassWithKwargsTest, TestExamples)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001651 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001652
1653 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001654 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001655 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001656 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00001657 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001658 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001659 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001660 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001661 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001662
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001663 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001664 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001665
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001666if __name__ == "__main__":
1667 test_main(verbose=True)