blob: 9ce54d0b06ffcfa90ab81ad862db1571c5204d20 [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)))
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +0000342 self.assertEqual(list(islice(count(10, maxsize+5), 3)),
343 list(range(10, 10+3*(maxsize+5), maxsize+5)))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000344 c = count(3)
345 self.assertEqual(repr(c), 'count(3)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000346 next(c)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000347 self.assertEqual(repr(c), 'count(4)')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000348 c = count(-9)
349 self.assertEqual(repr(c), 'count(-9)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000350 next(c)
Raymond Hettinger30729212009-02-12 06:28:27 +0000351 self.assertEqual(repr(count(10.25)), 'count(10.25)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000352 self.assertEqual(next(c), -8)
Christian Heimesa37d4c62007-12-04 23:02:19 +0000353 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 +0000354 # Test repr (ignoring the L in longs)
355 r1 = repr(count(i)).replace('L', '')
356 r2 = 'count(%r)'.__mod__(i).replace('L', '')
357 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000358
Raymond Hettinger6c8ee7a2009-11-30 21:55:17 +0000359 # check copy, deepcopy, pickle
360 for value in -3, 3, maxsize-5, maxsize+5:
361 c = count(value)
362 self.assertEqual(next(copy.copy(c)), value)
363 self.assertEqual(next(copy.deepcopy(c)), value)
364 self.assertEqual(next(pickle.loads(pickle.dumps(c))), value)
365
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +0000366 #check proper internal error handling for large "step' sizes
367 count(1, maxsize+5); sys.exc_info()
368
Raymond Hettinger30729212009-02-12 06:28:27 +0000369 def test_count_with_stride(self):
370 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingereb13fdd2009-02-21 22:30:12 +0000371 self.assertEqual(lzip('abc',count(start=2,step=3)),
372 [('a', 2), ('b', 5), ('c', 8)])
373 self.assertEqual(lzip('abc',count(step=-1)),
374 [('a', 0), ('b', -1), ('c', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000375 self.assertEqual(lzip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
376 self.assertEqual(lzip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000377 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000378 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
379 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
380 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettinger8a882a72009-02-12 12:08:18 +0000381 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
382 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerde09acf2009-02-12 12:53:53 +0000383 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
384 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000385 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
386 c = count(3, 5)
387 self.assertEqual(repr(c), 'count(3, 5)')
388 next(c)
389 self.assertEqual(repr(c), 'count(8, 5)')
390 c = count(-9, 0)
391 self.assertEqual(repr(c), 'count(-9, 0)')
392 next(c)
393 self.assertEqual(repr(c), 'count(-9, 0)')
394 c = count(-9, -3)
395 self.assertEqual(repr(c), 'count(-9, -3)')
396 next(c)
397 self.assertEqual(repr(c), 'count(-12, -3)')
398 self.assertEqual(repr(c), 'count(-12, -3)')
399 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
400 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
401 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
402 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
403 for j in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 1, 10, sys.maxsize-5, sys.maxsize+5):
404 # Test repr (ignoring the L in longs)
405 r1 = repr(count(i, j)).replace('L', '')
406 if j == 1:
407 r2 = ('count(%r)' % i).replace('L', '')
408 else:
409 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
410 self.assertEqual(r1, r2)
411
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000412 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000413 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000414 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000415 self.assertRaises(TypeError, cycle)
416 self.assertRaises(TypeError, cycle, 5)
417 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000418
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000419 def test_groupby(self):
420 # Check whether it accepts arguments correctly
421 self.assertEqual([], list(groupby([])))
422 self.assertEqual([], list(groupby([], key=id)))
423 self.assertRaises(TypeError, list, groupby('abc', []))
424 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000425 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000426
427 # Check normal input
428 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
429 (2,15,22), (3,16,23), (3,17,23)]
430 dup = []
431 for k, g in groupby(s, lambda r:r[0]):
432 for elem in g:
433 self.assertEqual(k, elem[0])
434 dup.append(elem)
435 self.assertEqual(s, dup)
436
437 # Check nested case
438 dup = []
439 for k, g in groupby(s, lambda r:r[0]):
440 for ik, ig in groupby(g, lambda r:r[2]):
441 for elem in ig:
442 self.assertEqual(k, elem[0])
443 self.assertEqual(ik, elem[2])
444 dup.append(elem)
445 self.assertEqual(s, dup)
446
447 # Check case where inner iterator is not used
448 keys = [k for k, g in groupby(s, lambda r:r[0])]
449 expectedkeys = set([r[0] for r in s])
450 self.assertEqual(set(keys), expectedkeys)
451 self.assertEqual(len(keys), len(expectedkeys))
452
453 # Exercise pipes and filters style
454 s = 'abracadabra'
455 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000456 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000457 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
458 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000459 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000460 self.assertEqual(r, ['a', 'b', 'r'])
461 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000462 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000463 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
464 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000465 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000466 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
467
Georg Brandla18af4e2007-04-21 15:47:16 +0000468 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000469 class ExpectedError(Exception):
470 pass
471 def delayed_raise(n=0):
472 for i in range(n):
473 yield 'yo'
474 raise ExpectedError
475 def gulp(iterable, keyp=None, func=list):
476 return [func(g) for k, g in groupby(iterable, keyp)]
477
Georg Brandla18af4e2007-04-21 15:47:16 +0000478 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000479 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000480 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000481 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
482
483 # __cmp__ failure
484 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000485 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000486 raise ExpectedError
487 s = [DummyCmp(), DummyCmp(), None]
488
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000489 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000490 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000491 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000492 self.assertRaises(ExpectedError, gulp, s)
493
494 # keyfunc failure
495 def keyfunc(obj):
496 if keyfunc.skip > 0:
497 keyfunc.skip -= 1
498 return obj
499 else:
500 raise ExpectedError
501
502 # keyfunc failure on outer object
503 keyfunc.skip = 0
504 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
505 keyfunc.skip = 1
506 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
507
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000508 def test_filter(self):
509 self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
510 self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])
511 self.assertEqual(list(filter(bool, [0,1,0,2,0])), [1,2])
512 self.assertEqual(take(4, filter(isEven, count())), [0,2,4,6])
513 self.assertRaises(TypeError, filter)
514 self.assertRaises(TypeError, filter, lambda x:x)
515 self.assertRaises(TypeError, filter, lambda x:x, range(6), 7)
516 self.assertRaises(TypeError, filter, isEven, 3)
517 self.assertRaises(TypeError, next, filter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000518
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000519 def test_filterfalse(self):
520 self.assertEqual(list(filterfalse(isEven, range(6))), [1,3,5])
521 self.assertEqual(list(filterfalse(None, [0,1,0,2,0])), [0,0,0])
522 self.assertEqual(list(filterfalse(bool, [0,1,0,2,0])), [0,0,0])
523 self.assertEqual(take(4, filterfalse(isEven, count())), [1,3,5,7])
524 self.assertRaises(TypeError, filterfalse)
525 self.assertRaises(TypeError, filterfalse, lambda x:x)
526 self.assertRaises(TypeError, filterfalse, lambda x:x, range(6), 7)
527 self.assertRaises(TypeError, filterfalse, isEven, 3)
528 self.assertRaises(TypeError, next, filterfalse(range(6), range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000529
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000530 def test_zip(self):
531 # XXX This is rather silly now that builtin zip() calls zip()...
532 ans = [(x,y) for x, y in zip('abc',count())]
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000533 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000534 self.assertEqual(list(zip('abc', range(6))), lzip('abc', range(6)))
535 self.assertEqual(list(zip('abcdef', range(3))), lzip('abcdef', range(3)))
536 self.assertEqual(take(3,zip('abcdef', count())), lzip('abcdef', range(3)))
537 self.assertEqual(list(zip('abcdef')), lzip('abcdef'))
538 self.assertEqual(list(zip()), lzip())
539 self.assertRaises(TypeError, zip, 3)
540 self.assertRaises(TypeError, zip, range(3), 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000541 # Check tuple re-use (implementation detail)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000542 self.assertEqual([tuple(list(pair)) for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000543 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000544 self.assertEqual([pair for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000545 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000546 ids = list(map(id, zip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000547 self.assertEqual(min(ids), max(ids))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000548 ids = list(map(id, list(zip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000549 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000550
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000551 def test_ziplongest(self):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000552 for args in [
553 ['abc', range(6)],
554 [range(6), 'abc'],
555 [range(1000), range(2000,2100), range(3000,3050)],
556 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
557 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
558 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000559 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
560 for i in range(max(map(len, args)))]
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000561 self.assertEqual(list(zip_longest(*args)), target)
562 self.assertEqual(list(zip_longest(*args, **{})), target)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000563 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 +0000564 self.assertEqual(list(zip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000565
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000566 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 +0000567
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000568 self.assertEqual(list(zip_longest()), list(zip()))
569 self.assertEqual(list(zip_longest([])), list(zip([])))
570 self.assertEqual(list(zip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000571
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000572 self.assertEqual(list(zip_longest('abc', 'defg', **{})),
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000573 list(zip(list('abc')+[None], 'defg'))) # empty keyword dict
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000574 self.assertRaises(TypeError, zip_longest, 3)
575 self.assertRaises(TypeError, zip_longest, range(3), 3)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000576
577 for stmt in [
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000578 "zip_longest('abc', fv=1)",
579 "zip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000580 ]:
581 try:
582 eval(stmt, globals(), locals())
583 except TypeError:
584 pass
585 else:
586 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000587
Thomas Wouterscf297e42007-02-23 15:07:44 +0000588 # Check tuple re-use (implementation detail)
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000589 self.assertEqual([tuple(list(pair)) for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000590 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000591 self.assertEqual([pair for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000592 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000593 ids = list(map(id, zip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000594 self.assertEqual(min(ids), max(ids))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000595 ids = list(map(id, list(zip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000596 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
597
Raymond Hettingera9311a32009-11-01 21:02:38 +0000598 def test_bug_7244(self):
599
600 class Repeater:
601 # this class is similar to itertools.repeat
602 def __init__(self, o, t, e):
603 self.o = o
604 self.t = int(t)
605 self.e = e
606 def __iter__(self): # its iterator is itself
607 return self
608 def __next__(self):
609 if self.t > 0:
610 self.t -= 1
611 return self.o
612 else:
613 raise self.e
614
615 # Formerly this code in would fail in debug mode
616 # with Undetected Error and Stop Iteration
617 r1 = Repeater(1, 3, StopIteration)
618 r2 = Repeater(2, 4, StopIteration)
619 def run(r1, r2):
620 result = []
621 for i, j in zip_longest(r1, r2, fillvalue=0):
622 with support.captured_output('stdout'):
623 print((i, j))
624 result.append((i, j))
625 return result
626 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
627
628 # Formerly, the RuntimeError would be lost
629 # and StopIteration would stop as expected
630 r1 = Repeater(1, 3, RuntimeError)
631 r2 = Repeater(2, 4, StopIteration)
632 it = zip_longest(r1, r2, fillvalue=0)
633 self.assertEqual(next(it), (1, 2))
634 self.assertEqual(next(it), (1, 2))
635 self.assertEqual(next(it), (1, 2))
636 self.assertRaises(RuntimeError, next, it)
637
Christian Heimesc3f30c42008-02-22 16:37:40 +0000638 def test_product(self):
639 for args, result in [
Christian Heimes78644762008-03-04 23:39:23 +0000640 ([], [()]), # zero iterables
Christian Heimesc3f30c42008-02-22 16:37:40 +0000641 (['ab'], [('a',), ('b',)]), # one iterable
642 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
643 ([range(0), range(2), range(3)], []), # first iterable with zero length
644 ([range(2), range(0), range(3)], []), # middle iterable with zero length
645 ([range(2), range(3), range(0)], []), # last iterable with zero length
646 ]:
647 self.assertEqual(list(product(*args)), result)
Christian Heimesf16baeb2008-02-29 14:57:44 +0000648 for r in range(4):
649 self.assertEqual(list(product(*(args*r))),
650 list(product(*args, **dict(repeat=r))))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000651 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
652 self.assertRaises(TypeError, product, range(6), None)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000653
654 def product1(*args, **kwds):
655 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
656 n = len(pools)
657 if n == 0:
658 yield ()
659 return
660 if any(len(pool) == 0 for pool in pools):
661 return
662 indices = [0] * n
663 yield tuple(pool[i] for pool, i in zip(pools, indices))
664 while 1:
665 for i in reversed(range(n)): # right to left
666 if indices[i] == len(pools[i]) - 1:
667 continue
668 indices[i] += 1
669 for j in range(i+1, n):
670 indices[j] = 0
671 yield tuple(pool[i] for pool, i in zip(pools, indices))
672 break
673 else:
674 return
675
676 def product2(*args, **kwds):
677 'Pure python version used in docs'
678 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
679 result = [[]]
680 for pool in pools:
681 result = [x+[y] for x in result for y in pool]
682 for prod in result:
683 yield tuple(prod)
684
Christian Heimesc3f30c42008-02-22 16:37:40 +0000685 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
686 set('abcdefg'), range(11), tuple(range(13))]
687 for i in range(100):
688 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Christian Heimes78644762008-03-04 23:39:23 +0000689 expected_len = prod(map(len, args))
690 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000691 self.assertEqual(list(product(*args)), list(product1(*args)))
692 self.assertEqual(list(product(*args)), list(product2(*args)))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000693 args = map(iter, args)
Christian Heimes78644762008-03-04 23:39:23 +0000694 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000695
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000696 # Test implementation detail: tuple re-use
697 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
698 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000699
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000700 def test_repeat(self):
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +0000701 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
Guido van Rossum805365e2007-05-07 22:24:25 +0000702 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000703 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000704 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000705 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000706 self.assertEqual(list(repeat('a', 0)), [])
707 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000708 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000709 self.assertRaises(TypeError, repeat, None, 3, 4)
710 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000711 r = repeat(1+0j)
712 self.assertEqual(repr(r), 'repeat((1+0j))')
713 r = repeat(1+0j, 5)
714 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
715 list(r)
716 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000717
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000718 def test_map(self):
719 self.assertEqual(list(map(operator.pow, range(3), range(1,7))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000720 [0**1, 1**2, 2**3])
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000721 def tupleize(*args):
722 return args
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000723 self.assertEqual(list(map(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000724 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000725 self.assertEqual(list(map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000726 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000727 self.assertEqual(take(2,map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000728 [('a',0),('b',1)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000729 self.assertEqual(list(map(operator.pow, [])), [])
730 self.assertRaises(TypeError, map)
731 self.assertRaises(TypeError, list, map(None, range(3), range(3)))
732 self.assertRaises(TypeError, map, operator.neg)
733 self.assertRaises(TypeError, next, map(10, range(5)))
734 self.assertRaises(ValueError, next, map(errfunc, [4], [5]))
735 self.assertRaises(TypeError, next, map(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000736
737 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000738 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
739 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000740 self.assertEqual(take(3, starmap(operator.pow, zip(count(), count(1)))),
Raymond Hettinger02420702003-06-29 20:36:23 +0000741 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000742 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +0000743 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
744 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000745 self.assertRaises(TypeError, starmap)
746 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000747 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
748 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
749 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000750
751 def test_islice(self):
752 for args in [ # islice(args) should agree with range(args)
753 (10, 20, 3),
754 (10, 3, 20),
755 (10, 20),
756 (10, 3),
757 (20,)
758 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000759 self.assertEqual(list(islice(range(100), *args)),
760 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000761
762 for args, tgtargs in [ # Stop when seqn is exhausted
763 ((10, 110, 3), ((10, 100, 3))),
764 ((10, 110), ((10, 100))),
765 ((110,), (100,))
766 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000767 self.assertEqual(list(islice(range(100), *args)),
768 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000769
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000770 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +0000771 self.assertEqual(list(islice(range(10), None)), list(range(10)))
772 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
773 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
774 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
775 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000776
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000777 # Test number of items consumed SF #1171417
778 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +0000779 self.assertEqual(list(islice(it, 3)), list(range(3)))
780 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000781
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000782 # Test invalid arguments
Guido van Rossum805365e2007-05-07 22:24:25 +0000783 self.assertRaises(TypeError, islice, range(10))
784 self.assertRaises(TypeError, islice, range(10), 1, 2, 3, 4)
785 self.assertRaises(ValueError, islice, range(10), -5, 10, 1)
786 self.assertRaises(ValueError, islice, range(10), 1, -5, -1)
787 self.assertRaises(ValueError, islice, range(10), 1, 10, -1)
788 self.assertRaises(ValueError, islice, range(10), 1, 10, 0)
789 self.assertRaises(ValueError, islice, range(10), 'a')
790 self.assertRaises(ValueError, islice, range(10), 'a', 1)
791 self.assertRaises(ValueError, islice, range(10), 1, 'a')
792 self.assertRaises(ValueError, islice, range(10), 'a', 1, 1)
793 self.assertRaises(ValueError, islice, range(10), 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +0000794 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000795
Raymond Hettinger101f09e2010-11-30 03:09:05 +0000796 # Issue #10323: Less islice in a predictable state
797 c = count()
798 self.assertEqual(list(islice(c, 1, 3, 50)), [1])
799 self.assertEqual(next(c), 3)
800
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000801 def test_takewhile(self):
802 data = [1, 3, 5, 20, 2, 4, 6, 8]
803 underten = lambda x: x<10
804 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000805 self.assertEqual(list(takewhile(underten, [])), [])
806 self.assertRaises(TypeError, takewhile)
807 self.assertRaises(TypeError, takewhile, operator.pow)
808 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000809 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
810 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000811 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
812 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +0000813 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000814
815 def test_dropwhile(self):
816 data = [1, 3, 5, 20, 2, 4, 6, 8]
817 underten = lambda x: x<10
818 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000819 self.assertEqual(list(dropwhile(underten, [])), [])
820 self.assertRaises(TypeError, dropwhile)
821 self.assertRaises(TypeError, dropwhile, operator.pow)
822 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000823 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
824 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000825
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000826 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000827 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000828 def irange(n):
Guido van Rossum805365e2007-05-07 22:24:25 +0000829 for i in range(n):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000830 yield i
831
832 a, b = tee([]) # test empty iterator
833 self.assertEqual(list(a), [])
834 self.assertEqual(list(b), [])
835
836 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +0000837 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000838
839 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +0000840 self.assertEqual(list(a), list(range(n)))
841 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000842
843 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000844 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000845 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000846 del a
Guido van Rossum805365e2007-05-07 22:24:25 +0000847 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000848
849 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000850 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000851 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000852 del b
Guido van Rossum805365e2007-05-07 22:24:25 +0000853 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000854
Guido van Rossum805365e2007-05-07 22:24:25 +0000855 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000856 order = [0]*n + [1]*n
857 random.shuffle(order)
858 lists = ([], [])
859 its = tee(irange(n))
860 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +0000861 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000862 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +0000863 self.assertEqual(lists[0], list(range(n)))
864 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000865
Raymond Hettingerad983e72003-11-12 14:32:26 +0000866 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000867 self.assertRaises(TypeError, tee)
868 self.assertRaises(TypeError, tee, 3)
869 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000870 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000871
Raymond Hettingerad983e72003-11-12 14:32:26 +0000872 # tee object should be instantiable
873 a, b = tee('abc')
874 c = type(a)('def')
875 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000876
Raymond Hettingerad983e72003-11-12 14:32:26 +0000877 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +0000878 a, b, c = tee(range(2000), 3)
879 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000880 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +0000881 self.assertEqual(list(b), list(range(2000)))
882 self.assertEqual([next(c), next(c)], list(range(2)))
883 self.assertEqual(list(a), list(range(100,2000)))
884 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000885
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000886 # test values of n
887 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000888 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +0000889 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000890 result = tee('abc', n)
891 self.assertEqual(type(result), tuple)
892 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000893 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000894
Raymond Hettingerad983e72003-11-12 14:32:26 +0000895 # tee pass-through to copyable iterator
896 a, b = tee('abc')
897 c, d = tee(a)
Georg Brandlab91fde2009-08-13 08:51:18 +0000898 self.assertTrue(a is c)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000899
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000900 # test tee_new
901 t1, t2 = tee('abc')
902 tnew = type(t1)
903 self.assertRaises(TypeError, tnew)
904 self.assertRaises(TypeError, tnew, 10)
905 t3 = tnew(t1)
Georg Brandlab91fde2009-08-13 08:51:18 +0000906 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000907
Raymond Hettingera9f60922004-10-17 16:40:14 +0000908 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +0000909 a, b = tee(range(10))
Raymond Hettingera9f60922004-10-17 16:40:14 +0000910 p = proxy(a)
911 self.assertEqual(getattr(p, '__class__'), type(b))
912 del a
913 self.assertRaises(ReferenceError, getattr, p, '__class__')
914
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000915 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000916 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000917
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000918 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +0000919 self.assertRaises(StopIteration, next, f([]))
920 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000921
Georg Brandla18af4e2007-04-21 15:47:16 +0000922 self.assertRaises(StopIteration, next, islice([], None))
923 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000924
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000925 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +0000926 self.assertRaises(StopIteration, next, p)
927 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000928 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +0000929 self.assertRaises(StopIteration, next, p)
930 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000931
Georg Brandla18af4e2007-04-21 15:47:16 +0000932 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000933
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000934 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +0000935 self.assertRaises(StopIteration, next, f(lambda x:x, []))
936 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000937
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000938class TestExamples(unittest.TestCase):
939
940 def test_chain(self):
941 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
942
943 def test_chain_from_iterable(self):
944 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
945
946 def test_combinations(self):
947 self.assertEqual(list(combinations('ABCD', 2)),
948 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
949 self.assertEqual(list(combinations(range(4), 3)),
950 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
951
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000952 def test_combinations_with_replacement(self):
953 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
954 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
955
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000956 def test_compress(self):
957 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
958
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000959 def test_count(self):
960 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
961
962 def test_cycle(self):
963 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
964
965 def test_dropwhile(self):
966 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
967
968 def test_groupby(self):
969 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
970 list('ABCDAB'))
971 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
972 [list('AAAA'), list('BBB'), list('CC'), list('D')])
973
974 def test_filter(self):
975 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
976
977 def test_filterfalse(self):
978 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
979
980 def test_map(self):
981 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
982
983 def test_islice(self):
984 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
985 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
986 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
987 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
988
989 def test_zip(self):
990 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
991
992 def test_zip_longest(self):
993 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
994 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
995
996 def test_permutations(self):
997 self.assertEqual(list(permutations('ABCD', 2)),
998 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
999 self.assertEqual(list(permutations(range(3))),
1000 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
1001
1002 def test_product(self):
1003 self.assertEqual(list(product('ABCD', 'xy')),
1004 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
1005 self.assertEqual(list(product(range(2), repeat=3)),
1006 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
1007 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
1008
1009 def test_repeat(self):
1010 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
1011
1012 def test_stapmap(self):
1013 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
1014 [32, 9, 1000])
1015
1016 def test_takewhile(self):
1017 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
1018
1019
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001020class TestGC(unittest.TestCase):
1021
1022 def makecycle(self, iterator, container):
1023 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +00001024 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001025 del container, iterator
1026
1027 def test_chain(self):
1028 a = []
1029 self.makecycle(chain(a), a)
1030
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001031 def test_chain_from_iterable(self):
1032 a = []
1033 self.makecycle(chain.from_iterable([a]), a)
1034
1035 def test_combinations(self):
1036 a = []
1037 self.makecycle(combinations([1,2,a,3], 3), a)
1038
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001039 def test_combinations_with_replacement(self):
1040 a = []
1041 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1042
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001043 def test_compress(self):
1044 a = []
1045 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1046
Raymond Hettingerd6280f42009-02-16 20:50:56 +00001047 def test_count(self):
1048 a = []
1049 Int = type('Int', (int,), dict(x=a))
1050 self.makecycle(count(Int(0), Int(1)), a)
1051
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001052 def test_cycle(self):
1053 a = []
1054 self.makecycle(cycle([a]*2), a)
1055
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001056 def test_dropwhile(self):
1057 a = []
1058 self.makecycle(dropwhile(bool, [0, a, a]), a)
1059
1060 def test_groupby(self):
1061 a = []
1062 self.makecycle(groupby([a]*2, lambda x:x), a)
1063
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001064 def test_issue2246(self):
1065 # Issue 2246 -- the _grouper iterator was not included in GC
1066 n = 10
1067 keyfunc = lambda x: x
1068 for i, j in groupby(range(n), key=keyfunc):
1069 keyfunc.__dict__.setdefault('x',[]).append(j)
1070
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001071 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001072 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001073 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001074
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001075 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001076 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001077 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001078
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001079 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001080 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001081 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001082
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001083 def test_zip_longest(self):
1084 a = []
1085 self.makecycle(zip_longest([a]*2, [a]*3), a)
1086 b = [a, None]
1087 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1088
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001089 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001090 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001091 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001092
1093 def test_islice(self):
1094 a = []
1095 self.makecycle(islice([a]*2, None), a)
1096
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001097 def test_permutations(self):
1098 a = []
1099 self.makecycle(permutations([1,2,a,3], 3), a)
1100
1101 def test_product(self):
1102 a = []
1103 self.makecycle(product([1,2,a,3], repeat=3), a)
1104
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001105 def test_repeat(self):
1106 a = []
1107 self.makecycle(repeat(a), a)
1108
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001109 def test_starmap(self):
1110 a = []
1111 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1112
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001113 def test_takewhile(self):
1114 a = []
1115 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1116
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001117def R(seqn):
1118 'Regular generator'
1119 for i in seqn:
1120 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001121
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001122class G:
1123 'Sequence using __getitem__'
1124 def __init__(self, seqn):
1125 self.seqn = seqn
1126 def __getitem__(self, i):
1127 return self.seqn[i]
1128
1129class I:
1130 'Sequence using iterator protocol'
1131 def __init__(self, seqn):
1132 self.seqn = seqn
1133 self.i = 0
1134 def __iter__(self):
1135 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001136 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001137 if self.i >= len(self.seqn): raise StopIteration
1138 v = self.seqn[self.i]
1139 self.i += 1
1140 return v
1141
1142class Ig:
1143 'Sequence using iterator protocol defined with a generator'
1144 def __init__(self, seqn):
1145 self.seqn = seqn
1146 self.i = 0
1147 def __iter__(self):
1148 for val in self.seqn:
1149 yield val
1150
1151class X:
1152 'Missing __getitem__ and __iter__'
1153 def __init__(self, seqn):
1154 self.seqn = seqn
1155 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001156 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001157 if self.i >= len(self.seqn): raise StopIteration
1158 v = self.seqn[self.i]
1159 self.i += 1
1160 return v
1161
1162class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001163 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001164 def __init__(self, seqn):
1165 self.seqn = seqn
1166 self.i = 0
1167 def __iter__(self):
1168 return self
1169
1170class E:
1171 'Test propagation of exceptions'
1172 def __init__(self, seqn):
1173 self.seqn = seqn
1174 self.i = 0
1175 def __iter__(self):
1176 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001177 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001178 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001179
1180class S:
1181 'Test immediate stop'
1182 def __init__(self, seqn):
1183 pass
1184 def __iter__(self):
1185 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001186 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001187 raise StopIteration
1188
1189def L(seqn):
1190 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001191 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001192
1193
1194class TestVariousIteratorArgs(unittest.TestCase):
1195
1196 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001197 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001198 for g in (G, I, Ig, S, L, R):
1199 self.assertEqual(list(chain(g(s))), list(g(s)))
1200 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001201 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001202 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001203 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1204
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001205 def test_compress(self):
1206 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1207 n = len(s)
1208 for g in (G, I, Ig, S, L, R):
1209 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1210 self.assertRaises(TypeError, compress, X(s), repeat(1))
1211 self.assertRaises(TypeError, compress, N(s), repeat(1))
1212 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1213
Christian Heimesc3f30c42008-02-22 16:37:40 +00001214 def test_product(self):
1215 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1216 self.assertRaises(TypeError, product, X(s))
1217 self.assertRaises(TypeError, product, N(s))
1218 self.assertRaises(ZeroDivisionError, product, E(s))
1219
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001220 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001221 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001222 for g in (G, I, Ig, S, L, R):
1223 tgtlen = len(s) * 3
1224 expected = list(g(s))*3
1225 actual = list(islice(cycle(g(s)), tgtlen))
1226 self.assertEqual(actual, expected)
1227 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001228 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001229 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1230
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001231 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001232 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001233 for g in (G, I, Ig, S, L, R):
1234 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1235 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001236 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001237 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1238
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001239 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001240 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001241 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001242 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001243 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001244 self.assertRaises(TypeError, filter, isEven, X(s))
1245 self.assertRaises(TypeError, filter, isEven, N(s))
1246 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001247
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001248 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001249 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001250 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001251 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001252 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001253 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1254 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1255 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001256
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001257 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001258 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001259 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001260 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1261 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1262 self.assertRaises(TypeError, zip, X(s))
1263 self.assertRaises(TypeError, zip, N(s))
1264 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001265
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001266 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001267 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001268 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001269 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1270 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1271 self.assertRaises(TypeError, zip_longest, X(s))
1272 self.assertRaises(TypeError, zip_longest, N(s))
1273 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001274
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001275 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001276 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001277 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001278 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001279 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001280 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001281 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001282 self.assertRaises(TypeError, map, onearg, X(s))
1283 self.assertRaises(TypeError, map, onearg, N(s))
1284 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001285
1286 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001287 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001288 for g in (G, I, Ig, S, L, R):
1289 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1290 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001291 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001292 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1293
1294 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001295 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001296 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001297 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001298 self.assertEqual(list(starmap(operator.pow, g(ss))),
1299 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001300 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001301 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001302 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1303
1304 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001305 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001306 for g in (G, I, Ig, S, L, R):
1307 tgt = []
1308 for elem in g(s):
1309 if not isEven(elem): break
1310 tgt.append(elem)
1311 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1312 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001313 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001314 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1315
1316 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001317 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001318 for g in (G, I, Ig, S, L, R):
1319 tgt = []
1320 for elem in g(s):
1321 if not tgt and isOdd(elem): continue
1322 tgt.append(elem)
1323 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1324 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001325 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001326 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1327
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001328 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001329 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001330 for g in (G, I, Ig, S, L, R):
1331 it1, it2 = tee(g(s))
1332 self.assertEqual(list(it1), list(g(s)))
1333 self.assertEqual(list(it2), list(g(s)))
1334 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001335 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001336 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1337
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001338class LengthTransparency(unittest.TestCase):
1339
1340 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001341 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001342 self.assertEqual(len(repeat(None, 50)), 50)
1343 self.assertRaises(TypeError, len, repeat(None))
1344
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001345class RegressionTests(unittest.TestCase):
1346
1347 def test_sf_793826(self):
1348 # Fix Armin Rigo's successful efforts to wreak havoc
1349
1350 def mutatingtuple(tuple1, f, tuple2):
1351 # this builds a tuple t which is a copy of tuple1,
1352 # then calls f(t), then mutates t to be equal to tuple2
1353 # (needs len(tuple1) == len(tuple2)).
1354 def g(value, first=[1]):
1355 if first:
1356 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00001357 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001358 return value
1359 items = list(tuple2)
1360 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001361 gen = map(g, items)
1362 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00001363 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001364
1365 def f(t):
1366 global T
1367 T = t
1368 first[:] = list(T)
1369
1370 first = []
1371 mutatingtuple((1,2,3), f, (4,5,6))
1372 second = list(T)
1373 self.assertEqual(first, second)
1374
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001375
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001376 def test_sf_950057(self):
1377 # Make sure that chain() and cycle() catch exceptions immediately
1378 # rather than when shifting between input sources
1379
1380 def gen1():
1381 hist.append(0)
1382 yield 1
1383 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001384 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001385 hist.append(2)
1386
1387 def gen2(x):
1388 hist.append(3)
1389 yield 2
1390 hist.append(4)
1391 if x:
1392 raise StopIteration
1393
1394 hist = []
1395 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1396 self.assertEqual(hist, [0,1])
1397
1398 hist = []
1399 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1400 self.assertEqual(hist, [0,1])
1401
1402 hist = []
1403 self.assertRaises(AssertionError, list, cycle(gen1()))
1404 self.assertEqual(hist, [0,1])
1405
Thomas Woutersb2137042007-02-01 18:02:27 +00001406class SubclassWithKwargsTest(unittest.TestCase):
1407 def test_keywords_in_subclass(self):
1408 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001409 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001410 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00001411 class Subclass(cls):
1412 def __init__(self, newarg=None, *args):
1413 cls.__init__(self, *args)
1414 try:
1415 Subclass(newarg=1)
1416 except TypeError as err:
1417 # we expect type errors because of wrong argument count
Georg Brandlab91fde2009-08-13 08:51:18 +00001418 self.assertFalse("does not take keyword arguments" in err.args[0])
Thomas Woutersb2137042007-02-01 18:02:27 +00001419
1420
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001421libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001422
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001423
1424>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001425>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00001426... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001427...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001428Check 1200 is for $120.15
1429Check 1201 is for $764.05
1430Check 1202 is for $823.14
1431
1432>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001433>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00001434... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00001435...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000014361
14378
143827
1439
1440>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001441>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00001442... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00001443...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001444Alex
1445Laura
1446Martin
1447Walter
1448Samuele
1449
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001450>>> from operator import itemgetter
1451>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00001452>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001453>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001454... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001455...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000014561 ['a', 'c', 'e']
14572 ['b', 'd', 'f']
14583 ['g']
1459
Raymond Hettinger734fb572004-01-20 20:04:40 +00001460# Find runs of consecutive numbers using groupby. The key to the solution
1461# is differencing with a range so that consecutive numbers all appear in
1462# same group.
1463>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00001464>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001465... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001466...
Raymond Hettinger734fb572004-01-20 20:04:40 +00001467[1]
1468[4, 5, 6]
1469[10]
1470[15, 16, 17, 18]
1471[22]
1472[25, 26, 27, 28]
1473
Georg Brandl3dbca812008-07-23 16:10:53 +00001474>>> def take(n, iterable):
1475... "Return first n items of the iterable as a list"
1476... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001477
Georg Brandl3dbca812008-07-23 16:10:53 +00001478>>> def enumerate(iterable, start=0):
1479... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001480
Georg Brandl3dbca812008-07-23 16:10:53 +00001481>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001482... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00001483... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001484
Raymond Hettingercdf8ba32009-02-19 04:45:07 +00001485>>> def nth(iterable, n, default=None):
1486... "Returns the nth item or a default value"
1487... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001488
Georg Brandl3dbca812008-07-23 16:10:53 +00001489>>> def quantify(iterable, pred=bool):
1490... "Count how many times the predicate is true"
1491... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001492
Georg Brandl3dbca812008-07-23 16:10:53 +00001493>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001494... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00001495... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001496
Georg Brandl3dbca812008-07-23 16:10:53 +00001497>>> def ncycles(iterable, n):
Ezio Melotti13925002011-03-16 11:05:33 +02001498... "Returns the sequence elements n times"
Georg Brandl3dbca812008-07-23 16:10:53 +00001499... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001500
1501>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001502... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001503
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001504>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001505... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001506
1507>>> def repeatfunc(func, times=None, *args):
1508... "Repeat calls to func with specified arguments."
1509... " Example: repeatfunc(random.random)"
1510... if times is None:
1511... return starmap(func, repeat(args))
1512... else:
1513... return starmap(func, repeat(args, times))
1514
Raymond Hettingerd591f662003-10-26 15:34:50 +00001515>>> def pairwise(iterable):
1516... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1517... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001518... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00001519... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001520... except StopIteration:
1521... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001522... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001523
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001524>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001525... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001526... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00001527... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001528
1529>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001530... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001531... # Recipe credited to George Sakkis
1532... pending = len(iterables)
1533... nexts = cycle(iter(it).__next__ for it in iterables)
1534... while pending:
1535... try:
1536... for next in nexts:
1537... yield next()
1538... except StopIteration:
1539... pending -= 1
1540... nexts = cycle(islice(nexts, pending))
1541
1542>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00001543... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1544... s = list(iterable)
1545... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001546
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001547>>> def unique_everseen(iterable, key=None):
1548... "List unique elements, preserving order. Remember all elements ever seen."
1549... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1550... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1551... seen = set()
1552... seen_add = seen.add
1553... if key is None:
1554... for element in iterable:
1555... if element not in seen:
1556... seen_add(element)
1557... yield element
1558... else:
1559... for element in iterable:
1560... k = key(element)
1561... if k not in seen:
1562... seen_add(k)
1563... yield element
1564
1565>>> def unique_justseen(iterable, key=None):
1566... "List unique elements, preserving order. Remember only the element just seen."
1567... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1568... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1569... return map(next, map(itemgetter(1), groupby(iterable, key)))
1570
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001571This is not part of the examples but it tests to make sure the definitions
1572perform as purported.
1573
Raymond Hettingera098b332003-09-08 23:58:40 +00001574>>> take(10, count())
1575[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1576
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001577>>> list(enumerate('abc'))
1578[(0, 'a'), (1, 'b'), (2, 'c')]
1579
1580>>> list(islice(tabulate(lambda x: 2*x), 4))
1581[0, 2, 4, 6]
1582
1583>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00001584'd'
1585
1586>>> nth('abcde', 9) is None
1587True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001588
Guido van Rossum805365e2007-05-07 22:24:25 +00001589>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000159050
1591
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001592>>> a = [[1, 2, 3], [4, 5, 6]]
1593>>> flatten(a)
1594[1, 2, 3, 4, 5, 6]
1595
1596>>> list(repeatfunc(pow, 5, 2, 3))
1597[8, 8, 8, 8, 8]
1598
1599>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001600>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001601[0, 0, 0, 0, 0]
1602
Raymond Hettingerd591f662003-10-26 15:34:50 +00001603>>> list(pairwise('abcd'))
1604[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001605
Raymond Hettingerd591f662003-10-26 15:34:50 +00001606>>> list(pairwise([]))
1607[]
1608
1609>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001610[]
1611
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001612>>> list(islice(padnone('abc'), 0, 6))
1613['a', 'b', 'c', None, None, None]
1614
1615>>> list(ncycles('abc', 3))
1616['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1617
1618>>> dotproduct([1,2,3], [4,5,6])
161932
1620
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001621>>> list(grouper(3, 'abcdefg', 'x'))
1622[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1623
1624>>> list(roundrobin('abc', 'd', 'ef'))
1625['a', 'd', 'e', 'b', 'f', 'c']
1626
Raymond Hettingerace67332009-01-26 02:23:50 +00001627>>> list(powerset([1,2,3]))
1628[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001629
Raymond Hettinger191e8502009-01-27 13:29:43 +00001630>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1631True
1632
1633>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1634True
1635
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001636>>> list(unique_everseen('AAAABBBCCDAABBB'))
1637['A', 'B', 'C', 'D']
1638
1639>>> list(unique_everseen('ABBCcAD', str.lower))
1640['A', 'B', 'C', 'D']
1641
1642>>> list(unique_justseen('AAAABBBCCDAABBB'))
1643['A', 'B', 'C', 'D', 'A', 'B']
1644
1645>>> list(unique_justseen('ABBCcAD', str.lower))
1646['A', 'B', 'C', 'A', 'D']
1647
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001648"""
1649
1650__test__ = {'libreftest' : libreftest}
1651
1652def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001653 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00001654 RegressionTests, LengthTransparency,
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001655 SubclassWithKwargsTest, TestExamples)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001656 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001657
1658 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001659 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001660 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001661 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00001662 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001663 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001664 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001665 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001666 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001667
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001668 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001669 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001670
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001671if __name__ == "__main__":
1672 test_main(verbose=True)