blob: 8a67cff60ce9db32e7c531b17227bdb6d5bfd0b8 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001import unittest
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002from test import support
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003from itertools import *
Raymond Hettingera9f60922004-10-17 16:40:14 +00004from weakref import proxy
Raymond Hettinger8a882a72009-02-12 12:08:18 +00005from decimal import Decimal
Raymond Hettingerde09acf2009-02-12 12:53:53 +00006from fractions import Fraction
Raymond Hettinger2012f172003-02-07 05:32:58 +00007import sys
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00008import operator
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00009import random
Raymond Hettingera3e1ad22009-11-30 22:02:31 +000010import copy
11import pickle
Christian Heimesc3f30c42008-02-22 16:37:40 +000012from functools import reduce
Benjamin Petersonee8712c2008-05-20 21:35:26 +000013maxsize = support.MAX_Py_ssize_t
Guido van Rossum360e4b82007-05-14 22:51:27 +000014minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000015
Guido van Rossum801f0d72006-08-24 19:48:10 +000016def lzip(*args):
17 return list(zip(*args))
18
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000019def onearg(x):
20 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000021 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000022
23def errfunc(*args):
24 'Test function that raises an error'
25 raise ValueError
26
27def gen3():
28 'Non-restartable source sequence'
29 for i in (0, 1, 2):
30 yield i
31
32def isEven(x):
33 'Test predicate'
34 return x%2==0
35
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000036def isOdd(x):
37 'Test predicate'
38 return x%2==1
39
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000040class StopNow:
41 'Class emulating an empty iterable.'
42 def __iter__(self):
43 return self
Georg Brandla18af4e2007-04-21 15:47:16 +000044 def __next__(self):
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000045 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000046
Raymond Hettinger02420702003-06-29 20:36:23 +000047def take(n, seq):
48 'Convenience function for partially consuming a long of infinite iterable'
49 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000050
Christian Heimes78644762008-03-04 23:39:23 +000051def prod(iterable):
52 return reduce(operator.mul, iterable, 1)
53
Christian Heimes380f7f22008-02-28 11:19:05 +000054def fact(n):
55 'Factorial'
Christian Heimes78644762008-03-04 23:39:23 +000056 return prod(range(1, n+1))
57
Raymond Hettinger96ef8112003-02-01 00:10:11 +000058class TestBasicOps(unittest.TestCase):
Raymond Hettinger482ba772010-12-01 22:48:00 +000059
60 def test_accumulate(self):
61 self.assertEqual(list(accumulate(range(10))), # one positional arg
62 [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])
63 self.assertEqual(list(accumulate(range(10), 100)), # two positional args
64 [100, 101, 103, 106, 110, 115, 121, 128, 136, 145])
65 self.assertEqual(list(accumulate(iterable=range(10), start=100)), # kw args
66 [100, 101, 103, 106, 110, 115, 121, 128, 136, 145])
67 for typ in int, complex, Decimal, Fraction: # multiple types
68 self.assertEqual(list(accumulate(range(10), typ(0))),
69 list(map(typ, [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])))
70 self.assertEqual(list(accumulate([])), []) # empty iterable
71 self.assertRaises(TypeError, accumulate, range(10), 0, 5) # too many args
72 self.assertRaises(TypeError, accumulate) # too few args
73 self.assertRaises(TypeError, accumulate, range(10), x=7) # unexpected kwd args
74 self.assertRaises(TypeError, list, accumulate([1, []])) # args that don't add
75
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000076 def test_chain(self):
Christian Heimesdd15f6c2008-03-16 00:07:10 +000077
78 def chain2(*iterables):
79 'Pure python version in the docs'
80 for it in iterables:
81 for element in it:
82 yield element
83
84 for c in (chain, chain2):
85 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
86 self.assertEqual(list(c('abc')), list('abc'))
87 self.assertEqual(list(c('')), [])
88 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
89 self.assertRaises(TypeError, list,c(2, 3))
Christian Heimesf16baeb2008-02-29 14:57:44 +000090
91 def test_chain_from_iterable(self):
92 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
93 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
94 self.assertEqual(list(chain.from_iterable([''])), [])
95 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
96 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000097
Christian Heimes380f7f22008-02-28 11:19:05 +000098 def test_combinations(self):
Raymond Hettinger5bad41e2009-01-08 21:01:54 +000099 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Christian Heimes380f7f22008-02-28 11:19:05 +0000100 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Christian Heimes78644762008-03-04 23:39:23 +0000101 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Christian Heimes380f7f22008-02-28 11:19:05 +0000102 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000103 self.assertEqual(list(combinations('abc', 32)), []) # r > n
Christian Heimes380f7f22008-02-28 11:19:05 +0000104 self.assertEqual(list(combinations(range(4), 3)),
105 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
Christian Heimes78644762008-03-04 23:39:23 +0000106
107 def combinations1(iterable, r):
108 'Pure python version shown in the docs'
109 pool = tuple(iterable)
110 n = len(pool)
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000111 if r > n:
112 return
Christian Heimes78644762008-03-04 23:39:23 +0000113 indices = list(range(r))
114 yield tuple(pool[i] for i in indices)
115 while 1:
116 for i in reversed(range(r)):
117 if indices[i] != i + n - r:
118 break
119 else:
120 return
121 indices[i] += 1
122 for j in range(i+1, r):
123 indices[j] = indices[j-1] + 1
124 yield tuple(pool[i] for i in indices)
125
126 def combinations2(iterable, r):
127 'Pure python version shown in the docs'
128 pool = tuple(iterable)
129 n = len(pool)
130 for indices in permutations(range(n), r):
131 if sorted(indices) == list(indices):
132 yield tuple(pool[i] for i in indices)
133
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000134 def combinations3(iterable, r):
135 'Pure python version from cwr()'
136 pool = tuple(iterable)
137 n = len(pool)
138 for indices in combinations_with_replacement(range(n), r):
139 if len(set(indices)) == r:
140 yield tuple(pool[i] for i in indices)
141
Christian Heimes78644762008-03-04 23:39:23 +0000142 for n in range(7):
Christian Heimes380f7f22008-02-28 11:19:05 +0000143 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000144 for r in range(n+2):
Christian Heimes380f7f22008-02-28 11:19:05 +0000145 result = list(combinations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000146 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 +0000147 self.assertEqual(len(result), len(set(result))) # no repeats
148 self.assertEqual(result, sorted(result)) # lexicographic order
149 for c in result:
150 self.assertEqual(len(c), r) # r-length combinations
151 self.assertEqual(len(set(c)), r) # no duplicate elements
152 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000153 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000154 self.assertEqual(list(c),
155 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000156 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000157 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000158 self.assertEqual(result, list(combinations3(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000159
160 # Test implementation detail: tuple re-use
161 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
162 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
163
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000164 def test_combinations_with_replacement(self):
165 cwr = combinations_with_replacement
166 self.assertRaises(TypeError, cwr, 'abc') # missing r argument
167 self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
168 self.assertRaises(TypeError, cwr, None) # pool is not iterable
169 self.assertRaises(ValueError, cwr, 'abc', -2) # r is negative
170 self.assertEqual(list(cwr('ABC', 2)),
171 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
172
173 def cwr1(iterable, r):
174 'Pure python version shown in the docs'
175 # number items returned: (n+r-1)! / r! / (n-1)! when n>0
176 pool = tuple(iterable)
177 n = len(pool)
178 if not n and r:
179 return
180 indices = [0] * r
181 yield tuple(pool[i] for i in indices)
182 while 1:
183 for i in reversed(range(r)):
184 if indices[i] != n - 1:
185 break
186 else:
187 return
188 indices[i:] = [indices[i] + 1] * (r - i)
189 yield tuple(pool[i] for i in indices)
190
191 def cwr2(iterable, r):
192 'Pure python version shown in the docs'
193 pool = tuple(iterable)
194 n = len(pool)
195 for indices in product(range(n), repeat=r):
196 if sorted(indices) == list(indices):
197 yield tuple(pool[i] for i in indices)
198
199 def numcombs(n, r):
200 if not n:
201 return 0 if r else 1
202 return fact(n+r-1) / fact(r)/ fact(n-1)
203
204 for n in range(7):
205 values = [5*x-12 for x in range(n)]
206 for r in range(n+2):
207 result = list(cwr(values, r))
208
209 self.assertEqual(len(result), numcombs(n, r)) # right number of combs
210 self.assertEqual(len(result), len(set(result))) # no repeats
211 self.assertEqual(result, sorted(result)) # lexicographic order
212
213 regular_combs = list(combinations(values, r)) # compare to combs without replacement
214 if n == 0 or r <= 1:
Ezio Melottib3aedd42010-11-20 19:04:17 +0000215 self.assertEqual(result, regular_combs) # cases that should be identical
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000216 else:
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000217 self.assertTrue(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000218
219 for c in result:
220 self.assertEqual(len(c), r) # r-length combinations
221 noruns = [k for k,v in groupby(c)] # combo without consecutive repeats
222 self.assertEqual(len(noruns), len(set(noruns))) # no repeats other than consecutive
223 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000224 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000225 self.assertEqual(noruns,
226 [e for e in values if e in c]) # comb is a subsequence of the input iterable
227 self.assertEqual(result, list(cwr1(values, r))) # matches first pure python version
228 self.assertEqual(result, list(cwr2(values, r))) # matches second pure python version
229
230 # Test implementation detail: tuple re-use
231 self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
232 self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
233
Christian Heimes78644762008-03-04 23:39:23 +0000234 def test_permutations(self):
235 self.assertRaises(TypeError, permutations) # too few arguments
236 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000237 self.assertRaises(TypeError, permutations, None) # pool is not iterable
238 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000239 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000240 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Christian Heimes78644762008-03-04 23:39:23 +0000241 self.assertEqual(list(permutations(range(3), 2)),
242 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
243
244 def permutations1(iterable, r=None):
245 'Pure python version shown in the docs'
246 pool = tuple(iterable)
247 n = len(pool)
248 r = n if r is None else r
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000249 if r > n:
250 return
Christian Heimes78644762008-03-04 23:39:23 +0000251 indices = list(range(n))
252 cycles = list(range(n-r+1, n+1))[::-1]
253 yield tuple(pool[i] for i in indices[:r])
254 while n:
255 for i in reversed(range(r)):
256 cycles[i] -= 1
257 if cycles[i] == 0:
258 indices[i:] = indices[i+1:] + indices[i:i+1]
259 cycles[i] = n - i
260 else:
261 j = cycles[i]
262 indices[i], indices[-j] = indices[-j], indices[i]
263 yield tuple(pool[i] for i in indices[:r])
264 break
265 else:
266 return
267
268 def permutations2(iterable, r=None):
269 'Pure python version shown in the docs'
270 pool = tuple(iterable)
271 n = len(pool)
272 r = n if r is None else r
273 for indices in product(range(n), repeat=r):
274 if len(set(indices)) == r:
275 yield tuple(pool[i] for i in indices)
276
277 for n in range(7):
278 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000279 for r in range(n+2):
Christian Heimes78644762008-03-04 23:39:23 +0000280 result = list(permutations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000281 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 +0000282 self.assertEqual(len(result), len(set(result))) # no repeats
283 self.assertEqual(result, sorted(result)) # lexicographic order
284 for p in result:
285 self.assertEqual(len(p), r) # r-length permutations
286 self.assertEqual(len(set(p)), r) # no duplicate elements
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000287 self.assertTrue(all(e in values for e in p)) # elements taken from input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000288 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000289 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000290 if r == n:
291 self.assertEqual(result, list(permutations(values, None))) # test r as None
292 self.assertEqual(result, list(permutations(values))) # test default r
293
294 # Test implementation detail: tuple re-use
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000295 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Christian Heimes78644762008-03-04 23:39:23 +0000296 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Christian Heimes380f7f22008-02-28 11:19:05 +0000297
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000298 def test_combinatorics(self):
299 # Test relationships between product(), permutations(),
300 # combinations() and combinations_with_replacement().
301
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000302 for n in range(6):
303 s = 'ABCDEFG'[:n]
304 for r in range(8):
305 prod = list(product(s, repeat=r))
306 cwr = list(combinations_with_replacement(s, r))
307 perm = list(permutations(s, r))
308 comb = list(combinations(s, r))
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000309
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000310 # Check size
Ezio Melottib3aedd42010-11-20 19:04:17 +0000311 self.assertEqual(len(prod), n**r)
312 self.assertEqual(len(cwr), (fact(n+r-1) / fact(r)/ fact(n-1)) if n else (not r))
313 self.assertEqual(len(perm), 0 if r>n else fact(n) / fact(n-r))
314 self.assertEqual(len(comb), 0 if r>n else fact(n) / fact(r) / fact(n-r))
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000315
316 # Check lexicographic order without repeated tuples
Ezio Melottib3aedd42010-11-20 19:04:17 +0000317 self.assertEqual(prod, sorted(set(prod)))
318 self.assertEqual(cwr, sorted(set(cwr)))
319 self.assertEqual(perm, sorted(set(perm)))
320 self.assertEqual(comb, sorted(set(comb)))
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000321
322 # Check interrelationships
Ezio Melottib3aedd42010-11-20 19:04:17 +0000323 self.assertEqual(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
324 self.assertEqual(perm, [t for t in prod if len(set(t))==r]) # perm: prods with no dups
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000325 self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
326 self.assertEqual(comb, [t for t in cwr if len(set(t))==r]) # comb: cwrs without dups
327 self.assertEqual(comb, list(filter(set(cwr).__contains__, perm))) # comb: perm that is a cwr
328 self.assertEqual(comb, list(filter(set(perm).__contains__, cwr))) # comb: cwr that is a perm
329 self.assertEqual(comb, sorted(set(cwr) & set(perm))) # comb: both a cwr and a perm
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000330
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000331 def test_compress(self):
Raymond Hettinger15a49502009-02-19 02:17:09 +0000332 self.assertEqual(list(compress(data='ABCDEF', selectors=[1,0,1,0,1,1])), list('ACEF'))
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000333 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
334 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
335 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
336 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
337 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
338 n = 10000
339 data = chain.from_iterable(repeat(range(6), n))
340 selectors = chain.from_iterable(repeat((0, 1)))
341 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
342 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
343 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
344 self.assertRaises(TypeError, compress, range(6)) # too few args
345 self.assertRaises(TypeError, compress, range(6), None) # too many args
346
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000347 def test_count(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000348 self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
349 self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
350 self.assertEqual(take(2, lzip('abc',count(3))), [('a', 3), ('b', 4)])
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000351 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
352 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000353 self.assertRaises(TypeError, count, 2, 3, 4)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000354 self.assertRaises(TypeError, count, 'a')
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000355 self.assertEqual(list(islice(count(maxsize-5), 10)),
356 list(range(maxsize-5, maxsize+5)))
357 self.assertEqual(list(islice(count(-maxsize-5), 10)),
358 list(range(-maxsize-5, -maxsize+5)))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000359 c = count(3)
360 self.assertEqual(repr(c), 'count(3)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000361 next(c)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000362 self.assertEqual(repr(c), 'count(4)')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000363 c = count(-9)
364 self.assertEqual(repr(c), 'count(-9)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000365 next(c)
Raymond Hettinger30729212009-02-12 06:28:27 +0000366 self.assertEqual(repr(count(10.25)), 'count(10.25)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000367 self.assertEqual(next(c), -8)
Christian Heimesa37d4c62007-12-04 23:02:19 +0000368 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 +0000369 # Test repr (ignoring the L in longs)
370 r1 = repr(count(i)).replace('L', '')
371 r2 = 'count(%r)'.__mod__(i).replace('L', '')
372 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000373
Raymond Hettingera3e1ad22009-11-30 22:02:31 +0000374 # check copy, deepcopy, pickle
375 for value in -3, 3, maxsize-5, maxsize+5:
376 c = count(value)
377 self.assertEqual(next(copy.copy(c)), value)
378 self.assertEqual(next(copy.deepcopy(c)), value)
379 self.assertEqual(next(pickle.loads(pickle.dumps(c))), value)
380
Raymond Hettinger30729212009-02-12 06:28:27 +0000381 def test_count_with_stride(self):
382 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingereb13fdd2009-02-21 22:30:12 +0000383 self.assertEqual(lzip('abc',count(start=2,step=3)),
384 [('a', 2), ('b', 5), ('c', 8)])
385 self.assertEqual(lzip('abc',count(step=-1)),
386 [('a', 0), ('b', -1), ('c', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000387 self.assertEqual(lzip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
388 self.assertEqual(lzip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000389 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000390 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
391 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
392 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettinger8a882a72009-02-12 12:08:18 +0000393 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
394 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerde09acf2009-02-12 12:53:53 +0000395 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
396 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000397 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
398 c = count(3, 5)
399 self.assertEqual(repr(c), 'count(3, 5)')
400 next(c)
401 self.assertEqual(repr(c), 'count(8, 5)')
402 c = count(-9, 0)
403 self.assertEqual(repr(c), 'count(-9, 0)')
404 next(c)
405 self.assertEqual(repr(c), 'count(-9, 0)')
406 c = count(-9, -3)
407 self.assertEqual(repr(c), 'count(-9, -3)')
408 next(c)
409 self.assertEqual(repr(c), 'count(-12, -3)')
410 self.assertEqual(repr(c), 'count(-12, -3)')
411 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
412 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
413 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
414 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
415 for j in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 1, 10, sys.maxsize-5, sys.maxsize+5):
416 # Test repr (ignoring the L in longs)
417 r1 = repr(count(i, j)).replace('L', '')
418 if j == 1:
419 r2 = ('count(%r)' % i).replace('L', '')
420 else:
421 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
422 self.assertEqual(r1, r2)
423
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000424 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000425 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000426 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000427 self.assertRaises(TypeError, cycle)
428 self.assertRaises(TypeError, cycle, 5)
429 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000430
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000431 def test_groupby(self):
432 # Check whether it accepts arguments correctly
433 self.assertEqual([], list(groupby([])))
434 self.assertEqual([], list(groupby([], key=id)))
435 self.assertRaises(TypeError, list, groupby('abc', []))
436 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000437 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000438
439 # Check normal input
440 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
441 (2,15,22), (3,16,23), (3,17,23)]
442 dup = []
443 for k, g in groupby(s, lambda r:r[0]):
444 for elem in g:
445 self.assertEqual(k, elem[0])
446 dup.append(elem)
447 self.assertEqual(s, dup)
448
449 # Check nested case
450 dup = []
451 for k, g in groupby(s, lambda r:r[0]):
452 for ik, ig in groupby(g, lambda r:r[2]):
453 for elem in ig:
454 self.assertEqual(k, elem[0])
455 self.assertEqual(ik, elem[2])
456 dup.append(elem)
457 self.assertEqual(s, dup)
458
459 # Check case where inner iterator is not used
460 keys = [k for k, g in groupby(s, lambda r:r[0])]
461 expectedkeys = set([r[0] for r in s])
462 self.assertEqual(set(keys), expectedkeys)
463 self.assertEqual(len(keys), len(expectedkeys))
464
465 # Exercise pipes and filters style
466 s = 'abracadabra'
467 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000468 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000469 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
470 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000471 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000472 self.assertEqual(r, ['a', 'b', 'r'])
473 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000474 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000475 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
476 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000477 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000478 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
479
Georg Brandla18af4e2007-04-21 15:47:16 +0000480 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000481 class ExpectedError(Exception):
482 pass
483 def delayed_raise(n=0):
484 for i in range(n):
485 yield 'yo'
486 raise ExpectedError
487 def gulp(iterable, keyp=None, func=list):
488 return [func(g) for k, g in groupby(iterable, keyp)]
489
Georg Brandla18af4e2007-04-21 15:47:16 +0000490 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000491 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000492 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000493 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
494
495 # __cmp__ failure
496 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000497 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000498 raise ExpectedError
499 s = [DummyCmp(), DummyCmp(), None]
500
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000501 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000502 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000503 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000504 self.assertRaises(ExpectedError, gulp, s)
505
506 # keyfunc failure
507 def keyfunc(obj):
508 if keyfunc.skip > 0:
509 keyfunc.skip -= 1
510 return obj
511 else:
512 raise ExpectedError
513
514 # keyfunc failure on outer object
515 keyfunc.skip = 0
516 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
517 keyfunc.skip = 1
518 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
519
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000520 def test_filter(self):
521 self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
522 self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])
523 self.assertEqual(list(filter(bool, [0,1,0,2,0])), [1,2])
524 self.assertEqual(take(4, filter(isEven, count())), [0,2,4,6])
525 self.assertRaises(TypeError, filter)
526 self.assertRaises(TypeError, filter, lambda x:x)
527 self.assertRaises(TypeError, filter, lambda x:x, range(6), 7)
528 self.assertRaises(TypeError, filter, isEven, 3)
529 self.assertRaises(TypeError, next, filter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000530
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000531 def test_filterfalse(self):
532 self.assertEqual(list(filterfalse(isEven, range(6))), [1,3,5])
533 self.assertEqual(list(filterfalse(None, [0,1,0,2,0])), [0,0,0])
534 self.assertEqual(list(filterfalse(bool, [0,1,0,2,0])), [0,0,0])
535 self.assertEqual(take(4, filterfalse(isEven, count())), [1,3,5,7])
536 self.assertRaises(TypeError, filterfalse)
537 self.assertRaises(TypeError, filterfalse, lambda x:x)
538 self.assertRaises(TypeError, filterfalse, lambda x:x, range(6), 7)
539 self.assertRaises(TypeError, filterfalse, isEven, 3)
540 self.assertRaises(TypeError, next, filterfalse(range(6), range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000541
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000542 def test_zip(self):
543 # XXX This is rather silly now that builtin zip() calls zip()...
544 ans = [(x,y) for x, y in zip('abc',count())]
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000545 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000546 self.assertEqual(list(zip('abc', range(6))), lzip('abc', range(6)))
547 self.assertEqual(list(zip('abcdef', range(3))), lzip('abcdef', range(3)))
548 self.assertEqual(take(3,zip('abcdef', count())), lzip('abcdef', range(3)))
549 self.assertEqual(list(zip('abcdef')), lzip('abcdef'))
550 self.assertEqual(list(zip()), lzip())
551 self.assertRaises(TypeError, zip, 3)
552 self.assertRaises(TypeError, zip, range(3), 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000553 # Check tuple re-use (implementation detail)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000554 self.assertEqual([tuple(list(pair)) for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000555 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000556 self.assertEqual([pair for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000557 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000558 ids = list(map(id, zip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000559 self.assertEqual(min(ids), max(ids))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000560 ids = list(map(id, list(zip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000561 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000562
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000563 def test_ziplongest(self):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000564 for args in [
565 ['abc', range(6)],
566 [range(6), 'abc'],
567 [range(1000), range(2000,2100), range(3000,3050)],
568 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
569 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
570 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000571 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
572 for i in range(max(map(len, args)))]
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000573 self.assertEqual(list(zip_longest(*args)), target)
574 self.assertEqual(list(zip_longest(*args, **{})), target)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000575 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 +0000576 self.assertEqual(list(zip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000577
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000578 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 +0000579
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000580 self.assertEqual(list(zip_longest()), list(zip()))
581 self.assertEqual(list(zip_longest([])), list(zip([])))
582 self.assertEqual(list(zip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000583
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000584 self.assertEqual(list(zip_longest('abc', 'defg', **{})),
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000585 list(zip(list('abc')+[None], 'defg'))) # empty keyword dict
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000586 self.assertRaises(TypeError, zip_longest, 3)
587 self.assertRaises(TypeError, zip_longest, range(3), 3)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000588
589 for stmt in [
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000590 "zip_longest('abc', fv=1)",
591 "zip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000592 ]:
593 try:
594 eval(stmt, globals(), locals())
595 except TypeError:
596 pass
597 else:
598 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000599
Thomas Wouterscf297e42007-02-23 15:07:44 +0000600 # Check tuple re-use (implementation detail)
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000601 self.assertEqual([tuple(list(pair)) for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000602 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000603 self.assertEqual([pair for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000604 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000605 ids = list(map(id, zip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000606 self.assertEqual(min(ids), max(ids))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000607 ids = list(map(id, list(zip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000608 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
609
Raymond Hettingerfc438512009-11-01 20:55:33 +0000610 def test_bug_7244(self):
611
612 class Repeater:
613 # this class is similar to itertools.repeat
614 def __init__(self, o, t, e):
615 self.o = o
616 self.t = int(t)
617 self.e = e
618 def __iter__(self): # its iterator is itself
619 return self
620 def __next__(self):
621 if self.t > 0:
622 self.t -= 1
623 return self.o
624 else:
625 raise self.e
626
627 # Formerly this code in would fail in debug mode
628 # with Undetected Error and Stop Iteration
629 r1 = Repeater(1, 3, StopIteration)
630 r2 = Repeater(2, 4, StopIteration)
631 def run(r1, r2):
632 result = []
633 for i, j in zip_longest(r1, r2, fillvalue=0):
634 with support.captured_output('stdout'):
635 print((i, j))
636 result.append((i, j))
637 return result
638 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
639
640 # Formerly, the RuntimeError would be lost
641 # and StopIteration would stop as expected
642 r1 = Repeater(1, 3, RuntimeError)
643 r2 = Repeater(2, 4, StopIteration)
644 it = zip_longest(r1, r2, fillvalue=0)
645 self.assertEqual(next(it), (1, 2))
646 self.assertEqual(next(it), (1, 2))
647 self.assertEqual(next(it), (1, 2))
648 self.assertRaises(RuntimeError, next, it)
649
Christian Heimesc3f30c42008-02-22 16:37:40 +0000650 def test_product(self):
651 for args, result in [
Christian Heimes78644762008-03-04 23:39:23 +0000652 ([], [()]), # zero iterables
Christian Heimesc3f30c42008-02-22 16:37:40 +0000653 (['ab'], [('a',), ('b',)]), # one iterable
654 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
655 ([range(0), range(2), range(3)], []), # first iterable with zero length
656 ([range(2), range(0), range(3)], []), # middle iterable with zero length
657 ([range(2), range(3), range(0)], []), # last iterable with zero length
658 ]:
659 self.assertEqual(list(product(*args)), result)
Christian Heimesf16baeb2008-02-29 14:57:44 +0000660 for r in range(4):
661 self.assertEqual(list(product(*(args*r))),
662 list(product(*args, **dict(repeat=r))))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000663 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
664 self.assertRaises(TypeError, product, range(6), None)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000665
666 def product1(*args, **kwds):
667 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
668 n = len(pools)
669 if n == 0:
670 yield ()
671 return
672 if any(len(pool) == 0 for pool in pools):
673 return
674 indices = [0] * n
675 yield tuple(pool[i] for pool, i in zip(pools, indices))
676 while 1:
677 for i in reversed(range(n)): # right to left
678 if indices[i] == len(pools[i]) - 1:
679 continue
680 indices[i] += 1
681 for j in range(i+1, n):
682 indices[j] = 0
683 yield tuple(pool[i] for pool, i in zip(pools, indices))
684 break
685 else:
686 return
687
688 def product2(*args, **kwds):
689 'Pure python version used in docs'
690 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
691 result = [[]]
692 for pool in pools:
693 result = [x+[y] for x in result for y in pool]
694 for prod in result:
695 yield tuple(prod)
696
Christian Heimesc3f30c42008-02-22 16:37:40 +0000697 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
698 set('abcdefg'), range(11), tuple(range(13))]
699 for i in range(100):
700 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Christian Heimes78644762008-03-04 23:39:23 +0000701 expected_len = prod(map(len, args))
702 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000703 self.assertEqual(list(product(*args)), list(product1(*args)))
704 self.assertEqual(list(product(*args)), list(product2(*args)))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000705 args = map(iter, args)
Christian Heimes78644762008-03-04 23:39:23 +0000706 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000707
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000708 # Test implementation detail: tuple re-use
709 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
710 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000711
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000712 def test_repeat(self):
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +0000713 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
Guido van Rossum805365e2007-05-07 22:24:25 +0000714 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000715 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000716 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000717 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000718 self.assertEqual(list(repeat('a', 0)), [])
719 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000720 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000721 self.assertRaises(TypeError, repeat, None, 3, 4)
722 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000723 r = repeat(1+0j)
724 self.assertEqual(repr(r), 'repeat((1+0j))')
725 r = repeat(1+0j, 5)
726 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
727 list(r)
728 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000729
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000730 def test_map(self):
731 self.assertEqual(list(map(operator.pow, range(3), range(1,7))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000732 [0**1, 1**2, 2**3])
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000733 def tupleize(*args):
734 return args
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000735 self.assertEqual(list(map(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000736 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000737 self.assertEqual(list(map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000738 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000739 self.assertEqual(take(2,map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000740 [('a',0),('b',1)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000741 self.assertEqual(list(map(operator.pow, [])), [])
742 self.assertRaises(TypeError, map)
743 self.assertRaises(TypeError, list, map(None, range(3), range(3)))
744 self.assertRaises(TypeError, map, operator.neg)
745 self.assertRaises(TypeError, next, map(10, range(5)))
746 self.assertRaises(ValueError, next, map(errfunc, [4], [5]))
747 self.assertRaises(TypeError, next, map(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000748
749 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000750 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
751 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000752 self.assertEqual(take(3, starmap(operator.pow, zip(count(), count(1)))),
Raymond Hettinger02420702003-06-29 20:36:23 +0000753 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000754 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +0000755 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
756 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000757 self.assertRaises(TypeError, starmap)
758 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000759 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
760 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
761 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000762
763 def test_islice(self):
764 for args in [ # islice(args) should agree with range(args)
765 (10, 20, 3),
766 (10, 3, 20),
767 (10, 20),
768 (10, 3),
769 (20,)
770 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000771 self.assertEqual(list(islice(range(100), *args)),
772 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000773
774 for args, tgtargs in [ # Stop when seqn is exhausted
775 ((10, 110, 3), ((10, 100, 3))),
776 ((10, 110), ((10, 100))),
777 ((110,), (100,))
778 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000779 self.assertEqual(list(islice(range(100), *args)),
780 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000781
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000782 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +0000783 self.assertEqual(list(islice(range(10), None)), list(range(10)))
784 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
785 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
786 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
787 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000788
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000789 # Test number of items consumed SF #1171417
790 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +0000791 self.assertEqual(list(islice(it, 3)), list(range(3)))
792 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000793
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000794 # Test invalid arguments
Guido van Rossum805365e2007-05-07 22:24:25 +0000795 self.assertRaises(TypeError, islice, range(10))
796 self.assertRaises(TypeError, islice, range(10), 1, 2, 3, 4)
797 self.assertRaises(ValueError, islice, range(10), -5, 10, 1)
798 self.assertRaises(ValueError, islice, range(10), 1, -5, -1)
799 self.assertRaises(ValueError, islice, range(10), 1, 10, -1)
800 self.assertRaises(ValueError, islice, range(10), 1, 10, 0)
801 self.assertRaises(ValueError, islice, range(10), 'a')
802 self.assertRaises(ValueError, islice, range(10), 'a', 1)
803 self.assertRaises(ValueError, islice, range(10), 1, 'a')
804 self.assertRaises(ValueError, islice, range(10), 'a', 1, 1)
805 self.assertRaises(ValueError, islice, range(10), 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +0000806 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000807
Raymond Hettinger69b34bf2010-11-30 02:49:29 +0000808 # Issue #10323: Less islice in a predictable state
809 c = count()
810 self.assertEqual(list(islice(c, 1, 3, 50)), [1])
811 self.assertEqual(next(c), 3)
812
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000813 def test_takewhile(self):
814 data = [1, 3, 5, 20, 2, 4, 6, 8]
815 underten = lambda x: x<10
816 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000817 self.assertEqual(list(takewhile(underten, [])), [])
818 self.assertRaises(TypeError, takewhile)
819 self.assertRaises(TypeError, takewhile, operator.pow)
820 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000821 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
822 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000823 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
824 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +0000825 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000826
827 def test_dropwhile(self):
828 data = [1, 3, 5, 20, 2, 4, 6, 8]
829 underten = lambda x: x<10
830 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000831 self.assertEqual(list(dropwhile(underten, [])), [])
832 self.assertRaises(TypeError, dropwhile)
833 self.assertRaises(TypeError, dropwhile, operator.pow)
834 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000835 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
836 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000837
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000838 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000839 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000840 def irange(n):
Guido van Rossum805365e2007-05-07 22:24:25 +0000841 for i in range(n):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000842 yield i
843
844 a, b = tee([]) # test empty iterator
845 self.assertEqual(list(a), [])
846 self.assertEqual(list(b), [])
847
848 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +0000849 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000850
851 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +0000852 self.assertEqual(list(a), list(range(n)))
853 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000854
855 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000856 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000857 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000858 del a
Guido van Rossum805365e2007-05-07 22:24:25 +0000859 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000860
861 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000862 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000863 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000864 del b
Guido van Rossum805365e2007-05-07 22:24:25 +0000865 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000866
Guido van Rossum805365e2007-05-07 22:24:25 +0000867 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000868 order = [0]*n + [1]*n
869 random.shuffle(order)
870 lists = ([], [])
871 its = tee(irange(n))
872 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +0000873 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000874 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +0000875 self.assertEqual(lists[0], list(range(n)))
876 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000877
Raymond Hettingerad983e72003-11-12 14:32:26 +0000878 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000879 self.assertRaises(TypeError, tee)
880 self.assertRaises(TypeError, tee, 3)
881 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000882 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000883
Raymond Hettingerad983e72003-11-12 14:32:26 +0000884 # tee object should be instantiable
885 a, b = tee('abc')
886 c = type(a)('def')
887 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000888
Raymond Hettingerad983e72003-11-12 14:32:26 +0000889 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +0000890 a, b, c = tee(range(2000), 3)
891 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000892 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +0000893 self.assertEqual(list(b), list(range(2000)))
894 self.assertEqual([next(c), next(c)], list(range(2)))
895 self.assertEqual(list(a), list(range(100,2000)))
896 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000897
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000898 # test values of n
899 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000900 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +0000901 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000902 result = tee('abc', n)
903 self.assertEqual(type(result), tuple)
904 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000905 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000906
Raymond Hettingerad983e72003-11-12 14:32:26 +0000907 # tee pass-through to copyable iterator
908 a, b = tee('abc')
909 c, d = tee(a)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000910 self.assertTrue(a is c)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000911
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000912 # test tee_new
913 t1, t2 = tee('abc')
914 tnew = type(t1)
915 self.assertRaises(TypeError, tnew)
916 self.assertRaises(TypeError, tnew, 10)
917 t3 = tnew(t1)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000918 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000919
Raymond Hettingera9f60922004-10-17 16:40:14 +0000920 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +0000921 a, b = tee(range(10))
Raymond Hettingera9f60922004-10-17 16:40:14 +0000922 p = proxy(a)
923 self.assertEqual(getattr(p, '__class__'), type(b))
924 del a
925 self.assertRaises(ReferenceError, getattr, p, '__class__')
926
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000927 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000928 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000929
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000930 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +0000931 self.assertRaises(StopIteration, next, f([]))
932 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000933
Georg Brandla18af4e2007-04-21 15:47:16 +0000934 self.assertRaises(StopIteration, next, islice([], None))
935 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000936
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000937 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +0000938 self.assertRaises(StopIteration, next, p)
939 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000940 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +0000941 self.assertRaises(StopIteration, next, p)
942 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000943
Georg Brandla18af4e2007-04-21 15:47:16 +0000944 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000945
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000946 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +0000947 self.assertRaises(StopIteration, next, f(lambda x:x, []))
948 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000949
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000950class TestExamples(unittest.TestCase):
951
Raymond Hettinger482ba772010-12-01 22:48:00 +0000952 def test_accumlate(self):
953 self.assertEqual(list(accumulate([1,2,3,4,5])), [1, 3, 6, 10, 15])
954
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000955 def test_chain(self):
956 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
957
958 def test_chain_from_iterable(self):
959 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
960
961 def test_combinations(self):
962 self.assertEqual(list(combinations('ABCD', 2)),
963 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
964 self.assertEqual(list(combinations(range(4), 3)),
965 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
966
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000967 def test_combinations_with_replacement(self):
968 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
969 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
970
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000971 def test_compress(self):
972 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
973
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000974 def test_count(self):
975 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
976
977 def test_cycle(self):
978 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
979
980 def test_dropwhile(self):
981 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
982
983 def test_groupby(self):
984 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
985 list('ABCDAB'))
986 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
987 [list('AAAA'), list('BBB'), list('CC'), list('D')])
988
989 def test_filter(self):
990 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
991
992 def test_filterfalse(self):
993 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
994
995 def test_map(self):
996 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
997
998 def test_islice(self):
999 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
1000 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
1001 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
1002 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1003
1004 def test_zip(self):
1005 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
1006
1007 def test_zip_longest(self):
1008 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
1009 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
1010
1011 def test_permutations(self):
1012 self.assertEqual(list(permutations('ABCD', 2)),
1013 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
1014 self.assertEqual(list(permutations(range(3))),
1015 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
1016
1017 def test_product(self):
1018 self.assertEqual(list(product('ABCD', 'xy')),
1019 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
1020 self.assertEqual(list(product(range(2), repeat=3)),
1021 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
1022 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
1023
1024 def test_repeat(self):
1025 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
1026
1027 def test_stapmap(self):
1028 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
1029 [32, 9, 1000])
1030
1031 def test_takewhile(self):
1032 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
1033
1034
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001035class TestGC(unittest.TestCase):
1036
1037 def makecycle(self, iterator, container):
1038 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +00001039 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001040 del container, iterator
1041
Raymond Hettinger482ba772010-12-01 22:48:00 +00001042 def test_accumulate(self):
1043 a = []
1044 self.makecycle(accumulate([1,2,a,3]), a)
1045
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001046 def test_chain(self):
1047 a = []
1048 self.makecycle(chain(a), a)
1049
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001050 def test_chain_from_iterable(self):
1051 a = []
1052 self.makecycle(chain.from_iterable([a]), a)
1053
1054 def test_combinations(self):
1055 a = []
1056 self.makecycle(combinations([1,2,a,3], 3), a)
1057
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001058 def test_combinations_with_replacement(self):
1059 a = []
1060 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1061
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001062 def test_compress(self):
1063 a = []
1064 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1065
Raymond Hettingerd6280f42009-02-16 20:50:56 +00001066 def test_count(self):
1067 a = []
1068 Int = type('Int', (int,), dict(x=a))
1069 self.makecycle(count(Int(0), Int(1)), a)
1070
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001071 def test_cycle(self):
1072 a = []
1073 self.makecycle(cycle([a]*2), a)
1074
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001075 def test_dropwhile(self):
1076 a = []
1077 self.makecycle(dropwhile(bool, [0, a, a]), a)
1078
1079 def test_groupby(self):
1080 a = []
1081 self.makecycle(groupby([a]*2, lambda x:x), a)
1082
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001083 def test_issue2246(self):
1084 # Issue 2246 -- the _grouper iterator was not included in GC
1085 n = 10
1086 keyfunc = lambda x: x
1087 for i, j in groupby(range(n), key=keyfunc):
1088 keyfunc.__dict__.setdefault('x',[]).append(j)
1089
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001090 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001091 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001092 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001093
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001094 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001095 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001096 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001097
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001098 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001099 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001100 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001101
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001102 def test_zip_longest(self):
1103 a = []
1104 self.makecycle(zip_longest([a]*2, [a]*3), a)
1105 b = [a, None]
1106 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1107
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001108 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001109 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001110 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001111
1112 def test_islice(self):
1113 a = []
1114 self.makecycle(islice([a]*2, None), a)
1115
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001116 def test_permutations(self):
1117 a = []
1118 self.makecycle(permutations([1,2,a,3], 3), a)
1119
1120 def test_product(self):
1121 a = []
1122 self.makecycle(product([1,2,a,3], repeat=3), a)
1123
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001124 def test_repeat(self):
1125 a = []
1126 self.makecycle(repeat(a), a)
1127
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001128 def test_starmap(self):
1129 a = []
1130 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1131
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001132 def test_takewhile(self):
1133 a = []
1134 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1135
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001136def R(seqn):
1137 'Regular generator'
1138 for i in seqn:
1139 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001140
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001141class G:
1142 'Sequence using __getitem__'
1143 def __init__(self, seqn):
1144 self.seqn = seqn
1145 def __getitem__(self, i):
1146 return self.seqn[i]
1147
1148class I:
1149 'Sequence using iterator protocol'
1150 def __init__(self, seqn):
1151 self.seqn = seqn
1152 self.i = 0
1153 def __iter__(self):
1154 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001155 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001156 if self.i >= len(self.seqn): raise StopIteration
1157 v = self.seqn[self.i]
1158 self.i += 1
1159 return v
1160
1161class Ig:
1162 'Sequence using iterator protocol defined with a generator'
1163 def __init__(self, seqn):
1164 self.seqn = seqn
1165 self.i = 0
1166 def __iter__(self):
1167 for val in self.seqn:
1168 yield val
1169
1170class X:
1171 'Missing __getitem__ and __iter__'
1172 def __init__(self, seqn):
1173 self.seqn = seqn
1174 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001175 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001176 if self.i >= len(self.seqn): raise StopIteration
1177 v = self.seqn[self.i]
1178 self.i += 1
1179 return v
1180
1181class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001182 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001183 def __init__(self, seqn):
1184 self.seqn = seqn
1185 self.i = 0
1186 def __iter__(self):
1187 return self
1188
1189class E:
1190 'Test propagation of exceptions'
1191 def __init__(self, seqn):
1192 self.seqn = seqn
1193 self.i = 0
1194 def __iter__(self):
1195 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001196 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001197 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001198
1199class S:
1200 'Test immediate stop'
1201 def __init__(self, seqn):
1202 pass
1203 def __iter__(self):
1204 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001205 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001206 raise StopIteration
1207
1208def L(seqn):
1209 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001210 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001211
1212
1213class TestVariousIteratorArgs(unittest.TestCase):
1214
Raymond Hettinger482ba772010-12-01 22:48:00 +00001215 def test_accumulate(self):
1216 s = [1,2,3,4,5]
1217 r = [1,3,6,10,15]
1218 n = len(s)
1219 for g in (G, I, Ig, L, R):
1220 self.assertEqual(list(accumulate(g(s))), r)
1221 self.assertEqual(list(accumulate(S(s))), [])
1222 self.assertRaises(TypeError, accumulate, X(s))
1223 self.assertRaises(TypeError, accumulate, N(s))
1224 self.assertRaises(ZeroDivisionError, list, accumulate(E(s)))
1225
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001226 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001227 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001228 for g in (G, I, Ig, S, L, R):
1229 self.assertEqual(list(chain(g(s))), list(g(s)))
1230 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001231 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001232 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001233 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1234
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001235 def test_compress(self):
1236 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1237 n = len(s)
1238 for g in (G, I, Ig, S, L, R):
1239 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1240 self.assertRaises(TypeError, compress, X(s), repeat(1))
1241 self.assertRaises(TypeError, compress, N(s), repeat(1))
1242 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1243
Christian Heimesc3f30c42008-02-22 16:37:40 +00001244 def test_product(self):
1245 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1246 self.assertRaises(TypeError, product, X(s))
1247 self.assertRaises(TypeError, product, N(s))
1248 self.assertRaises(ZeroDivisionError, product, E(s))
1249
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001250 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001251 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001252 for g in (G, I, Ig, S, L, R):
1253 tgtlen = len(s) * 3
1254 expected = list(g(s))*3
1255 actual = list(islice(cycle(g(s)), tgtlen))
1256 self.assertEqual(actual, expected)
1257 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001258 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001259 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1260
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001261 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001262 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001263 for g in (G, I, Ig, S, L, R):
1264 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1265 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001266 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001267 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1268
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001269 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001270 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001271 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001272 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001273 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001274 self.assertRaises(TypeError, filter, isEven, X(s))
1275 self.assertRaises(TypeError, filter, isEven, N(s))
1276 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001277
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001278 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001279 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001280 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001281 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001282 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001283 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1284 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1285 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001286
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001287 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001288 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001289 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001290 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1291 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1292 self.assertRaises(TypeError, zip, X(s))
1293 self.assertRaises(TypeError, zip, N(s))
1294 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001295
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001296 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001297 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001298 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001299 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1300 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1301 self.assertRaises(TypeError, zip_longest, X(s))
1302 self.assertRaises(TypeError, zip_longest, N(s))
1303 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001304
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001305 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001306 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001307 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001308 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001309 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001310 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001311 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001312 self.assertRaises(TypeError, map, onearg, X(s))
1313 self.assertRaises(TypeError, map, onearg, N(s))
1314 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001315
1316 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001317 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001318 for g in (G, I, Ig, S, L, R):
1319 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1320 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001321 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001322 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1323
1324 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001325 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001326 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001327 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001328 self.assertEqual(list(starmap(operator.pow, g(ss))),
1329 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001330 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001331 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001332 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1333
1334 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001335 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001336 for g in (G, I, Ig, S, L, R):
1337 tgt = []
1338 for elem in g(s):
1339 if not isEven(elem): break
1340 tgt.append(elem)
1341 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1342 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001343 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001344 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1345
1346 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001347 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001348 for g in (G, I, Ig, S, L, R):
1349 tgt = []
1350 for elem in g(s):
1351 if not tgt and isOdd(elem): continue
1352 tgt.append(elem)
1353 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1354 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001355 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001356 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1357
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001358 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001359 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001360 for g in (G, I, Ig, S, L, R):
1361 it1, it2 = tee(g(s))
1362 self.assertEqual(list(it1), list(g(s)))
1363 self.assertEqual(list(it2), list(g(s)))
1364 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001365 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001366 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1367
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001368class LengthTransparency(unittest.TestCase):
1369
1370 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001371 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001372 self.assertEqual(len(repeat(None, 50)), 50)
1373 self.assertRaises(TypeError, len, repeat(None))
1374
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001375class RegressionTests(unittest.TestCase):
1376
1377 def test_sf_793826(self):
1378 # Fix Armin Rigo's successful efforts to wreak havoc
1379
1380 def mutatingtuple(tuple1, f, tuple2):
1381 # this builds a tuple t which is a copy of tuple1,
1382 # then calls f(t), then mutates t to be equal to tuple2
1383 # (needs len(tuple1) == len(tuple2)).
1384 def g(value, first=[1]):
1385 if first:
1386 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00001387 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001388 return value
1389 items = list(tuple2)
1390 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001391 gen = map(g, items)
1392 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00001393 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001394
1395 def f(t):
1396 global T
1397 T = t
1398 first[:] = list(T)
1399
1400 first = []
1401 mutatingtuple((1,2,3), f, (4,5,6))
1402 second = list(T)
1403 self.assertEqual(first, second)
1404
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001405
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001406 def test_sf_950057(self):
1407 # Make sure that chain() and cycle() catch exceptions immediately
1408 # rather than when shifting between input sources
1409
1410 def gen1():
1411 hist.append(0)
1412 yield 1
1413 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001414 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001415 hist.append(2)
1416
1417 def gen2(x):
1418 hist.append(3)
1419 yield 2
1420 hist.append(4)
1421 if x:
1422 raise StopIteration
1423
1424 hist = []
1425 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1426 self.assertEqual(hist, [0,1])
1427
1428 hist = []
1429 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1430 self.assertEqual(hist, [0,1])
1431
1432 hist = []
1433 self.assertRaises(AssertionError, list, cycle(gen1()))
1434 self.assertEqual(hist, [0,1])
1435
Thomas Woutersb2137042007-02-01 18:02:27 +00001436class SubclassWithKwargsTest(unittest.TestCase):
1437 def test_keywords_in_subclass(self):
1438 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001439 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001440 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00001441 class Subclass(cls):
1442 def __init__(self, newarg=None, *args):
1443 cls.__init__(self, *args)
1444 try:
1445 Subclass(newarg=1)
1446 except TypeError as err:
1447 # we expect type errors because of wrong argument count
Ezio Melottib58e0bd2010-01-23 15:40:09 +00001448 self.assertNotIn("does not take keyword arguments", err.args[0])
Thomas Woutersb2137042007-02-01 18:02:27 +00001449
1450
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001451libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001452
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001453
1454>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001455>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00001456... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001457...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001458Check 1200 is for $120.15
1459Check 1201 is for $764.05
1460Check 1202 is for $823.14
1461
1462>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001463>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00001464... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00001465...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000014661
14678
146827
1469
1470>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001471>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00001472... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00001473...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001474Alex
1475Laura
1476Martin
1477Walter
1478Samuele
1479
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001480>>> from operator import itemgetter
1481>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00001482>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001483>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001484... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001485...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000014861 ['a', 'c', 'e']
14872 ['b', 'd', 'f']
14883 ['g']
1489
Raymond Hettinger734fb572004-01-20 20:04:40 +00001490# Find runs of consecutive numbers using groupby. The key to the solution
1491# is differencing with a range so that consecutive numbers all appear in
1492# same group.
1493>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00001494>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001495... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001496...
Raymond Hettinger734fb572004-01-20 20:04:40 +00001497[1]
1498[4, 5, 6]
1499[10]
1500[15, 16, 17, 18]
1501[22]
1502[25, 26, 27, 28]
1503
Georg Brandl3dbca812008-07-23 16:10:53 +00001504>>> def take(n, iterable):
1505... "Return first n items of the iterable as a list"
1506... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001507
Georg Brandl3dbca812008-07-23 16:10:53 +00001508>>> def enumerate(iterable, start=0):
1509... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001510
Georg Brandl3dbca812008-07-23 16:10:53 +00001511>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001512... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00001513... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001514
Raymond Hettingercdf8ba32009-02-19 04:45:07 +00001515>>> def nth(iterable, n, default=None):
1516... "Returns the nth item or a default value"
1517... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001518
Georg Brandl3dbca812008-07-23 16:10:53 +00001519>>> def quantify(iterable, pred=bool):
1520... "Count how many times the predicate is true"
1521... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001522
Georg Brandl3dbca812008-07-23 16:10:53 +00001523>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001524... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00001525... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001526
Georg Brandl3dbca812008-07-23 16:10:53 +00001527>>> def ncycles(iterable, n):
1528... "Returns the seqeuence elements n times"
1529... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001530
1531>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001532... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001533
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001534>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001535... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001536
1537>>> def repeatfunc(func, times=None, *args):
1538... "Repeat calls to func with specified arguments."
1539... " Example: repeatfunc(random.random)"
1540... if times is None:
1541... return starmap(func, repeat(args))
1542... else:
1543... return starmap(func, repeat(args, times))
1544
Raymond Hettingerd591f662003-10-26 15:34:50 +00001545>>> def pairwise(iterable):
1546... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1547... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001548... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00001549... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001550... except StopIteration:
1551... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001552... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001553
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001554>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001555... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001556... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00001557... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001558
1559>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001560... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001561... # Recipe credited to George Sakkis
1562... pending = len(iterables)
1563... nexts = cycle(iter(it).__next__ for it in iterables)
1564... while pending:
1565... try:
1566... for next in nexts:
1567... yield next()
1568... except StopIteration:
1569... pending -= 1
1570... nexts = cycle(islice(nexts, pending))
1571
1572>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00001573... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1574... s = list(iterable)
1575... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001576
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001577>>> def unique_everseen(iterable, key=None):
1578... "List unique elements, preserving order. Remember all elements ever seen."
1579... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1580... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1581... seen = set()
1582... seen_add = seen.add
1583... if key is None:
1584... for element in iterable:
1585... if element not in seen:
1586... seen_add(element)
1587... yield element
1588... else:
1589... for element in iterable:
1590... k = key(element)
1591... if k not in seen:
1592... seen_add(k)
1593... yield element
1594
1595>>> def unique_justseen(iterable, key=None):
1596... "List unique elements, preserving order. Remember only the element just seen."
1597... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1598... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1599... return map(next, map(itemgetter(1), groupby(iterable, key)))
1600
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001601This is not part of the examples but it tests to make sure the definitions
1602perform as purported.
1603
Raymond Hettingera098b332003-09-08 23:58:40 +00001604>>> take(10, count())
1605[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1606
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001607>>> list(enumerate('abc'))
1608[(0, 'a'), (1, 'b'), (2, 'c')]
1609
1610>>> list(islice(tabulate(lambda x: 2*x), 4))
1611[0, 2, 4, 6]
1612
1613>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00001614'd'
1615
1616>>> nth('abcde', 9) is None
1617True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001618
Guido van Rossum805365e2007-05-07 22:24:25 +00001619>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000162050
1621
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001622>>> a = [[1, 2, 3], [4, 5, 6]]
1623>>> flatten(a)
1624[1, 2, 3, 4, 5, 6]
1625
1626>>> list(repeatfunc(pow, 5, 2, 3))
1627[8, 8, 8, 8, 8]
1628
1629>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001630>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001631[0, 0, 0, 0, 0]
1632
Raymond Hettingerd591f662003-10-26 15:34:50 +00001633>>> list(pairwise('abcd'))
1634[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001635
Raymond Hettingerd591f662003-10-26 15:34:50 +00001636>>> list(pairwise([]))
1637[]
1638
1639>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001640[]
1641
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001642>>> list(islice(padnone('abc'), 0, 6))
1643['a', 'b', 'c', None, None, None]
1644
1645>>> list(ncycles('abc', 3))
1646['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1647
1648>>> dotproduct([1,2,3], [4,5,6])
164932
1650
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001651>>> list(grouper(3, 'abcdefg', 'x'))
1652[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1653
1654>>> list(roundrobin('abc', 'd', 'ef'))
1655['a', 'd', 'e', 'b', 'f', 'c']
1656
Raymond Hettingerace67332009-01-26 02:23:50 +00001657>>> list(powerset([1,2,3]))
1658[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001659
Raymond Hettinger191e8502009-01-27 13:29:43 +00001660>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1661True
1662
1663>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1664True
1665
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001666>>> list(unique_everseen('AAAABBBCCDAABBB'))
1667['A', 'B', 'C', 'D']
1668
1669>>> list(unique_everseen('ABBCcAD', str.lower))
1670['A', 'B', 'C', 'D']
1671
1672>>> list(unique_justseen('AAAABBBCCDAABBB'))
1673['A', 'B', 'C', 'D', 'A', 'B']
1674
1675>>> list(unique_justseen('ABBCcAD', str.lower))
1676['A', 'B', 'C', 'A', 'D']
1677
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001678"""
1679
1680__test__ = {'libreftest' : libreftest}
1681
1682def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001683 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00001684 RegressionTests, LengthTransparency,
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001685 SubclassWithKwargsTest, TestExamples)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001686 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001687
1688 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001689 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001690 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001691 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00001692 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001693 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001694 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001695 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001696 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001697
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001698 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001699 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001700
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001701if __name__ == "__main__":
1702 test_main(verbose=True)