blob: 2f449da0bbb6ed7b3e1f8d419d63e53c889e2e3b [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001import unittest
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002from test import support
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003from itertools import *
Raymond Hettingera9f60922004-10-17 16:40:14 +00004from weakref import proxy
Raymond Hettinger8a882a72009-02-12 12:08:18 +00005from decimal import Decimal
Raymond Hettingerde09acf2009-02-12 12:53:53 +00006from fractions import Fraction
Raymond Hettinger2012f172003-02-07 05:32:58 +00007import sys
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00008import operator
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00009import random
Christian Heimesc3f30c42008-02-22 16:37:40 +000010from functools import reduce
Benjamin Petersonee8712c2008-05-20 21:35:26 +000011maxsize = support.MAX_Py_ssize_t
Guido van Rossum360e4b82007-05-14 22:51:27 +000012minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000013
Guido van Rossum801f0d72006-08-24 19:48:10 +000014def lzip(*args):
15 return list(zip(*args))
16
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000017def onearg(x):
18 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000019 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000020
21def errfunc(*args):
22 'Test function that raises an error'
23 raise ValueError
24
25def gen3():
26 'Non-restartable source sequence'
27 for i in (0, 1, 2):
28 yield i
29
30def isEven(x):
31 'Test predicate'
32 return x%2==0
33
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000034def isOdd(x):
35 'Test predicate'
36 return x%2==1
37
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000038class StopNow:
39 'Class emulating an empty iterable.'
40 def __iter__(self):
41 return self
Georg Brandla18af4e2007-04-21 15:47:16 +000042 def __next__(self):
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000043 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000044
Raymond Hettinger02420702003-06-29 20:36:23 +000045def take(n, seq):
46 'Convenience function for partially consuming a long of infinite iterable'
47 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000048
Christian Heimes78644762008-03-04 23:39:23 +000049def prod(iterable):
50 return reduce(operator.mul, iterable, 1)
51
Christian Heimes380f7f22008-02-28 11:19:05 +000052def fact(n):
53 'Factorial'
Christian Heimes78644762008-03-04 23:39:23 +000054 return prod(range(1, n+1))
55
Raymond Hettinger96ef8112003-02-01 00:10:11 +000056class TestBasicOps(unittest.TestCase):
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000057 def test_chain(self):
Christian Heimesdd15f6c2008-03-16 00:07:10 +000058
59 def chain2(*iterables):
60 'Pure python version in the docs'
61 for it in iterables:
62 for element in it:
63 yield element
64
65 for c in (chain, chain2):
66 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
67 self.assertEqual(list(c('abc')), list('abc'))
68 self.assertEqual(list(c('')), [])
69 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
70 self.assertRaises(TypeError, list,c(2, 3))
Christian Heimesf16baeb2008-02-29 14:57:44 +000071
72 def test_chain_from_iterable(self):
73 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
74 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
75 self.assertEqual(list(chain.from_iterable([''])), [])
76 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
77 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000078
Christian Heimes380f7f22008-02-28 11:19:05 +000079 def test_combinations(self):
Raymond Hettinger5bad41e2009-01-08 21:01:54 +000080 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Christian Heimes380f7f22008-02-28 11:19:05 +000081 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Christian Heimes78644762008-03-04 23:39:23 +000082 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Christian Heimes380f7f22008-02-28 11:19:05 +000083 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Raymond Hettinger5bad41e2009-01-08 21:01:54 +000084 self.assertEqual(list(combinations('abc', 32)), []) # r > n
Christian Heimes380f7f22008-02-28 11:19:05 +000085 self.assertEqual(list(combinations(range(4), 3)),
86 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
Christian Heimes78644762008-03-04 23:39:23 +000087
88 def combinations1(iterable, r):
89 'Pure python version shown in the docs'
90 pool = tuple(iterable)
91 n = len(pool)
Raymond Hettinger5bad41e2009-01-08 21:01:54 +000092 if r > n:
93 return
Christian Heimes78644762008-03-04 23:39:23 +000094 indices = list(range(r))
95 yield tuple(pool[i] for i in indices)
96 while 1:
97 for i in reversed(range(r)):
98 if indices[i] != i + n - r:
99 break
100 else:
101 return
102 indices[i] += 1
103 for j in range(i+1, r):
104 indices[j] = indices[j-1] + 1
105 yield tuple(pool[i] for i in indices)
106
107 def combinations2(iterable, r):
108 'Pure python version shown in the docs'
109 pool = tuple(iterable)
110 n = len(pool)
111 for indices in permutations(range(n), r):
112 if sorted(indices) == list(indices):
113 yield tuple(pool[i] for i in indices)
114
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000115 def combinations3(iterable, r):
116 'Pure python version from cwr()'
117 pool = tuple(iterable)
118 n = len(pool)
119 for indices in combinations_with_replacement(range(n), r):
120 if len(set(indices)) == r:
121 yield tuple(pool[i] for i in indices)
122
Christian Heimes78644762008-03-04 23:39:23 +0000123 for n in range(7):
Christian Heimes380f7f22008-02-28 11:19:05 +0000124 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000125 for r in range(n+2):
Christian Heimes380f7f22008-02-28 11:19:05 +0000126 result = list(combinations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000127 self.assertEqual(len(result), 0 if r>n else fact(n) / fact(r) / fact(n-r)) # right number of combs
Christian Heimes380f7f22008-02-28 11:19:05 +0000128 self.assertEqual(len(result), len(set(result))) # no repeats
129 self.assertEqual(result, sorted(result)) # lexicographic order
130 for c in result:
131 self.assertEqual(len(c), r) # r-length combinations
132 self.assertEqual(len(set(c)), r) # no duplicate elements
133 self.assertEqual(list(c), sorted(c)) # keep original ordering
Georg Brandlab91fde2009-08-13 08:51:18 +0000134 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000135 self.assertEqual(list(c),
136 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000137 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000138 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000139 self.assertEqual(result, list(combinations3(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000140
141 # Test implementation detail: tuple re-use
142 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
143 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
144
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000145 def test_combinations_with_replacement(self):
146 cwr = combinations_with_replacement
147 self.assertRaises(TypeError, cwr, 'abc') # missing r argument
148 self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
149 self.assertRaises(TypeError, cwr, None) # pool is not iterable
150 self.assertRaises(ValueError, cwr, 'abc', -2) # r is negative
151 self.assertEqual(list(cwr('ABC', 2)),
152 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
153
154 def cwr1(iterable, r):
155 'Pure python version shown in the docs'
156 # number items returned: (n+r-1)! / r! / (n-1)! when n>0
157 pool = tuple(iterable)
158 n = len(pool)
159 if not n and r:
160 return
161 indices = [0] * r
162 yield tuple(pool[i] for i in indices)
163 while 1:
164 for i in reversed(range(r)):
165 if indices[i] != n - 1:
166 break
167 else:
168 return
169 indices[i:] = [indices[i] + 1] * (r - i)
170 yield tuple(pool[i] for i in indices)
171
172 def cwr2(iterable, r):
173 'Pure python version shown in the docs'
174 pool = tuple(iterable)
175 n = len(pool)
176 for indices in product(range(n), repeat=r):
177 if sorted(indices) == list(indices):
178 yield tuple(pool[i] for i in indices)
179
180 def numcombs(n, r):
181 if not n:
182 return 0 if r else 1
183 return fact(n+r-1) / fact(r)/ fact(n-1)
184
185 for n in range(7):
186 values = [5*x-12 for x in range(n)]
187 for r in range(n+2):
188 result = list(cwr(values, r))
189
190 self.assertEqual(len(result), numcombs(n, r)) # right number of combs
191 self.assertEqual(len(result), len(set(result))) # no repeats
192 self.assertEqual(result, sorted(result)) # lexicographic order
193
194 regular_combs = list(combinations(values, r)) # compare to combs without replacement
195 if n == 0 or r <= 1:
196 self.assertEquals(result, regular_combs) # cases that should be identical
197 else:
Georg Brandlab91fde2009-08-13 08:51:18 +0000198 self.assertTrue(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000199
200 for c in result:
201 self.assertEqual(len(c), r) # r-length combinations
202 noruns = [k for k,v in groupby(c)] # combo without consecutive repeats
203 self.assertEqual(len(noruns), len(set(noruns))) # no repeats other than consecutive
204 self.assertEqual(list(c), sorted(c)) # keep original ordering
Georg Brandlab91fde2009-08-13 08:51:18 +0000205 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000206 self.assertEqual(noruns,
207 [e for e in values if e in c]) # comb is a subsequence of the input iterable
208 self.assertEqual(result, list(cwr1(values, r))) # matches first pure python version
209 self.assertEqual(result, list(cwr2(values, r))) # matches second pure python version
210
211 # Test implementation detail: tuple re-use
212 self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
213 self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
214
Christian Heimes78644762008-03-04 23:39:23 +0000215 def test_permutations(self):
216 self.assertRaises(TypeError, permutations) # too few arguments
217 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000218 self.assertRaises(TypeError, permutations, None) # pool is not iterable
219 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000220 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000221 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Christian Heimes78644762008-03-04 23:39:23 +0000222 self.assertEqual(list(permutations(range(3), 2)),
223 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
224
225 def permutations1(iterable, r=None):
226 'Pure python version shown in the docs'
227 pool = tuple(iterable)
228 n = len(pool)
229 r = n if r is None else r
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000230 if r > n:
231 return
Christian Heimes78644762008-03-04 23:39:23 +0000232 indices = list(range(n))
233 cycles = list(range(n-r+1, n+1))[::-1]
234 yield tuple(pool[i] for i in indices[:r])
235 while n:
236 for i in reversed(range(r)):
237 cycles[i] -= 1
238 if cycles[i] == 0:
239 indices[i:] = indices[i+1:] + indices[i:i+1]
240 cycles[i] = n - i
241 else:
242 j = cycles[i]
243 indices[i], indices[-j] = indices[-j], indices[i]
244 yield tuple(pool[i] for i in indices[:r])
245 break
246 else:
247 return
248
249 def permutations2(iterable, r=None):
250 'Pure python version shown in the docs'
251 pool = tuple(iterable)
252 n = len(pool)
253 r = n if r is None else r
254 for indices in product(range(n), repeat=r):
255 if len(set(indices)) == r:
256 yield tuple(pool[i] for i in indices)
257
258 for n in range(7):
259 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000260 for r in range(n+2):
Christian Heimes78644762008-03-04 23:39:23 +0000261 result = list(permutations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000262 self.assertEqual(len(result), 0 if r>n else fact(n) / fact(n-r)) # right number of perms
Christian Heimes78644762008-03-04 23:39:23 +0000263 self.assertEqual(len(result), len(set(result))) # no repeats
264 self.assertEqual(result, sorted(result)) # lexicographic order
265 for p in result:
266 self.assertEqual(len(p), r) # r-length permutations
267 self.assertEqual(len(set(p)), r) # no duplicate elements
Georg Brandlab91fde2009-08-13 08:51:18 +0000268 self.assertTrue(all(e in values for e in p)) # elements taken from input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000269 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000270 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000271 if r == n:
272 self.assertEqual(result, list(permutations(values, None))) # test r as None
273 self.assertEqual(result, list(permutations(values))) # test default r
274
275 # Test implementation detail: tuple re-use
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000276 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Christian Heimes78644762008-03-04 23:39:23 +0000277 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Christian Heimes380f7f22008-02-28 11:19:05 +0000278
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000279 def test_combinatorics(self):
280 # Test relationships between product(), permutations(),
281 # combinations() and combinations_with_replacement().
282
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000283 for n in range(6):
284 s = 'ABCDEFG'[:n]
285 for r in range(8):
286 prod = list(product(s, repeat=r))
287 cwr = list(combinations_with_replacement(s, r))
288 perm = list(permutations(s, r))
289 comb = list(combinations(s, r))
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000290
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000291 # Check size
292 self.assertEquals(len(prod), n**r)
293 self.assertEquals(len(cwr), (fact(n+r-1) / fact(r)/ fact(n-1)) if n else (not r))
294 self.assertEquals(len(perm), 0 if r>n else fact(n) / fact(n-r))
295 self.assertEquals(len(comb), 0 if r>n else fact(n) / fact(r) / fact(n-r))
296
297 # Check lexicographic order without repeated tuples
298 self.assertEquals(prod, sorted(set(prod)))
299 self.assertEquals(cwr, sorted(set(cwr)))
300 self.assertEquals(perm, sorted(set(perm)))
301 self.assertEquals(comb, sorted(set(comb)))
302
303 # Check interrelationships
304 self.assertEquals(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
305 self.assertEquals(perm, [t for t in prod if len(set(t))==r]) # perm: prods with no dups
306 self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
307 self.assertEqual(comb, [t for t in cwr if len(set(t))==r]) # comb: cwrs without dups
308 self.assertEqual(comb, list(filter(set(cwr).__contains__, perm))) # comb: perm that is a cwr
309 self.assertEqual(comb, list(filter(set(perm).__contains__, cwr))) # comb: cwr that is a perm
310 self.assertEqual(comb, sorted(set(cwr) & set(perm))) # comb: both a cwr and a perm
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000311
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000312 def test_compress(self):
Raymond Hettinger15a49502009-02-19 02:17:09 +0000313 self.assertEqual(list(compress(data='ABCDEF', selectors=[1,0,1,0,1,1])), list('ACEF'))
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000314 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
315 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
316 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
317 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
318 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
319 n = 10000
320 data = chain.from_iterable(repeat(range(6), n))
321 selectors = chain.from_iterable(repeat((0, 1)))
322 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
323 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
324 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
325 self.assertRaises(TypeError, compress, range(6)) # too few args
326 self.assertRaises(TypeError, compress, range(6), None) # too many args
327
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000328 def test_count(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000329 self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
330 self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
331 self.assertEqual(take(2, lzip('abc',count(3))), [('a', 3), ('b', 4)])
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000332 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
333 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000334 self.assertRaises(TypeError, count, 2, 3, 4)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000335 self.assertRaises(TypeError, count, 'a')
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000336 self.assertEqual(list(islice(count(maxsize-5), 10)),
337 list(range(maxsize-5, maxsize+5)))
338 self.assertEqual(list(islice(count(-maxsize-5), 10)),
339 list(range(-maxsize-5, -maxsize+5)))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000340 c = count(3)
341 self.assertEqual(repr(c), 'count(3)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000342 next(c)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000343 self.assertEqual(repr(c), 'count(4)')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000344 c = count(-9)
345 self.assertEqual(repr(c), 'count(-9)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000346 next(c)
Raymond Hettinger30729212009-02-12 06:28:27 +0000347 self.assertEqual(repr(count(10.25)), 'count(10.25)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000348 self.assertEqual(next(c), -8)
Christian Heimesa37d4c62007-12-04 23:02:19 +0000349 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000350 # Test repr (ignoring the L in longs)
351 r1 = repr(count(i)).replace('L', '')
352 r2 = 'count(%r)'.__mod__(i).replace('L', '')
353 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000354
Raymond Hettinger30729212009-02-12 06:28:27 +0000355 def test_count_with_stride(self):
356 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingereb13fdd2009-02-21 22:30:12 +0000357 self.assertEqual(lzip('abc',count(start=2,step=3)),
358 [('a', 2), ('b', 5), ('c', 8)])
359 self.assertEqual(lzip('abc',count(step=-1)),
360 [('a', 0), ('b', -1), ('c', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000361 self.assertEqual(lzip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
362 self.assertEqual(lzip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000363 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000364 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
365 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
366 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettinger8a882a72009-02-12 12:08:18 +0000367 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
368 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerde09acf2009-02-12 12:53:53 +0000369 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
370 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000371 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
372 c = count(3, 5)
373 self.assertEqual(repr(c), 'count(3, 5)')
374 next(c)
375 self.assertEqual(repr(c), 'count(8, 5)')
376 c = count(-9, 0)
377 self.assertEqual(repr(c), 'count(-9, 0)')
378 next(c)
379 self.assertEqual(repr(c), 'count(-9, 0)')
380 c = count(-9, -3)
381 self.assertEqual(repr(c), 'count(-9, -3)')
382 next(c)
383 self.assertEqual(repr(c), 'count(-12, -3)')
384 self.assertEqual(repr(c), 'count(-12, -3)')
385 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
386 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
387 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
388 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
389 for j in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 1, 10, sys.maxsize-5, sys.maxsize+5):
390 # Test repr (ignoring the L in longs)
391 r1 = repr(count(i, j)).replace('L', '')
392 if j == 1:
393 r2 = ('count(%r)' % i).replace('L', '')
394 else:
395 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
396 self.assertEqual(r1, r2)
397
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000398 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000399 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000400 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000401 self.assertRaises(TypeError, cycle)
402 self.assertRaises(TypeError, cycle, 5)
403 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000404
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000405 def test_groupby(self):
406 # Check whether it accepts arguments correctly
407 self.assertEqual([], list(groupby([])))
408 self.assertEqual([], list(groupby([], key=id)))
409 self.assertRaises(TypeError, list, groupby('abc', []))
410 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000411 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000412
413 # Check normal input
414 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
415 (2,15,22), (3,16,23), (3,17,23)]
416 dup = []
417 for k, g in groupby(s, lambda r:r[0]):
418 for elem in g:
419 self.assertEqual(k, elem[0])
420 dup.append(elem)
421 self.assertEqual(s, dup)
422
423 # Check nested case
424 dup = []
425 for k, g in groupby(s, lambda r:r[0]):
426 for ik, ig in groupby(g, lambda r:r[2]):
427 for elem in ig:
428 self.assertEqual(k, elem[0])
429 self.assertEqual(ik, elem[2])
430 dup.append(elem)
431 self.assertEqual(s, dup)
432
433 # Check case where inner iterator is not used
434 keys = [k for k, g in groupby(s, lambda r:r[0])]
435 expectedkeys = set([r[0] for r in s])
436 self.assertEqual(set(keys), expectedkeys)
437 self.assertEqual(len(keys), len(expectedkeys))
438
439 # Exercise pipes and filters style
440 s = 'abracadabra'
441 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000442 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000443 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
444 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000445 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000446 self.assertEqual(r, ['a', 'b', 'r'])
447 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000448 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000449 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
450 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000451 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000452 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
453
Georg Brandla18af4e2007-04-21 15:47:16 +0000454 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000455 class ExpectedError(Exception):
456 pass
457 def delayed_raise(n=0):
458 for i in range(n):
459 yield 'yo'
460 raise ExpectedError
461 def gulp(iterable, keyp=None, func=list):
462 return [func(g) for k, g in groupby(iterable, keyp)]
463
Georg Brandla18af4e2007-04-21 15:47:16 +0000464 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000465 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000466 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000467 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
468
469 # __cmp__ failure
470 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000471 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000472 raise ExpectedError
473 s = [DummyCmp(), DummyCmp(), None]
474
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000475 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000476 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000477 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000478 self.assertRaises(ExpectedError, gulp, s)
479
480 # keyfunc failure
481 def keyfunc(obj):
482 if keyfunc.skip > 0:
483 keyfunc.skip -= 1
484 return obj
485 else:
486 raise ExpectedError
487
488 # keyfunc failure on outer object
489 keyfunc.skip = 0
490 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
491 keyfunc.skip = 1
492 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
493
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000494 def test_filter(self):
495 self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
496 self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])
497 self.assertEqual(list(filter(bool, [0,1,0,2,0])), [1,2])
498 self.assertEqual(take(4, filter(isEven, count())), [0,2,4,6])
499 self.assertRaises(TypeError, filter)
500 self.assertRaises(TypeError, filter, lambda x:x)
501 self.assertRaises(TypeError, filter, lambda x:x, range(6), 7)
502 self.assertRaises(TypeError, filter, isEven, 3)
503 self.assertRaises(TypeError, next, filter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000504
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000505 def test_filterfalse(self):
506 self.assertEqual(list(filterfalse(isEven, range(6))), [1,3,5])
507 self.assertEqual(list(filterfalse(None, [0,1,0,2,0])), [0,0,0])
508 self.assertEqual(list(filterfalse(bool, [0,1,0,2,0])), [0,0,0])
509 self.assertEqual(take(4, filterfalse(isEven, count())), [1,3,5,7])
510 self.assertRaises(TypeError, filterfalse)
511 self.assertRaises(TypeError, filterfalse, lambda x:x)
512 self.assertRaises(TypeError, filterfalse, lambda x:x, range(6), 7)
513 self.assertRaises(TypeError, filterfalse, isEven, 3)
514 self.assertRaises(TypeError, next, filterfalse(range(6), range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000515
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000516 def test_zip(self):
517 # XXX This is rather silly now that builtin zip() calls zip()...
518 ans = [(x,y) for x, y in zip('abc',count())]
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000519 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000520 self.assertEqual(list(zip('abc', range(6))), lzip('abc', range(6)))
521 self.assertEqual(list(zip('abcdef', range(3))), lzip('abcdef', range(3)))
522 self.assertEqual(take(3,zip('abcdef', count())), lzip('abcdef', range(3)))
523 self.assertEqual(list(zip('abcdef')), lzip('abcdef'))
524 self.assertEqual(list(zip()), lzip())
525 self.assertRaises(TypeError, zip, 3)
526 self.assertRaises(TypeError, zip, range(3), 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000527 # Check tuple re-use (implementation detail)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000528 self.assertEqual([tuple(list(pair)) for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000529 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000530 self.assertEqual([pair for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000531 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000532 ids = list(map(id, zip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000533 self.assertEqual(min(ids), max(ids))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000534 ids = list(map(id, list(zip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000535 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000536
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000537 def test_ziplongest(self):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000538 for args in [
539 ['abc', range(6)],
540 [range(6), 'abc'],
541 [range(1000), range(2000,2100), range(3000,3050)],
542 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
543 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
544 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000545 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
546 for i in range(max(map(len, args)))]
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000547 self.assertEqual(list(zip_longest(*args)), target)
548 self.assertEqual(list(zip_longest(*args, **{})), target)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000549 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 +0000550 self.assertEqual(list(zip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000551
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000552 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 +0000553
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000554 self.assertEqual(list(zip_longest()), list(zip()))
555 self.assertEqual(list(zip_longest([])), list(zip([])))
556 self.assertEqual(list(zip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000557
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000558 self.assertEqual(list(zip_longest('abc', 'defg', **{})),
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000559 list(zip(list('abc')+[None], 'defg'))) # empty keyword dict
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000560 self.assertRaises(TypeError, zip_longest, 3)
561 self.assertRaises(TypeError, zip_longest, range(3), 3)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000562
563 for stmt in [
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000564 "zip_longest('abc', fv=1)",
565 "zip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000566 ]:
567 try:
568 eval(stmt, globals(), locals())
569 except TypeError:
570 pass
571 else:
572 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000573
Thomas Wouterscf297e42007-02-23 15:07:44 +0000574 # Check tuple re-use (implementation detail)
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000575 self.assertEqual([tuple(list(pair)) for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000576 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000577 self.assertEqual([pair for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000578 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000579 ids = list(map(id, zip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000580 self.assertEqual(min(ids), max(ids))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000581 ids = list(map(id, list(zip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000582 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
583
Raymond Hettingera9311a32009-11-01 21:02:38 +0000584 def test_bug_7244(self):
585
586 class Repeater:
587 # this class is similar to itertools.repeat
588 def __init__(self, o, t, e):
589 self.o = o
590 self.t = int(t)
591 self.e = e
592 def __iter__(self): # its iterator is itself
593 return self
594 def __next__(self):
595 if self.t > 0:
596 self.t -= 1
597 return self.o
598 else:
599 raise self.e
600
601 # Formerly this code in would fail in debug mode
602 # with Undetected Error and Stop Iteration
603 r1 = Repeater(1, 3, StopIteration)
604 r2 = Repeater(2, 4, StopIteration)
605 def run(r1, r2):
606 result = []
607 for i, j in zip_longest(r1, r2, fillvalue=0):
608 with support.captured_output('stdout'):
609 print((i, j))
610 result.append((i, j))
611 return result
612 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
613
614 # Formerly, the RuntimeError would be lost
615 # and StopIteration would stop as expected
616 r1 = Repeater(1, 3, RuntimeError)
617 r2 = Repeater(2, 4, StopIteration)
618 it = zip_longest(r1, r2, fillvalue=0)
619 self.assertEqual(next(it), (1, 2))
620 self.assertEqual(next(it), (1, 2))
621 self.assertEqual(next(it), (1, 2))
622 self.assertRaises(RuntimeError, next, it)
623
Christian Heimesc3f30c42008-02-22 16:37:40 +0000624 def test_product(self):
625 for args, result in [
Christian Heimes78644762008-03-04 23:39:23 +0000626 ([], [()]), # zero iterables
Christian Heimesc3f30c42008-02-22 16:37:40 +0000627 (['ab'], [('a',), ('b',)]), # one iterable
628 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
629 ([range(0), range(2), range(3)], []), # first iterable with zero length
630 ([range(2), range(0), range(3)], []), # middle iterable with zero length
631 ([range(2), range(3), range(0)], []), # last iterable with zero length
632 ]:
633 self.assertEqual(list(product(*args)), result)
Christian Heimesf16baeb2008-02-29 14:57:44 +0000634 for r in range(4):
635 self.assertEqual(list(product(*(args*r))),
636 list(product(*args, **dict(repeat=r))))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000637 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
638 self.assertRaises(TypeError, product, range(6), None)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000639
640 def product1(*args, **kwds):
641 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
642 n = len(pools)
643 if n == 0:
644 yield ()
645 return
646 if any(len(pool) == 0 for pool in pools):
647 return
648 indices = [0] * n
649 yield tuple(pool[i] for pool, i in zip(pools, indices))
650 while 1:
651 for i in reversed(range(n)): # right to left
652 if indices[i] == len(pools[i]) - 1:
653 continue
654 indices[i] += 1
655 for j in range(i+1, n):
656 indices[j] = 0
657 yield tuple(pool[i] for pool, i in zip(pools, indices))
658 break
659 else:
660 return
661
662 def product2(*args, **kwds):
663 'Pure python version used in docs'
664 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
665 result = [[]]
666 for pool in pools:
667 result = [x+[y] for x in result for y in pool]
668 for prod in result:
669 yield tuple(prod)
670
Christian Heimesc3f30c42008-02-22 16:37:40 +0000671 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
672 set('abcdefg'), range(11), tuple(range(13))]
673 for i in range(100):
674 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Christian Heimes78644762008-03-04 23:39:23 +0000675 expected_len = prod(map(len, args))
676 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000677 self.assertEqual(list(product(*args)), list(product1(*args)))
678 self.assertEqual(list(product(*args)), list(product2(*args)))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000679 args = map(iter, args)
Christian Heimes78644762008-03-04 23:39:23 +0000680 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000681
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000682 # Test implementation detail: tuple re-use
683 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
684 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000685
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000686 def test_repeat(self):
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +0000687 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
Guido van Rossum805365e2007-05-07 22:24:25 +0000688 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000689 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000690 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000691 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000692 self.assertEqual(list(repeat('a', 0)), [])
693 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000694 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000695 self.assertRaises(TypeError, repeat, None, 3, 4)
696 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000697 r = repeat(1+0j)
698 self.assertEqual(repr(r), 'repeat((1+0j))')
699 r = repeat(1+0j, 5)
700 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
701 list(r)
702 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000703
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000704 def test_map(self):
705 self.assertEqual(list(map(operator.pow, range(3), range(1,7))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000706 [0**1, 1**2, 2**3])
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000707 def tupleize(*args):
708 return args
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000709 self.assertEqual(list(map(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000710 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000711 self.assertEqual(list(map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000712 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000713 self.assertEqual(take(2,map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000714 [('a',0),('b',1)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000715 self.assertEqual(list(map(operator.pow, [])), [])
716 self.assertRaises(TypeError, map)
717 self.assertRaises(TypeError, list, map(None, range(3), range(3)))
718 self.assertRaises(TypeError, map, operator.neg)
719 self.assertRaises(TypeError, next, map(10, range(5)))
720 self.assertRaises(ValueError, next, map(errfunc, [4], [5]))
721 self.assertRaises(TypeError, next, map(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000722
723 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000724 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
725 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000726 self.assertEqual(take(3, starmap(operator.pow, zip(count(), count(1)))),
Raymond Hettinger02420702003-06-29 20:36:23 +0000727 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000728 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +0000729 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
730 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000731 self.assertRaises(TypeError, starmap)
732 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000733 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
734 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
735 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000736
737 def test_islice(self):
738 for args in [ # islice(args) should agree with range(args)
739 (10, 20, 3),
740 (10, 3, 20),
741 (10, 20),
742 (10, 3),
743 (20,)
744 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000745 self.assertEqual(list(islice(range(100), *args)),
746 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000747
748 for args, tgtargs in [ # Stop when seqn is exhausted
749 ((10, 110, 3), ((10, 100, 3))),
750 ((10, 110), ((10, 100))),
751 ((110,), (100,))
752 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000753 self.assertEqual(list(islice(range(100), *args)),
754 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000755
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000756 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +0000757 self.assertEqual(list(islice(range(10), None)), list(range(10)))
758 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
759 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
760 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
761 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000762
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000763 # Test number of items consumed SF #1171417
764 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +0000765 self.assertEqual(list(islice(it, 3)), list(range(3)))
766 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000767
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000768 # Test invalid arguments
Guido van Rossum805365e2007-05-07 22:24:25 +0000769 self.assertRaises(TypeError, islice, range(10))
770 self.assertRaises(TypeError, islice, range(10), 1, 2, 3, 4)
771 self.assertRaises(ValueError, islice, range(10), -5, 10, 1)
772 self.assertRaises(ValueError, islice, range(10), 1, -5, -1)
773 self.assertRaises(ValueError, islice, range(10), 1, 10, -1)
774 self.assertRaises(ValueError, islice, range(10), 1, 10, 0)
775 self.assertRaises(ValueError, islice, range(10), 'a')
776 self.assertRaises(ValueError, islice, range(10), 'a', 1)
777 self.assertRaises(ValueError, islice, range(10), 1, 'a')
778 self.assertRaises(ValueError, islice, range(10), 'a', 1, 1)
779 self.assertRaises(ValueError, islice, range(10), 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +0000780 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000781
782 def test_takewhile(self):
783 data = [1, 3, 5, 20, 2, 4, 6, 8]
784 underten = lambda x: x<10
785 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000786 self.assertEqual(list(takewhile(underten, [])), [])
787 self.assertRaises(TypeError, takewhile)
788 self.assertRaises(TypeError, takewhile, operator.pow)
789 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000790 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
791 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000792 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
793 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +0000794 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000795
796 def test_dropwhile(self):
797 data = [1, 3, 5, 20, 2, 4, 6, 8]
798 underten = lambda x: x<10
799 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000800 self.assertEqual(list(dropwhile(underten, [])), [])
801 self.assertRaises(TypeError, dropwhile)
802 self.assertRaises(TypeError, dropwhile, operator.pow)
803 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000804 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
805 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000806
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000807 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000808 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000809 def irange(n):
Guido van Rossum805365e2007-05-07 22:24:25 +0000810 for i in range(n):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000811 yield i
812
813 a, b = tee([]) # test empty iterator
814 self.assertEqual(list(a), [])
815 self.assertEqual(list(b), [])
816
817 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +0000818 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000819
820 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +0000821 self.assertEqual(list(a), list(range(n)))
822 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000823
824 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000825 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000826 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000827 del a
Guido van Rossum805365e2007-05-07 22:24:25 +0000828 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000829
830 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000831 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000832 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000833 del b
Guido van Rossum805365e2007-05-07 22:24:25 +0000834 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000835
Guido van Rossum805365e2007-05-07 22:24:25 +0000836 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000837 order = [0]*n + [1]*n
838 random.shuffle(order)
839 lists = ([], [])
840 its = tee(irange(n))
841 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +0000842 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000843 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +0000844 self.assertEqual(lists[0], list(range(n)))
845 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000846
Raymond Hettingerad983e72003-11-12 14:32:26 +0000847 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000848 self.assertRaises(TypeError, tee)
849 self.assertRaises(TypeError, tee, 3)
850 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000851 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000852
Raymond Hettingerad983e72003-11-12 14:32:26 +0000853 # tee object should be instantiable
854 a, b = tee('abc')
855 c = type(a)('def')
856 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000857
Raymond Hettingerad983e72003-11-12 14:32:26 +0000858 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +0000859 a, b, c = tee(range(2000), 3)
860 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000861 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +0000862 self.assertEqual(list(b), list(range(2000)))
863 self.assertEqual([next(c), next(c)], list(range(2)))
864 self.assertEqual(list(a), list(range(100,2000)))
865 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000866
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000867 # test values of n
868 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000869 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +0000870 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000871 result = tee('abc', n)
872 self.assertEqual(type(result), tuple)
873 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000874 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000875
Raymond Hettingerad983e72003-11-12 14:32:26 +0000876 # tee pass-through to copyable iterator
877 a, b = tee('abc')
878 c, d = tee(a)
Georg Brandlab91fde2009-08-13 08:51:18 +0000879 self.assertTrue(a is c)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000880
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000881 # test tee_new
882 t1, t2 = tee('abc')
883 tnew = type(t1)
884 self.assertRaises(TypeError, tnew)
885 self.assertRaises(TypeError, tnew, 10)
886 t3 = tnew(t1)
Georg Brandlab91fde2009-08-13 08:51:18 +0000887 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000888
Raymond Hettingera9f60922004-10-17 16:40:14 +0000889 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +0000890 a, b = tee(range(10))
Raymond Hettingera9f60922004-10-17 16:40:14 +0000891 p = proxy(a)
892 self.assertEqual(getattr(p, '__class__'), type(b))
893 del a
894 self.assertRaises(ReferenceError, getattr, p, '__class__')
895
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000896 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000897 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000898
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000899 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +0000900 self.assertRaises(StopIteration, next, f([]))
901 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000902
Georg Brandla18af4e2007-04-21 15:47:16 +0000903 self.assertRaises(StopIteration, next, islice([], None))
904 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000905
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000906 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +0000907 self.assertRaises(StopIteration, next, p)
908 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000909 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +0000910 self.assertRaises(StopIteration, next, p)
911 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000912
Georg Brandla18af4e2007-04-21 15:47:16 +0000913 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000914
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000915 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +0000916 self.assertRaises(StopIteration, next, f(lambda x:x, []))
917 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000918
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000919class TestExamples(unittest.TestCase):
920
921 def test_chain(self):
922 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
923
924 def test_chain_from_iterable(self):
925 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
926
927 def test_combinations(self):
928 self.assertEqual(list(combinations('ABCD', 2)),
929 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
930 self.assertEqual(list(combinations(range(4), 3)),
931 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
932
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000933 def test_combinations_with_replacement(self):
934 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
935 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
936
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000937 def test_compress(self):
938 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
939
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000940 def test_count(self):
941 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
942
943 def test_cycle(self):
944 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
945
946 def test_dropwhile(self):
947 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
948
949 def test_groupby(self):
950 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
951 list('ABCDAB'))
952 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
953 [list('AAAA'), list('BBB'), list('CC'), list('D')])
954
955 def test_filter(self):
956 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
957
958 def test_filterfalse(self):
959 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
960
961 def test_map(self):
962 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
963
964 def test_islice(self):
965 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
966 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
967 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
968 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
969
970 def test_zip(self):
971 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
972
973 def test_zip_longest(self):
974 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
975 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
976
977 def test_permutations(self):
978 self.assertEqual(list(permutations('ABCD', 2)),
979 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
980 self.assertEqual(list(permutations(range(3))),
981 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
982
983 def test_product(self):
984 self.assertEqual(list(product('ABCD', 'xy')),
985 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
986 self.assertEqual(list(product(range(2), repeat=3)),
987 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
988 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
989
990 def test_repeat(self):
991 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
992
993 def test_stapmap(self):
994 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
995 [32, 9, 1000])
996
997 def test_takewhile(self):
998 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
999
1000
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001001class TestGC(unittest.TestCase):
1002
1003 def makecycle(self, iterator, container):
1004 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +00001005 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001006 del container, iterator
1007
1008 def test_chain(self):
1009 a = []
1010 self.makecycle(chain(a), a)
1011
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001012 def test_chain_from_iterable(self):
1013 a = []
1014 self.makecycle(chain.from_iterable([a]), a)
1015
1016 def test_combinations(self):
1017 a = []
1018 self.makecycle(combinations([1,2,a,3], 3), a)
1019
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001020 def test_combinations_with_replacement(self):
1021 a = []
1022 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1023
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001024 def test_compress(self):
1025 a = []
1026 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1027
Raymond Hettingerd6280f42009-02-16 20:50:56 +00001028 def test_count(self):
1029 a = []
1030 Int = type('Int', (int,), dict(x=a))
1031 self.makecycle(count(Int(0), Int(1)), a)
1032
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001033 def test_cycle(self):
1034 a = []
1035 self.makecycle(cycle([a]*2), a)
1036
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001037 def test_dropwhile(self):
1038 a = []
1039 self.makecycle(dropwhile(bool, [0, a, a]), a)
1040
1041 def test_groupby(self):
1042 a = []
1043 self.makecycle(groupby([a]*2, lambda x:x), a)
1044
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001045 def test_issue2246(self):
1046 # Issue 2246 -- the _grouper iterator was not included in GC
1047 n = 10
1048 keyfunc = lambda x: x
1049 for i, j in groupby(range(n), key=keyfunc):
1050 keyfunc.__dict__.setdefault('x',[]).append(j)
1051
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001052 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001053 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001054 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001055
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001056 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001057 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001058 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001059
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001060 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001061 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001062 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001063
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001064 def test_zip_longest(self):
1065 a = []
1066 self.makecycle(zip_longest([a]*2, [a]*3), a)
1067 b = [a, None]
1068 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1069
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001070 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001071 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001072 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001073
1074 def test_islice(self):
1075 a = []
1076 self.makecycle(islice([a]*2, None), a)
1077
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001078 def test_permutations(self):
1079 a = []
1080 self.makecycle(permutations([1,2,a,3], 3), a)
1081
1082 def test_product(self):
1083 a = []
1084 self.makecycle(product([1,2,a,3], repeat=3), a)
1085
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001086 def test_repeat(self):
1087 a = []
1088 self.makecycle(repeat(a), a)
1089
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001090 def test_starmap(self):
1091 a = []
1092 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1093
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001094 def test_takewhile(self):
1095 a = []
1096 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1097
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001098def R(seqn):
1099 'Regular generator'
1100 for i in seqn:
1101 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001102
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001103class G:
1104 'Sequence using __getitem__'
1105 def __init__(self, seqn):
1106 self.seqn = seqn
1107 def __getitem__(self, i):
1108 return self.seqn[i]
1109
1110class I:
1111 'Sequence using iterator protocol'
1112 def __init__(self, seqn):
1113 self.seqn = seqn
1114 self.i = 0
1115 def __iter__(self):
1116 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001117 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001118 if self.i >= len(self.seqn): raise StopIteration
1119 v = self.seqn[self.i]
1120 self.i += 1
1121 return v
1122
1123class Ig:
1124 'Sequence using iterator protocol defined with a generator'
1125 def __init__(self, seqn):
1126 self.seqn = seqn
1127 self.i = 0
1128 def __iter__(self):
1129 for val in self.seqn:
1130 yield val
1131
1132class X:
1133 'Missing __getitem__ and __iter__'
1134 def __init__(self, seqn):
1135 self.seqn = seqn
1136 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001137 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001138 if self.i >= len(self.seqn): raise StopIteration
1139 v = self.seqn[self.i]
1140 self.i += 1
1141 return v
1142
1143class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001144 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001145 def __init__(self, seqn):
1146 self.seqn = seqn
1147 self.i = 0
1148 def __iter__(self):
1149 return self
1150
1151class E:
1152 'Test propagation of exceptions'
1153 def __init__(self, seqn):
1154 self.seqn = seqn
1155 self.i = 0
1156 def __iter__(self):
1157 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001158 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001159 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001160
1161class S:
1162 'Test immediate stop'
1163 def __init__(self, seqn):
1164 pass
1165 def __iter__(self):
1166 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001167 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001168 raise StopIteration
1169
1170def L(seqn):
1171 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001172 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001173
1174
1175class TestVariousIteratorArgs(unittest.TestCase):
1176
1177 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001178 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001179 for g in (G, I, Ig, S, L, R):
1180 self.assertEqual(list(chain(g(s))), list(g(s)))
1181 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001182 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001183 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001184 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1185
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001186 def test_compress(self):
1187 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1188 n = len(s)
1189 for g in (G, I, Ig, S, L, R):
1190 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1191 self.assertRaises(TypeError, compress, X(s), repeat(1))
1192 self.assertRaises(TypeError, compress, N(s), repeat(1))
1193 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1194
Christian Heimesc3f30c42008-02-22 16:37:40 +00001195 def test_product(self):
1196 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1197 self.assertRaises(TypeError, product, X(s))
1198 self.assertRaises(TypeError, product, N(s))
1199 self.assertRaises(ZeroDivisionError, product, E(s))
1200
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001201 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001202 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001203 for g in (G, I, Ig, S, L, R):
1204 tgtlen = len(s) * 3
1205 expected = list(g(s))*3
1206 actual = list(islice(cycle(g(s)), tgtlen))
1207 self.assertEqual(actual, expected)
1208 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001209 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001210 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1211
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001212 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001213 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001214 for g in (G, I, Ig, S, L, R):
1215 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1216 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001217 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001218 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1219
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001220 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001221 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001222 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001223 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001224 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001225 self.assertRaises(TypeError, filter, isEven, X(s))
1226 self.assertRaises(TypeError, filter, isEven, N(s))
1227 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001228
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001229 def test_filterfalse(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 Hettingerb0002d22008-03-13 01:41:43 +00001232 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001233 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001234 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1235 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1236 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001237
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001238 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001239 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001240 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001241 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1242 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1243 self.assertRaises(TypeError, zip, X(s))
1244 self.assertRaises(TypeError, zip, N(s))
1245 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001246
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001247 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001248 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001249 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001250 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1251 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1252 self.assertRaises(TypeError, zip_longest, X(s))
1253 self.assertRaises(TypeError, zip_longest, N(s))
1254 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001255
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001256 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001257 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001258 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001259 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001260 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001261 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001262 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001263 self.assertRaises(TypeError, map, onearg, X(s))
1264 self.assertRaises(TypeError, map, onearg, N(s))
1265 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001266
1267 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001268 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001269 for g in (G, I, Ig, S, L, R):
1270 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1271 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001272 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001273 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1274
1275 def test_starmap(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):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001278 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001279 self.assertEqual(list(starmap(operator.pow, g(ss))),
1280 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001281 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001282 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001283 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1284
1285 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001286 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001287 for g in (G, I, Ig, S, L, R):
1288 tgt = []
1289 for elem in g(s):
1290 if not isEven(elem): break
1291 tgt.append(elem)
1292 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1293 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001294 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001295 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1296
1297 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001298 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001299 for g in (G, I, Ig, S, L, R):
1300 tgt = []
1301 for elem in g(s):
1302 if not tgt and isOdd(elem): continue
1303 tgt.append(elem)
1304 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1305 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001306 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001307 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1308
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001309 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001310 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001311 for g in (G, I, Ig, S, L, R):
1312 it1, it2 = tee(g(s))
1313 self.assertEqual(list(it1), list(g(s)))
1314 self.assertEqual(list(it2), list(g(s)))
1315 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001316 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001317 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1318
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001319class LengthTransparency(unittest.TestCase):
1320
1321 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001322 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001323 self.assertEqual(len(repeat(None, 50)), 50)
1324 self.assertRaises(TypeError, len, repeat(None))
1325
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001326class RegressionTests(unittest.TestCase):
1327
1328 def test_sf_793826(self):
1329 # Fix Armin Rigo's successful efforts to wreak havoc
1330
1331 def mutatingtuple(tuple1, f, tuple2):
1332 # this builds a tuple t which is a copy of tuple1,
1333 # then calls f(t), then mutates t to be equal to tuple2
1334 # (needs len(tuple1) == len(tuple2)).
1335 def g(value, first=[1]):
1336 if first:
1337 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00001338 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001339 return value
1340 items = list(tuple2)
1341 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001342 gen = map(g, items)
1343 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00001344 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001345
1346 def f(t):
1347 global T
1348 T = t
1349 first[:] = list(T)
1350
1351 first = []
1352 mutatingtuple((1,2,3), f, (4,5,6))
1353 second = list(T)
1354 self.assertEqual(first, second)
1355
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001356
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001357 def test_sf_950057(self):
1358 # Make sure that chain() and cycle() catch exceptions immediately
1359 # rather than when shifting between input sources
1360
1361 def gen1():
1362 hist.append(0)
1363 yield 1
1364 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001365 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001366 hist.append(2)
1367
1368 def gen2(x):
1369 hist.append(3)
1370 yield 2
1371 hist.append(4)
1372 if x:
1373 raise StopIteration
1374
1375 hist = []
1376 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1377 self.assertEqual(hist, [0,1])
1378
1379 hist = []
1380 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1381 self.assertEqual(hist, [0,1])
1382
1383 hist = []
1384 self.assertRaises(AssertionError, list, cycle(gen1()))
1385 self.assertEqual(hist, [0,1])
1386
Thomas Woutersb2137042007-02-01 18:02:27 +00001387class SubclassWithKwargsTest(unittest.TestCase):
1388 def test_keywords_in_subclass(self):
1389 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001390 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001391 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00001392 class Subclass(cls):
1393 def __init__(self, newarg=None, *args):
1394 cls.__init__(self, *args)
1395 try:
1396 Subclass(newarg=1)
1397 except TypeError as err:
1398 # we expect type errors because of wrong argument count
Georg Brandlab91fde2009-08-13 08:51:18 +00001399 self.assertFalse("does not take keyword arguments" in err.args[0])
Thomas Woutersb2137042007-02-01 18:02:27 +00001400
1401
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001402libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001403
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001404
1405>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001406>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00001407... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001408...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001409Check 1200 is for $120.15
1410Check 1201 is for $764.05
1411Check 1202 is for $823.14
1412
1413>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001414>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00001415... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00001416...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000014171
14188
141927
1420
1421>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001422>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00001423... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00001424...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001425Alex
1426Laura
1427Martin
1428Walter
1429Samuele
1430
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001431>>> from operator import itemgetter
1432>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00001433>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001434>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001435... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001436...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000014371 ['a', 'c', 'e']
14382 ['b', 'd', 'f']
14393 ['g']
1440
Raymond Hettinger734fb572004-01-20 20:04:40 +00001441# Find runs of consecutive numbers using groupby. The key to the solution
1442# is differencing with a range so that consecutive numbers all appear in
1443# same group.
1444>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00001445>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001446... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001447...
Raymond Hettinger734fb572004-01-20 20:04:40 +00001448[1]
1449[4, 5, 6]
1450[10]
1451[15, 16, 17, 18]
1452[22]
1453[25, 26, 27, 28]
1454
Georg Brandl3dbca812008-07-23 16:10:53 +00001455>>> def take(n, iterable):
1456... "Return first n items of the iterable as a list"
1457... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001458
Georg Brandl3dbca812008-07-23 16:10:53 +00001459>>> def enumerate(iterable, start=0):
1460... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001461
Georg Brandl3dbca812008-07-23 16:10:53 +00001462>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001463... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00001464... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001465
Raymond Hettingercdf8ba32009-02-19 04:45:07 +00001466>>> def nth(iterable, n, default=None):
1467... "Returns the nth item or a default value"
1468... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001469
Georg Brandl3dbca812008-07-23 16:10:53 +00001470>>> def quantify(iterable, pred=bool):
1471... "Count how many times the predicate is true"
1472... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001473
Georg Brandl3dbca812008-07-23 16:10:53 +00001474>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001475... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00001476... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001477
Georg Brandl3dbca812008-07-23 16:10:53 +00001478>>> def ncycles(iterable, n):
1479... "Returns the seqeuence elements n times"
1480... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001481
1482>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001483... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001484
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001485>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001486... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001487
1488>>> def repeatfunc(func, times=None, *args):
1489... "Repeat calls to func with specified arguments."
1490... " Example: repeatfunc(random.random)"
1491... if times is None:
1492... return starmap(func, repeat(args))
1493... else:
1494... return starmap(func, repeat(args, times))
1495
Raymond Hettingerd591f662003-10-26 15:34:50 +00001496>>> def pairwise(iterable):
1497... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1498... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001499... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00001500... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001501... except StopIteration:
1502... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001503... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001504
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001505>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001506... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001507... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00001508... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001509
1510>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001511... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001512... # Recipe credited to George Sakkis
1513... pending = len(iterables)
1514... nexts = cycle(iter(it).__next__ for it in iterables)
1515... while pending:
1516... try:
1517... for next in nexts:
1518... yield next()
1519... except StopIteration:
1520... pending -= 1
1521... nexts = cycle(islice(nexts, pending))
1522
1523>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00001524... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1525... s = list(iterable)
1526... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001527
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001528>>> def unique_everseen(iterable, key=None):
1529... "List unique elements, preserving order. Remember all elements ever seen."
1530... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1531... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1532... seen = set()
1533... seen_add = seen.add
1534... if key is None:
1535... for element in iterable:
1536... if element not in seen:
1537... seen_add(element)
1538... yield element
1539... else:
1540... for element in iterable:
1541... k = key(element)
1542... if k not in seen:
1543... seen_add(k)
1544... yield element
1545
1546>>> def unique_justseen(iterable, key=None):
1547... "List unique elements, preserving order. Remember only the element just seen."
1548... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1549... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1550... return map(next, map(itemgetter(1), groupby(iterable, key)))
1551
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001552This is not part of the examples but it tests to make sure the definitions
1553perform as purported.
1554
Raymond Hettingera098b332003-09-08 23:58:40 +00001555>>> take(10, count())
1556[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1557
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001558>>> list(enumerate('abc'))
1559[(0, 'a'), (1, 'b'), (2, 'c')]
1560
1561>>> list(islice(tabulate(lambda x: 2*x), 4))
1562[0, 2, 4, 6]
1563
1564>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00001565'd'
1566
1567>>> nth('abcde', 9) is None
1568True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001569
Guido van Rossum805365e2007-05-07 22:24:25 +00001570>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000157150
1572
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001573>>> a = [[1, 2, 3], [4, 5, 6]]
1574>>> flatten(a)
1575[1, 2, 3, 4, 5, 6]
1576
1577>>> list(repeatfunc(pow, 5, 2, 3))
1578[8, 8, 8, 8, 8]
1579
1580>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001581>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001582[0, 0, 0, 0, 0]
1583
Raymond Hettingerd591f662003-10-26 15:34:50 +00001584>>> list(pairwise('abcd'))
1585[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001586
Raymond Hettingerd591f662003-10-26 15:34:50 +00001587>>> list(pairwise([]))
1588[]
1589
1590>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001591[]
1592
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001593>>> list(islice(padnone('abc'), 0, 6))
1594['a', 'b', 'c', None, None, None]
1595
1596>>> list(ncycles('abc', 3))
1597['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1598
1599>>> dotproduct([1,2,3], [4,5,6])
160032
1601
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001602>>> list(grouper(3, 'abcdefg', 'x'))
1603[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1604
1605>>> list(roundrobin('abc', 'd', 'ef'))
1606['a', 'd', 'e', 'b', 'f', 'c']
1607
Raymond Hettingerace67332009-01-26 02:23:50 +00001608>>> list(powerset([1,2,3]))
1609[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001610
Raymond Hettinger191e8502009-01-27 13:29:43 +00001611>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1612True
1613
1614>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1615True
1616
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001617>>> list(unique_everseen('AAAABBBCCDAABBB'))
1618['A', 'B', 'C', 'D']
1619
1620>>> list(unique_everseen('ABBCcAD', str.lower))
1621['A', 'B', 'C', 'D']
1622
1623>>> list(unique_justseen('AAAABBBCCDAABBB'))
1624['A', 'B', 'C', 'D', 'A', 'B']
1625
1626>>> list(unique_justseen('ABBCcAD', str.lower))
1627['A', 'B', 'C', 'A', 'D']
1628
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001629"""
1630
1631__test__ = {'libreftest' : libreftest}
1632
1633def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001634 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00001635 RegressionTests, LengthTransparency,
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001636 SubclassWithKwargsTest, TestExamples)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001637 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001638
1639 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001640 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001641 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001642 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00001643 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001644 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001645 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001646 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001647 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001648
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001649 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001650 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001651
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001652if __name__ == "__main__":
1653 test_main(verbose=True)