blob: d917262fd444c3256340256ca3a16ed3c9076425 [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
Christian Heimesc3f30c42008-02-22 16:37:40 +0000584 def test_product(self):
585 for args, result in [
Christian Heimes78644762008-03-04 23:39:23 +0000586 ([], [()]), # zero iterables
Christian Heimesc3f30c42008-02-22 16:37:40 +0000587 (['ab'], [('a',), ('b',)]), # one iterable
588 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
589 ([range(0), range(2), range(3)], []), # first iterable with zero length
590 ([range(2), range(0), range(3)], []), # middle iterable with zero length
591 ([range(2), range(3), range(0)], []), # last iterable with zero length
592 ]:
593 self.assertEqual(list(product(*args)), result)
Christian Heimesf16baeb2008-02-29 14:57:44 +0000594 for r in range(4):
595 self.assertEqual(list(product(*(args*r))),
596 list(product(*args, **dict(repeat=r))))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000597 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
598 self.assertRaises(TypeError, product, range(6), None)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000599
600 def product1(*args, **kwds):
601 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
602 n = len(pools)
603 if n == 0:
604 yield ()
605 return
606 if any(len(pool) == 0 for pool in pools):
607 return
608 indices = [0] * n
609 yield tuple(pool[i] for pool, i in zip(pools, indices))
610 while 1:
611 for i in reversed(range(n)): # right to left
612 if indices[i] == len(pools[i]) - 1:
613 continue
614 indices[i] += 1
615 for j in range(i+1, n):
616 indices[j] = 0
617 yield tuple(pool[i] for pool, i in zip(pools, indices))
618 break
619 else:
620 return
621
622 def product2(*args, **kwds):
623 'Pure python version used in docs'
624 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
625 result = [[]]
626 for pool in pools:
627 result = [x+[y] for x in result for y in pool]
628 for prod in result:
629 yield tuple(prod)
630
Christian Heimesc3f30c42008-02-22 16:37:40 +0000631 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
632 set('abcdefg'), range(11), tuple(range(13))]
633 for i in range(100):
634 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Christian Heimes78644762008-03-04 23:39:23 +0000635 expected_len = prod(map(len, args))
636 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000637 self.assertEqual(list(product(*args)), list(product1(*args)))
638 self.assertEqual(list(product(*args)), list(product2(*args)))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000639 args = map(iter, args)
Christian Heimes78644762008-03-04 23:39:23 +0000640 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000641
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000642 # Test implementation detail: tuple re-use
643 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
644 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000645
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000646 def test_repeat(self):
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +0000647 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
Guido van Rossum805365e2007-05-07 22:24:25 +0000648 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000649 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000650 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000651 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000652 self.assertEqual(list(repeat('a', 0)), [])
653 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000654 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000655 self.assertRaises(TypeError, repeat, None, 3, 4)
656 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000657 r = repeat(1+0j)
658 self.assertEqual(repr(r), 'repeat((1+0j))')
659 r = repeat(1+0j, 5)
660 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
661 list(r)
662 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000663
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000664 def test_map(self):
665 self.assertEqual(list(map(operator.pow, range(3), range(1,7))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000666 [0**1, 1**2, 2**3])
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000667 def tupleize(*args):
668 return args
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000669 self.assertEqual(list(map(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000670 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000671 self.assertEqual(list(map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000672 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000673 self.assertEqual(take(2,map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000674 [('a',0),('b',1)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000675 self.assertEqual(list(map(operator.pow, [])), [])
676 self.assertRaises(TypeError, map)
677 self.assertRaises(TypeError, list, map(None, range(3), range(3)))
678 self.assertRaises(TypeError, map, operator.neg)
679 self.assertRaises(TypeError, next, map(10, range(5)))
680 self.assertRaises(ValueError, next, map(errfunc, [4], [5]))
681 self.assertRaises(TypeError, next, map(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000682
683 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000684 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
685 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000686 self.assertEqual(take(3, starmap(operator.pow, zip(count(), count(1)))),
Raymond Hettinger02420702003-06-29 20:36:23 +0000687 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000688 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +0000689 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
690 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000691 self.assertRaises(TypeError, starmap)
692 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000693 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
694 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
695 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000696
697 def test_islice(self):
698 for args in [ # islice(args) should agree with range(args)
699 (10, 20, 3),
700 (10, 3, 20),
701 (10, 20),
702 (10, 3),
703 (20,)
704 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000705 self.assertEqual(list(islice(range(100), *args)),
706 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000707
708 for args, tgtargs in [ # Stop when seqn is exhausted
709 ((10, 110, 3), ((10, 100, 3))),
710 ((10, 110), ((10, 100))),
711 ((110,), (100,))
712 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000713 self.assertEqual(list(islice(range(100), *args)),
714 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000715
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000716 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +0000717 self.assertEqual(list(islice(range(10), None)), list(range(10)))
718 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
719 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
720 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
721 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000722
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000723 # Test number of items consumed SF #1171417
724 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +0000725 self.assertEqual(list(islice(it, 3)), list(range(3)))
726 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000727
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000728 # Test invalid arguments
Guido van Rossum805365e2007-05-07 22:24:25 +0000729 self.assertRaises(TypeError, islice, range(10))
730 self.assertRaises(TypeError, islice, range(10), 1, 2, 3, 4)
731 self.assertRaises(ValueError, islice, range(10), -5, 10, 1)
732 self.assertRaises(ValueError, islice, range(10), 1, -5, -1)
733 self.assertRaises(ValueError, islice, range(10), 1, 10, -1)
734 self.assertRaises(ValueError, islice, range(10), 1, 10, 0)
735 self.assertRaises(ValueError, islice, range(10), 'a')
736 self.assertRaises(ValueError, islice, range(10), 'a', 1)
737 self.assertRaises(ValueError, islice, range(10), 1, 'a')
738 self.assertRaises(ValueError, islice, range(10), 'a', 1, 1)
739 self.assertRaises(ValueError, islice, range(10), 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +0000740 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000741
742 def test_takewhile(self):
743 data = [1, 3, 5, 20, 2, 4, 6, 8]
744 underten = lambda x: x<10
745 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000746 self.assertEqual(list(takewhile(underten, [])), [])
747 self.assertRaises(TypeError, takewhile)
748 self.assertRaises(TypeError, takewhile, operator.pow)
749 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000750 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
751 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000752 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
753 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +0000754 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000755
756 def test_dropwhile(self):
757 data = [1, 3, 5, 20, 2, 4, 6, 8]
758 underten = lambda x: x<10
759 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000760 self.assertEqual(list(dropwhile(underten, [])), [])
761 self.assertRaises(TypeError, dropwhile)
762 self.assertRaises(TypeError, dropwhile, operator.pow)
763 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000764 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
765 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000766
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000767 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000768 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000769 def irange(n):
Guido van Rossum805365e2007-05-07 22:24:25 +0000770 for i in range(n):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000771 yield i
772
773 a, b = tee([]) # test empty iterator
774 self.assertEqual(list(a), [])
775 self.assertEqual(list(b), [])
776
777 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +0000778 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000779
780 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +0000781 self.assertEqual(list(a), list(range(n)))
782 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000783
784 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000785 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000786 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000787 del a
Guido van Rossum805365e2007-05-07 22:24:25 +0000788 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000789
790 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000791 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000792 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000793 del b
Guido van Rossum805365e2007-05-07 22:24:25 +0000794 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000795
Guido van Rossum805365e2007-05-07 22:24:25 +0000796 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000797 order = [0]*n + [1]*n
798 random.shuffle(order)
799 lists = ([], [])
800 its = tee(irange(n))
801 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +0000802 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000803 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +0000804 self.assertEqual(lists[0], list(range(n)))
805 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000806
Raymond Hettingerad983e72003-11-12 14:32:26 +0000807 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000808 self.assertRaises(TypeError, tee)
809 self.assertRaises(TypeError, tee, 3)
810 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000811 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000812
Raymond Hettingerad983e72003-11-12 14:32:26 +0000813 # tee object should be instantiable
814 a, b = tee('abc')
815 c = type(a)('def')
816 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000817
Raymond Hettingerad983e72003-11-12 14:32:26 +0000818 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +0000819 a, b, c = tee(range(2000), 3)
820 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000821 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +0000822 self.assertEqual(list(b), list(range(2000)))
823 self.assertEqual([next(c), next(c)], list(range(2)))
824 self.assertEqual(list(a), list(range(100,2000)))
825 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000826
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000827 # test values of n
828 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000829 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +0000830 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000831 result = tee('abc', n)
832 self.assertEqual(type(result), tuple)
833 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000834 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000835
Raymond Hettingerad983e72003-11-12 14:32:26 +0000836 # tee pass-through to copyable iterator
837 a, b = tee('abc')
838 c, d = tee(a)
Georg Brandlab91fde2009-08-13 08:51:18 +0000839 self.assertTrue(a is c)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000840
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000841 # test tee_new
842 t1, t2 = tee('abc')
843 tnew = type(t1)
844 self.assertRaises(TypeError, tnew)
845 self.assertRaises(TypeError, tnew, 10)
846 t3 = tnew(t1)
Georg Brandlab91fde2009-08-13 08:51:18 +0000847 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000848
Raymond Hettingera9f60922004-10-17 16:40:14 +0000849 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +0000850 a, b = tee(range(10))
Raymond Hettingera9f60922004-10-17 16:40:14 +0000851 p = proxy(a)
852 self.assertEqual(getattr(p, '__class__'), type(b))
853 del a
854 self.assertRaises(ReferenceError, getattr, p, '__class__')
855
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000856 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000857 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000858
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000859 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +0000860 self.assertRaises(StopIteration, next, f([]))
861 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000862
Georg Brandla18af4e2007-04-21 15:47:16 +0000863 self.assertRaises(StopIteration, next, islice([], None))
864 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000865
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000866 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +0000867 self.assertRaises(StopIteration, next, p)
868 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000869 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +0000870 self.assertRaises(StopIteration, next, p)
871 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000872
Georg Brandla18af4e2007-04-21 15:47:16 +0000873 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000874
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000875 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +0000876 self.assertRaises(StopIteration, next, f(lambda x:x, []))
877 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000878
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000879class TestExamples(unittest.TestCase):
880
881 def test_chain(self):
882 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
883
884 def test_chain_from_iterable(self):
885 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
886
887 def test_combinations(self):
888 self.assertEqual(list(combinations('ABCD', 2)),
889 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
890 self.assertEqual(list(combinations(range(4), 3)),
891 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
892
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000893 def test_combinations_with_replacement(self):
894 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
895 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
896
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000897 def test_compress(self):
898 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
899
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000900 def test_count(self):
901 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
902
903 def test_cycle(self):
904 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
905
906 def test_dropwhile(self):
907 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
908
909 def test_groupby(self):
910 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
911 list('ABCDAB'))
912 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
913 [list('AAAA'), list('BBB'), list('CC'), list('D')])
914
915 def test_filter(self):
916 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
917
918 def test_filterfalse(self):
919 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
920
921 def test_map(self):
922 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
923
924 def test_islice(self):
925 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
926 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
927 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
928 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
929
930 def test_zip(self):
931 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
932
933 def test_zip_longest(self):
934 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
935 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
936
937 def test_permutations(self):
938 self.assertEqual(list(permutations('ABCD', 2)),
939 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
940 self.assertEqual(list(permutations(range(3))),
941 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
942
943 def test_product(self):
944 self.assertEqual(list(product('ABCD', 'xy')),
945 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
946 self.assertEqual(list(product(range(2), repeat=3)),
947 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
948 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
949
950 def test_repeat(self):
951 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
952
953 def test_stapmap(self):
954 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
955 [32, 9, 1000])
956
957 def test_takewhile(self):
958 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
959
960
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000961class TestGC(unittest.TestCase):
962
963 def makecycle(self, iterator, container):
964 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +0000965 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000966 del container, iterator
967
968 def test_chain(self):
969 a = []
970 self.makecycle(chain(a), a)
971
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000972 def test_chain_from_iterable(self):
973 a = []
974 self.makecycle(chain.from_iterable([a]), a)
975
976 def test_combinations(self):
977 a = []
978 self.makecycle(combinations([1,2,a,3], 3), a)
979
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000980 def test_combinations_with_replacement(self):
981 a = []
982 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
983
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000984 def test_compress(self):
985 a = []
986 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
987
Raymond Hettingerd6280f42009-02-16 20:50:56 +0000988 def test_count(self):
989 a = []
990 Int = type('Int', (int,), dict(x=a))
991 self.makecycle(count(Int(0), Int(1)), a)
992
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000993 def test_cycle(self):
994 a = []
995 self.makecycle(cycle([a]*2), a)
996
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000997 def test_dropwhile(self):
998 a = []
999 self.makecycle(dropwhile(bool, [0, a, a]), a)
1000
1001 def test_groupby(self):
1002 a = []
1003 self.makecycle(groupby([a]*2, lambda x:x), a)
1004
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001005 def test_issue2246(self):
1006 # Issue 2246 -- the _grouper iterator was not included in GC
1007 n = 10
1008 keyfunc = lambda x: x
1009 for i, j in groupby(range(n), key=keyfunc):
1010 keyfunc.__dict__.setdefault('x',[]).append(j)
1011
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001012 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001013 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001014 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001015
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001016 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001017 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001018 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001019
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001020 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001021 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001022 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001023
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001024 def test_zip_longest(self):
1025 a = []
1026 self.makecycle(zip_longest([a]*2, [a]*3), a)
1027 b = [a, None]
1028 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1029
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001030 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001031 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001032 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001033
1034 def test_islice(self):
1035 a = []
1036 self.makecycle(islice([a]*2, None), a)
1037
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001038 def test_permutations(self):
1039 a = []
1040 self.makecycle(permutations([1,2,a,3], 3), a)
1041
1042 def test_product(self):
1043 a = []
1044 self.makecycle(product([1,2,a,3], repeat=3), a)
1045
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001046 def test_repeat(self):
1047 a = []
1048 self.makecycle(repeat(a), a)
1049
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001050 def test_starmap(self):
1051 a = []
1052 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1053
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001054 def test_takewhile(self):
1055 a = []
1056 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1057
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001058def R(seqn):
1059 'Regular generator'
1060 for i in seqn:
1061 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001062
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001063class G:
1064 'Sequence using __getitem__'
1065 def __init__(self, seqn):
1066 self.seqn = seqn
1067 def __getitem__(self, i):
1068 return self.seqn[i]
1069
1070class I:
1071 'Sequence using iterator protocol'
1072 def __init__(self, seqn):
1073 self.seqn = seqn
1074 self.i = 0
1075 def __iter__(self):
1076 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001077 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001078 if self.i >= len(self.seqn): raise StopIteration
1079 v = self.seqn[self.i]
1080 self.i += 1
1081 return v
1082
1083class Ig:
1084 'Sequence using iterator protocol defined with a generator'
1085 def __init__(self, seqn):
1086 self.seqn = seqn
1087 self.i = 0
1088 def __iter__(self):
1089 for val in self.seqn:
1090 yield val
1091
1092class X:
1093 'Missing __getitem__ and __iter__'
1094 def __init__(self, seqn):
1095 self.seqn = seqn
1096 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001097 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001098 if self.i >= len(self.seqn): raise StopIteration
1099 v = self.seqn[self.i]
1100 self.i += 1
1101 return v
1102
1103class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001104 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001105 def __init__(self, seqn):
1106 self.seqn = seqn
1107 self.i = 0
1108 def __iter__(self):
1109 return self
1110
1111class E:
1112 'Test propagation of exceptions'
1113 def __init__(self, seqn):
1114 self.seqn = seqn
1115 self.i = 0
1116 def __iter__(self):
1117 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001118 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001119 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001120
1121class S:
1122 'Test immediate stop'
1123 def __init__(self, seqn):
1124 pass
1125 def __iter__(self):
1126 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001127 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001128 raise StopIteration
1129
1130def L(seqn):
1131 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001132 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001133
1134
1135class TestVariousIteratorArgs(unittest.TestCase):
1136
1137 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001138 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001139 for g in (G, I, Ig, S, L, R):
1140 self.assertEqual(list(chain(g(s))), list(g(s)))
1141 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001142 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001143 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001144 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1145
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001146 def test_compress(self):
1147 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1148 n = len(s)
1149 for g in (G, I, Ig, S, L, R):
1150 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1151 self.assertRaises(TypeError, compress, X(s), repeat(1))
1152 self.assertRaises(TypeError, compress, N(s), repeat(1))
1153 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1154
Christian Heimesc3f30c42008-02-22 16:37:40 +00001155 def test_product(self):
1156 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1157 self.assertRaises(TypeError, product, X(s))
1158 self.assertRaises(TypeError, product, N(s))
1159 self.assertRaises(ZeroDivisionError, product, E(s))
1160
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001161 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001162 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001163 for g in (G, I, Ig, S, L, R):
1164 tgtlen = len(s) * 3
1165 expected = list(g(s))*3
1166 actual = list(islice(cycle(g(s)), tgtlen))
1167 self.assertEqual(actual, expected)
1168 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001169 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001170 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1171
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001172 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001173 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001174 for g in (G, I, Ig, S, L, R):
1175 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1176 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001177 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001178 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1179
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001180 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001181 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001182 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001183 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001184 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001185 self.assertRaises(TypeError, filter, isEven, X(s))
1186 self.assertRaises(TypeError, filter, isEven, N(s))
1187 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001188
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001189 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001190 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001191 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001192 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001193 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001194 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1195 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1196 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001197
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001198 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001199 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001200 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001201 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1202 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1203 self.assertRaises(TypeError, zip, X(s))
1204 self.assertRaises(TypeError, zip, N(s))
1205 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001206
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001207 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001208 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001209 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001210 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1211 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1212 self.assertRaises(TypeError, zip_longest, X(s))
1213 self.assertRaises(TypeError, zip_longest, N(s))
1214 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001215
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001216 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001217 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001218 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001219 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001220 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001221 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001222 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001223 self.assertRaises(TypeError, map, onearg, X(s))
1224 self.assertRaises(TypeError, map, onearg, N(s))
1225 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001226
1227 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001228 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001229 for g in (G, I, Ig, S, L, R):
1230 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1231 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001232 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001233 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1234
1235 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001236 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001237 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001238 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001239 self.assertEqual(list(starmap(operator.pow, g(ss))),
1240 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001241 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001242 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001243 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1244
1245 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001246 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001247 for g in (G, I, Ig, S, L, R):
1248 tgt = []
1249 for elem in g(s):
1250 if not isEven(elem): break
1251 tgt.append(elem)
1252 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1253 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001254 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001255 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1256
1257 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001258 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001259 for g in (G, I, Ig, S, L, R):
1260 tgt = []
1261 for elem in g(s):
1262 if not tgt and isOdd(elem): continue
1263 tgt.append(elem)
1264 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1265 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001266 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001267 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1268
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001269 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001270 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001271 for g in (G, I, Ig, S, L, R):
1272 it1, it2 = tee(g(s))
1273 self.assertEqual(list(it1), list(g(s)))
1274 self.assertEqual(list(it2), list(g(s)))
1275 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001276 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001277 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1278
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001279class LengthTransparency(unittest.TestCase):
1280
1281 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001282 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001283 self.assertEqual(len(repeat(None, 50)), 50)
1284 self.assertRaises(TypeError, len, repeat(None))
1285
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001286class RegressionTests(unittest.TestCase):
1287
1288 def test_sf_793826(self):
1289 # Fix Armin Rigo's successful efforts to wreak havoc
1290
1291 def mutatingtuple(tuple1, f, tuple2):
1292 # this builds a tuple t which is a copy of tuple1,
1293 # then calls f(t), then mutates t to be equal to tuple2
1294 # (needs len(tuple1) == len(tuple2)).
1295 def g(value, first=[1]):
1296 if first:
1297 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00001298 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001299 return value
1300 items = list(tuple2)
1301 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001302 gen = map(g, items)
1303 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00001304 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001305
1306 def f(t):
1307 global T
1308 T = t
1309 first[:] = list(T)
1310
1311 first = []
1312 mutatingtuple((1,2,3), f, (4,5,6))
1313 second = list(T)
1314 self.assertEqual(first, second)
1315
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001316
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001317 def test_sf_950057(self):
1318 # Make sure that chain() and cycle() catch exceptions immediately
1319 # rather than when shifting between input sources
1320
1321 def gen1():
1322 hist.append(0)
1323 yield 1
1324 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001325 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001326 hist.append(2)
1327
1328 def gen2(x):
1329 hist.append(3)
1330 yield 2
1331 hist.append(4)
1332 if x:
1333 raise StopIteration
1334
1335 hist = []
1336 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1337 self.assertEqual(hist, [0,1])
1338
1339 hist = []
1340 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1341 self.assertEqual(hist, [0,1])
1342
1343 hist = []
1344 self.assertRaises(AssertionError, list, cycle(gen1()))
1345 self.assertEqual(hist, [0,1])
1346
Thomas Woutersb2137042007-02-01 18:02:27 +00001347class SubclassWithKwargsTest(unittest.TestCase):
1348 def test_keywords_in_subclass(self):
1349 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001350 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001351 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00001352 class Subclass(cls):
1353 def __init__(self, newarg=None, *args):
1354 cls.__init__(self, *args)
1355 try:
1356 Subclass(newarg=1)
1357 except TypeError as err:
1358 # we expect type errors because of wrong argument count
Georg Brandlab91fde2009-08-13 08:51:18 +00001359 self.assertFalse("does not take keyword arguments" in err.args[0])
Thomas Woutersb2137042007-02-01 18:02:27 +00001360
1361
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001362libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001363
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001364
1365>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001366>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00001367... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001368...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001369Check 1200 is for $120.15
1370Check 1201 is for $764.05
1371Check 1202 is for $823.14
1372
1373>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001374>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00001375... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00001376...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000013771
13788
137927
1380
1381>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001382>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00001383... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00001384...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001385Alex
1386Laura
1387Martin
1388Walter
1389Samuele
1390
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001391>>> from operator import itemgetter
1392>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00001393>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001394>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001395... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001396...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000013971 ['a', 'c', 'e']
13982 ['b', 'd', 'f']
13993 ['g']
1400
Raymond Hettinger734fb572004-01-20 20:04:40 +00001401# Find runs of consecutive numbers using groupby. The key to the solution
1402# is differencing with a range so that consecutive numbers all appear in
1403# same group.
1404>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00001405>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001406... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001407...
Raymond Hettinger734fb572004-01-20 20:04:40 +00001408[1]
1409[4, 5, 6]
1410[10]
1411[15, 16, 17, 18]
1412[22]
1413[25, 26, 27, 28]
1414
Georg Brandl3dbca812008-07-23 16:10:53 +00001415>>> def take(n, iterable):
1416... "Return first n items of the iterable as a list"
1417... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001418
Georg Brandl3dbca812008-07-23 16:10:53 +00001419>>> def enumerate(iterable, start=0):
1420... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001421
Georg Brandl3dbca812008-07-23 16:10:53 +00001422>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001423... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00001424... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001425
Raymond Hettingercdf8ba32009-02-19 04:45:07 +00001426>>> def nth(iterable, n, default=None):
1427... "Returns the nth item or a default value"
1428... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001429
Georg Brandl3dbca812008-07-23 16:10:53 +00001430>>> def quantify(iterable, pred=bool):
1431... "Count how many times the predicate is true"
1432... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001433
Georg Brandl3dbca812008-07-23 16:10:53 +00001434>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001435... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00001436... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001437
Georg Brandl3dbca812008-07-23 16:10:53 +00001438>>> def ncycles(iterable, n):
1439... "Returns the seqeuence elements n times"
1440... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001441
1442>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001443... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001444
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001445>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001446... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001447
1448>>> def repeatfunc(func, times=None, *args):
1449... "Repeat calls to func with specified arguments."
1450... " Example: repeatfunc(random.random)"
1451... if times is None:
1452... return starmap(func, repeat(args))
1453... else:
1454... return starmap(func, repeat(args, times))
1455
Raymond Hettingerd591f662003-10-26 15:34:50 +00001456>>> def pairwise(iterable):
1457... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1458... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001459... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00001460... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001461... except StopIteration:
1462... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001463... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001464
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001465>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001466... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001467... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00001468... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001469
1470>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001471... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001472... # Recipe credited to George Sakkis
1473... pending = len(iterables)
1474... nexts = cycle(iter(it).__next__ for it in iterables)
1475... while pending:
1476... try:
1477... for next in nexts:
1478... yield next()
1479... except StopIteration:
1480... pending -= 1
1481... nexts = cycle(islice(nexts, pending))
1482
1483>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00001484... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1485... s = list(iterable)
1486... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001487
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001488>>> def unique_everseen(iterable, key=None):
1489... "List unique elements, preserving order. Remember all elements ever seen."
1490... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1491... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1492... seen = set()
1493... seen_add = seen.add
1494... if key is None:
1495... for element in iterable:
1496... if element not in seen:
1497... seen_add(element)
1498... yield element
1499... else:
1500... for element in iterable:
1501... k = key(element)
1502... if k not in seen:
1503... seen_add(k)
1504... yield element
1505
1506>>> def unique_justseen(iterable, key=None):
1507... "List unique elements, preserving order. Remember only the element just seen."
1508... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1509... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1510... return map(next, map(itemgetter(1), groupby(iterable, key)))
1511
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001512This is not part of the examples but it tests to make sure the definitions
1513perform as purported.
1514
Raymond Hettingera098b332003-09-08 23:58:40 +00001515>>> take(10, count())
1516[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1517
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001518>>> list(enumerate('abc'))
1519[(0, 'a'), (1, 'b'), (2, 'c')]
1520
1521>>> list(islice(tabulate(lambda x: 2*x), 4))
1522[0, 2, 4, 6]
1523
1524>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00001525'd'
1526
1527>>> nth('abcde', 9) is None
1528True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001529
Guido van Rossum805365e2007-05-07 22:24:25 +00001530>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000153150
1532
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001533>>> a = [[1, 2, 3], [4, 5, 6]]
1534>>> flatten(a)
1535[1, 2, 3, 4, 5, 6]
1536
1537>>> list(repeatfunc(pow, 5, 2, 3))
1538[8, 8, 8, 8, 8]
1539
1540>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001541>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001542[0, 0, 0, 0, 0]
1543
Raymond Hettingerd591f662003-10-26 15:34:50 +00001544>>> list(pairwise('abcd'))
1545[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001546
Raymond Hettingerd591f662003-10-26 15:34:50 +00001547>>> list(pairwise([]))
1548[]
1549
1550>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001551[]
1552
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001553>>> list(islice(padnone('abc'), 0, 6))
1554['a', 'b', 'c', None, None, None]
1555
1556>>> list(ncycles('abc', 3))
1557['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1558
1559>>> dotproduct([1,2,3], [4,5,6])
156032
1561
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001562>>> list(grouper(3, 'abcdefg', 'x'))
1563[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1564
1565>>> list(roundrobin('abc', 'd', 'ef'))
1566['a', 'd', 'e', 'b', 'f', 'c']
1567
Raymond Hettingerace67332009-01-26 02:23:50 +00001568>>> list(powerset([1,2,3]))
1569[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001570
Raymond Hettinger191e8502009-01-27 13:29:43 +00001571>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1572True
1573
1574>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1575True
1576
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001577>>> list(unique_everseen('AAAABBBCCDAABBB'))
1578['A', 'B', 'C', 'D']
1579
1580>>> list(unique_everseen('ABBCcAD', str.lower))
1581['A', 'B', 'C', 'D']
1582
1583>>> list(unique_justseen('AAAABBBCCDAABBB'))
1584['A', 'B', 'C', 'D', 'A', 'B']
1585
1586>>> list(unique_justseen('ABBCcAD', str.lower))
1587['A', 'B', 'C', 'A', 'D']
1588
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001589"""
1590
1591__test__ = {'libreftest' : libreftest}
1592
1593def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001594 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00001595 RegressionTests, LengthTransparency,
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001596 SubclassWithKwargsTest, TestExamples)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001597 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001598
1599 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001600 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001601 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001602 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00001603 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001604 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001605 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001606 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001607 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001608
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001609 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001610 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001611
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001612if __name__ == "__main__":
1613 test_main(verbose=True)