blob: 16e65cab65eb337da079b15766e4f7a552c57cb1 [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:
198 self.assertEquals(result, regular_combs) # cases that should be identical
199 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
294 self.assertEquals(len(prod), n**r)
295 self.assertEquals(len(cwr), (fact(n+r-1) / fact(r)/ fact(n-1)) if n else (not r))
296 self.assertEquals(len(perm), 0 if r>n else fact(n) / fact(n-r))
297 self.assertEquals(len(comb), 0 if r>n else fact(n) / fact(r) / fact(n-r))
298
299 # Check lexicographic order without repeated tuples
300 self.assertEquals(prod, sorted(set(prod)))
301 self.assertEquals(cwr, sorted(set(cwr)))
302 self.assertEquals(perm, sorted(set(perm)))
303 self.assertEquals(comb, sorted(set(comb)))
304
305 # Check interrelationships
306 self.assertEquals(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
307 self.assertEquals(perm, [t for t in prod if len(set(t))==r]) # perm: prods with no dups
308 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
791 def test_takewhile(self):
792 data = [1, 3, 5, 20, 2, 4, 6, 8]
793 underten = lambda x: x<10
794 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000795 self.assertEqual(list(takewhile(underten, [])), [])
796 self.assertRaises(TypeError, takewhile)
797 self.assertRaises(TypeError, takewhile, operator.pow)
798 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000799 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
800 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000801 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
802 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +0000803 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000804
805 def test_dropwhile(self):
806 data = [1, 3, 5, 20, 2, 4, 6, 8]
807 underten = lambda x: x<10
808 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000809 self.assertEqual(list(dropwhile(underten, [])), [])
810 self.assertRaises(TypeError, dropwhile)
811 self.assertRaises(TypeError, dropwhile, operator.pow)
812 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000813 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
814 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000815
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000816 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000817 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000818 def irange(n):
Guido van Rossum805365e2007-05-07 22:24:25 +0000819 for i in range(n):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000820 yield i
821
822 a, b = tee([]) # test empty iterator
823 self.assertEqual(list(a), [])
824 self.assertEqual(list(b), [])
825
826 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +0000827 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000828
829 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +0000830 self.assertEqual(list(a), list(range(n)))
831 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000832
833 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000834 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000835 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000836 del a
Guido van Rossum805365e2007-05-07 22:24:25 +0000837 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000838
839 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000840 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000841 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000842 del b
Guido van Rossum805365e2007-05-07 22:24:25 +0000843 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000844
Guido van Rossum805365e2007-05-07 22:24:25 +0000845 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000846 order = [0]*n + [1]*n
847 random.shuffle(order)
848 lists = ([], [])
849 its = tee(irange(n))
850 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +0000851 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000852 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +0000853 self.assertEqual(lists[0], list(range(n)))
854 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000855
Raymond Hettingerad983e72003-11-12 14:32:26 +0000856 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000857 self.assertRaises(TypeError, tee)
858 self.assertRaises(TypeError, tee, 3)
859 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000860 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000861
Raymond Hettingerad983e72003-11-12 14:32:26 +0000862 # tee object should be instantiable
863 a, b = tee('abc')
864 c = type(a)('def')
865 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000866
Raymond Hettingerad983e72003-11-12 14:32:26 +0000867 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +0000868 a, b, c = tee(range(2000), 3)
869 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000870 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +0000871 self.assertEqual(list(b), list(range(2000)))
872 self.assertEqual([next(c), next(c)], list(range(2)))
873 self.assertEqual(list(a), list(range(100,2000)))
874 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000875
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000876 # test values of n
877 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000878 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +0000879 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000880 result = tee('abc', n)
881 self.assertEqual(type(result), tuple)
882 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000883 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000884
Raymond Hettingerad983e72003-11-12 14:32:26 +0000885 # tee pass-through to copyable iterator
886 a, b = tee('abc')
887 c, d = tee(a)
Georg Brandlab91fde2009-08-13 08:51:18 +0000888 self.assertTrue(a is c)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000889
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000890 # test tee_new
891 t1, t2 = tee('abc')
892 tnew = type(t1)
893 self.assertRaises(TypeError, tnew)
894 self.assertRaises(TypeError, tnew, 10)
895 t3 = tnew(t1)
Georg Brandlab91fde2009-08-13 08:51:18 +0000896 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000897
Raymond Hettingera9f60922004-10-17 16:40:14 +0000898 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +0000899 a, b = tee(range(10))
Raymond Hettingera9f60922004-10-17 16:40:14 +0000900 p = proxy(a)
901 self.assertEqual(getattr(p, '__class__'), type(b))
902 del a
903 self.assertRaises(ReferenceError, getattr, p, '__class__')
904
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000905 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000906 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000907
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000908 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +0000909 self.assertRaises(StopIteration, next, f([]))
910 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000911
Georg Brandla18af4e2007-04-21 15:47:16 +0000912 self.assertRaises(StopIteration, next, islice([], None))
913 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000914
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000915 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +0000916 self.assertRaises(StopIteration, next, p)
917 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000918 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +0000919 self.assertRaises(StopIteration, next, p)
920 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000921
Georg Brandla18af4e2007-04-21 15:47:16 +0000922 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000923
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000924 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +0000925 self.assertRaises(StopIteration, next, f(lambda x:x, []))
926 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000927
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000928class TestExamples(unittest.TestCase):
929
930 def test_chain(self):
931 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
932
933 def test_chain_from_iterable(self):
934 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
935
936 def test_combinations(self):
937 self.assertEqual(list(combinations('ABCD', 2)),
938 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
939 self.assertEqual(list(combinations(range(4), 3)),
940 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
941
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000942 def test_combinations_with_replacement(self):
943 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
944 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
945
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000946 def test_compress(self):
947 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
948
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000949 def test_count(self):
950 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
951
952 def test_cycle(self):
953 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
954
955 def test_dropwhile(self):
956 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
957
958 def test_groupby(self):
959 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
960 list('ABCDAB'))
961 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
962 [list('AAAA'), list('BBB'), list('CC'), list('D')])
963
964 def test_filter(self):
965 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
966
967 def test_filterfalse(self):
968 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
969
970 def test_map(self):
971 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
972
973 def test_islice(self):
974 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
975 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
976 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
977 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
978
979 def test_zip(self):
980 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
981
982 def test_zip_longest(self):
983 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
984 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
985
986 def test_permutations(self):
987 self.assertEqual(list(permutations('ABCD', 2)),
988 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
989 self.assertEqual(list(permutations(range(3))),
990 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
991
992 def test_product(self):
993 self.assertEqual(list(product('ABCD', 'xy')),
994 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
995 self.assertEqual(list(product(range(2), repeat=3)),
996 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
997 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
998
999 def test_repeat(self):
1000 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
1001
1002 def test_stapmap(self):
1003 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
1004 [32, 9, 1000])
1005
1006 def test_takewhile(self):
1007 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
1008
1009
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001010class TestGC(unittest.TestCase):
1011
1012 def makecycle(self, iterator, container):
1013 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +00001014 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001015 del container, iterator
1016
1017 def test_chain(self):
1018 a = []
1019 self.makecycle(chain(a), a)
1020
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001021 def test_chain_from_iterable(self):
1022 a = []
1023 self.makecycle(chain.from_iterable([a]), a)
1024
1025 def test_combinations(self):
1026 a = []
1027 self.makecycle(combinations([1,2,a,3], 3), a)
1028
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001029 def test_combinations_with_replacement(self):
1030 a = []
1031 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1032
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001033 def test_compress(self):
1034 a = []
1035 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1036
Raymond Hettingerd6280f42009-02-16 20:50:56 +00001037 def test_count(self):
1038 a = []
1039 Int = type('Int', (int,), dict(x=a))
1040 self.makecycle(count(Int(0), Int(1)), a)
1041
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001042 def test_cycle(self):
1043 a = []
1044 self.makecycle(cycle([a]*2), a)
1045
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001046 def test_dropwhile(self):
1047 a = []
1048 self.makecycle(dropwhile(bool, [0, a, a]), a)
1049
1050 def test_groupby(self):
1051 a = []
1052 self.makecycle(groupby([a]*2, lambda x:x), a)
1053
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001054 def test_issue2246(self):
1055 # Issue 2246 -- the _grouper iterator was not included in GC
1056 n = 10
1057 keyfunc = lambda x: x
1058 for i, j in groupby(range(n), key=keyfunc):
1059 keyfunc.__dict__.setdefault('x',[]).append(j)
1060
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001061 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001062 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001063 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001064
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001065 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001066 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001067 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001068
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001069 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001070 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001071 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001072
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001073 def test_zip_longest(self):
1074 a = []
1075 self.makecycle(zip_longest([a]*2, [a]*3), a)
1076 b = [a, None]
1077 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1078
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001079 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001080 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001081 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001082
1083 def test_islice(self):
1084 a = []
1085 self.makecycle(islice([a]*2, None), a)
1086
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001087 def test_permutations(self):
1088 a = []
1089 self.makecycle(permutations([1,2,a,3], 3), a)
1090
1091 def test_product(self):
1092 a = []
1093 self.makecycle(product([1,2,a,3], repeat=3), a)
1094
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001095 def test_repeat(self):
1096 a = []
1097 self.makecycle(repeat(a), a)
1098
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001099 def test_starmap(self):
1100 a = []
1101 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1102
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001103 def test_takewhile(self):
1104 a = []
1105 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1106
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001107def R(seqn):
1108 'Regular generator'
1109 for i in seqn:
1110 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001111
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001112class G:
1113 'Sequence using __getitem__'
1114 def __init__(self, seqn):
1115 self.seqn = seqn
1116 def __getitem__(self, i):
1117 return self.seqn[i]
1118
1119class I:
1120 'Sequence using iterator protocol'
1121 def __init__(self, seqn):
1122 self.seqn = seqn
1123 self.i = 0
1124 def __iter__(self):
1125 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001126 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001127 if self.i >= len(self.seqn): raise StopIteration
1128 v = self.seqn[self.i]
1129 self.i += 1
1130 return v
1131
1132class Ig:
1133 'Sequence using iterator protocol defined with a generator'
1134 def __init__(self, seqn):
1135 self.seqn = seqn
1136 self.i = 0
1137 def __iter__(self):
1138 for val in self.seqn:
1139 yield val
1140
1141class X:
1142 'Missing __getitem__ and __iter__'
1143 def __init__(self, seqn):
1144 self.seqn = seqn
1145 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001146 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001147 if self.i >= len(self.seqn): raise StopIteration
1148 v = self.seqn[self.i]
1149 self.i += 1
1150 return v
1151
1152class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001153 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001154 def __init__(self, seqn):
1155 self.seqn = seqn
1156 self.i = 0
1157 def __iter__(self):
1158 return self
1159
1160class E:
1161 'Test propagation of exceptions'
1162 def __init__(self, seqn):
1163 self.seqn = seqn
1164 self.i = 0
1165 def __iter__(self):
1166 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001167 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001168 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001169
1170class S:
1171 'Test immediate stop'
1172 def __init__(self, seqn):
1173 pass
1174 def __iter__(self):
1175 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001176 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001177 raise StopIteration
1178
1179def L(seqn):
1180 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001181 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001182
1183
1184class TestVariousIteratorArgs(unittest.TestCase):
1185
1186 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001187 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001188 for g in (G, I, Ig, S, L, R):
1189 self.assertEqual(list(chain(g(s))), list(g(s)))
1190 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001191 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001192 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001193 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1194
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001195 def test_compress(self):
1196 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1197 n = len(s)
1198 for g in (G, I, Ig, S, L, R):
1199 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1200 self.assertRaises(TypeError, compress, X(s), repeat(1))
1201 self.assertRaises(TypeError, compress, N(s), repeat(1))
1202 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1203
Christian Heimesc3f30c42008-02-22 16:37:40 +00001204 def test_product(self):
1205 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1206 self.assertRaises(TypeError, product, X(s))
1207 self.assertRaises(TypeError, product, N(s))
1208 self.assertRaises(ZeroDivisionError, product, E(s))
1209
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001210 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001211 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001212 for g in (G, I, Ig, S, L, R):
1213 tgtlen = len(s) * 3
1214 expected = list(g(s))*3
1215 actual = list(islice(cycle(g(s)), tgtlen))
1216 self.assertEqual(actual, expected)
1217 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001218 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001219 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1220
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001221 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001222 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001223 for g in (G, I, Ig, S, L, R):
1224 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1225 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001226 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001227 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1228
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001229 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001230 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001231 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001232 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001233 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001234 self.assertRaises(TypeError, filter, isEven, X(s))
1235 self.assertRaises(TypeError, filter, isEven, N(s))
1236 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001237
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001238 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001239 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001240 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001241 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001242 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001243 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1244 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1245 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001246
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001247 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001248 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001249 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001250 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1251 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1252 self.assertRaises(TypeError, zip, X(s))
1253 self.assertRaises(TypeError, zip, N(s))
1254 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001255
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001256 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001257 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001258 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001259 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1260 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1261 self.assertRaises(TypeError, zip_longest, X(s))
1262 self.assertRaises(TypeError, zip_longest, N(s))
1263 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001264
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001265 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001266 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001267 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001268 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001269 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001270 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001271 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001272 self.assertRaises(TypeError, map, onearg, X(s))
1273 self.assertRaises(TypeError, map, onearg, N(s))
1274 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001275
1276 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001277 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001278 for g in (G, I, Ig, S, L, R):
1279 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1280 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001281 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001282 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1283
1284 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001285 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001286 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001287 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001288 self.assertEqual(list(starmap(operator.pow, g(ss))),
1289 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001290 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001291 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001292 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1293
1294 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001295 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001296 for g in (G, I, Ig, S, L, R):
1297 tgt = []
1298 for elem in g(s):
1299 if not isEven(elem): break
1300 tgt.append(elem)
1301 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1302 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001303 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001304 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1305
1306 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001307 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001308 for g in (G, I, Ig, S, L, R):
1309 tgt = []
1310 for elem in g(s):
1311 if not tgt and isOdd(elem): continue
1312 tgt.append(elem)
1313 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1314 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001315 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001316 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1317
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001318 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001319 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001320 for g in (G, I, Ig, S, L, R):
1321 it1, it2 = tee(g(s))
1322 self.assertEqual(list(it1), list(g(s)))
1323 self.assertEqual(list(it2), list(g(s)))
1324 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001325 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001326 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1327
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001328class LengthTransparency(unittest.TestCase):
1329
1330 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001331 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001332 self.assertEqual(len(repeat(None, 50)), 50)
1333 self.assertRaises(TypeError, len, repeat(None))
1334
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001335class RegressionTests(unittest.TestCase):
1336
1337 def test_sf_793826(self):
1338 # Fix Armin Rigo's successful efforts to wreak havoc
1339
1340 def mutatingtuple(tuple1, f, tuple2):
1341 # this builds a tuple t which is a copy of tuple1,
1342 # then calls f(t), then mutates t to be equal to tuple2
1343 # (needs len(tuple1) == len(tuple2)).
1344 def g(value, first=[1]):
1345 if first:
1346 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00001347 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001348 return value
1349 items = list(tuple2)
1350 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001351 gen = map(g, items)
1352 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00001353 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001354
1355 def f(t):
1356 global T
1357 T = t
1358 first[:] = list(T)
1359
1360 first = []
1361 mutatingtuple((1,2,3), f, (4,5,6))
1362 second = list(T)
1363 self.assertEqual(first, second)
1364
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001365
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001366 def test_sf_950057(self):
1367 # Make sure that chain() and cycle() catch exceptions immediately
1368 # rather than when shifting between input sources
1369
1370 def gen1():
1371 hist.append(0)
1372 yield 1
1373 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001374 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001375 hist.append(2)
1376
1377 def gen2(x):
1378 hist.append(3)
1379 yield 2
1380 hist.append(4)
1381 if x:
1382 raise StopIteration
1383
1384 hist = []
1385 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1386 self.assertEqual(hist, [0,1])
1387
1388 hist = []
1389 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1390 self.assertEqual(hist, [0,1])
1391
1392 hist = []
1393 self.assertRaises(AssertionError, list, cycle(gen1()))
1394 self.assertEqual(hist, [0,1])
1395
Thomas Woutersb2137042007-02-01 18:02:27 +00001396class SubclassWithKwargsTest(unittest.TestCase):
1397 def test_keywords_in_subclass(self):
1398 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001399 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001400 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00001401 class Subclass(cls):
1402 def __init__(self, newarg=None, *args):
1403 cls.__init__(self, *args)
1404 try:
1405 Subclass(newarg=1)
1406 except TypeError as err:
1407 # we expect type errors because of wrong argument count
Georg Brandlab91fde2009-08-13 08:51:18 +00001408 self.assertFalse("does not take keyword arguments" in err.args[0])
Thomas Woutersb2137042007-02-01 18:02:27 +00001409
1410
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001411libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001412
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001413
1414>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001415>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00001416... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001417...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001418Check 1200 is for $120.15
1419Check 1201 is for $764.05
1420Check 1202 is for $823.14
1421
1422>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001423>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00001424... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00001425...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000014261
14278
142827
1429
1430>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001431>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00001432... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00001433...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001434Alex
1435Laura
1436Martin
1437Walter
1438Samuele
1439
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001440>>> from operator import itemgetter
1441>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00001442>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001443>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001444... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001445...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000014461 ['a', 'c', 'e']
14472 ['b', 'd', 'f']
14483 ['g']
1449
Raymond Hettinger734fb572004-01-20 20:04:40 +00001450# Find runs of consecutive numbers using groupby. The key to the solution
1451# is differencing with a range so that consecutive numbers all appear in
1452# same group.
1453>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00001454>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001455... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001456...
Raymond Hettinger734fb572004-01-20 20:04:40 +00001457[1]
1458[4, 5, 6]
1459[10]
1460[15, 16, 17, 18]
1461[22]
1462[25, 26, 27, 28]
1463
Georg Brandl3dbca812008-07-23 16:10:53 +00001464>>> def take(n, iterable):
1465... "Return first n items of the iterable as a list"
1466... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001467
Georg Brandl3dbca812008-07-23 16:10:53 +00001468>>> def enumerate(iterable, start=0):
1469... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001470
Georg Brandl3dbca812008-07-23 16:10:53 +00001471>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001472... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00001473... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001474
Raymond Hettingercdf8ba32009-02-19 04:45:07 +00001475>>> def nth(iterable, n, default=None):
1476... "Returns the nth item or a default value"
1477... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001478
Georg Brandl3dbca812008-07-23 16:10:53 +00001479>>> def quantify(iterable, pred=bool):
1480... "Count how many times the predicate is true"
1481... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001482
Georg Brandl3dbca812008-07-23 16:10:53 +00001483>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001484... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00001485... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001486
Georg Brandl3dbca812008-07-23 16:10:53 +00001487>>> def ncycles(iterable, n):
1488... "Returns the seqeuence elements n times"
1489... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001490
1491>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001492... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001493
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001494>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001495... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001496
1497>>> def repeatfunc(func, times=None, *args):
1498... "Repeat calls to func with specified arguments."
1499... " Example: repeatfunc(random.random)"
1500... if times is None:
1501... return starmap(func, repeat(args))
1502... else:
1503... return starmap(func, repeat(args, times))
1504
Raymond Hettingerd591f662003-10-26 15:34:50 +00001505>>> def pairwise(iterable):
1506... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1507... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001508... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00001509... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001510... except StopIteration:
1511... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001512... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001513
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001514>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001515... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001516... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00001517... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001518
1519>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001520... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001521... # Recipe credited to George Sakkis
1522... pending = len(iterables)
1523... nexts = cycle(iter(it).__next__ for it in iterables)
1524... while pending:
1525... try:
1526... for next in nexts:
1527... yield next()
1528... except StopIteration:
1529... pending -= 1
1530... nexts = cycle(islice(nexts, pending))
1531
1532>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00001533... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1534... s = list(iterable)
1535... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001536
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001537>>> def unique_everseen(iterable, key=None):
1538... "List unique elements, preserving order. Remember all elements ever seen."
1539... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1540... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1541... seen = set()
1542... seen_add = seen.add
1543... if key is None:
1544... for element in iterable:
1545... if element not in seen:
1546... seen_add(element)
1547... yield element
1548... else:
1549... for element in iterable:
1550... k = key(element)
1551... if k not in seen:
1552... seen_add(k)
1553... yield element
1554
1555>>> def unique_justseen(iterable, key=None):
1556... "List unique elements, preserving order. Remember only the element just seen."
1557... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1558... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1559... return map(next, map(itemgetter(1), groupby(iterable, key)))
1560
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001561This is not part of the examples but it tests to make sure the definitions
1562perform as purported.
1563
Raymond Hettingera098b332003-09-08 23:58:40 +00001564>>> take(10, count())
1565[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1566
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001567>>> list(enumerate('abc'))
1568[(0, 'a'), (1, 'b'), (2, 'c')]
1569
1570>>> list(islice(tabulate(lambda x: 2*x), 4))
1571[0, 2, 4, 6]
1572
1573>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00001574'd'
1575
1576>>> nth('abcde', 9) is None
1577True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001578
Guido van Rossum805365e2007-05-07 22:24:25 +00001579>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000158050
1581
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001582>>> a = [[1, 2, 3], [4, 5, 6]]
1583>>> flatten(a)
1584[1, 2, 3, 4, 5, 6]
1585
1586>>> list(repeatfunc(pow, 5, 2, 3))
1587[8, 8, 8, 8, 8]
1588
1589>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001590>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001591[0, 0, 0, 0, 0]
1592
Raymond Hettingerd591f662003-10-26 15:34:50 +00001593>>> list(pairwise('abcd'))
1594[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001595
Raymond Hettingerd591f662003-10-26 15:34:50 +00001596>>> list(pairwise([]))
1597[]
1598
1599>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001600[]
1601
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001602>>> list(islice(padnone('abc'), 0, 6))
1603['a', 'b', 'c', None, None, None]
1604
1605>>> list(ncycles('abc', 3))
1606['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1607
1608>>> dotproduct([1,2,3], [4,5,6])
160932
1610
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001611>>> list(grouper(3, 'abcdefg', 'x'))
1612[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1613
1614>>> list(roundrobin('abc', 'd', 'ef'))
1615['a', 'd', 'e', 'b', 'f', 'c']
1616
Raymond Hettingerace67332009-01-26 02:23:50 +00001617>>> list(powerset([1,2,3]))
1618[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001619
Raymond Hettinger191e8502009-01-27 13:29:43 +00001620>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1621True
1622
1623>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1624True
1625
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001626>>> list(unique_everseen('AAAABBBCCDAABBB'))
1627['A', 'B', 'C', 'D']
1628
1629>>> list(unique_everseen('ABBCcAD', str.lower))
1630['A', 'B', 'C', 'D']
1631
1632>>> list(unique_justseen('AAAABBBCCDAABBB'))
1633['A', 'B', 'C', 'D', 'A', 'B']
1634
1635>>> list(unique_justseen('ABBCcAD', str.lower))
1636['A', 'B', 'C', 'A', 'D']
1637
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001638"""
1639
1640__test__ = {'libreftest' : libreftest}
1641
1642def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001643 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00001644 RegressionTests, LengthTransparency,
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001645 SubclassWithKwargsTest, TestExamples)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001646 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001647
1648 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001649 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001650 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001651 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00001652 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001653 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001654 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001655 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001656 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001657
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001658 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001659 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001660
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001661if __name__ == "__main__":
1662 test_main(verbose=True)