blob: 66e307da5a9321441557d030673ff6be58386def [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
Raymond Hettingerd8ff4652010-12-03 02:09:34 +000062 [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])
63 self.assertEqual(list(accumulate(iterable=range(10))), # kw arg
64 [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])
Raymond Hettinger482ba772010-12-01 22:48:00 +000065 for typ in int, complex, Decimal, Fraction: # multiple types
Raymond Hettingerd8ff4652010-12-03 02:09:34 +000066 self.assertEqual(
67 list(accumulate(map(typ, range(10)))),
Raymond Hettinger482ba772010-12-01 22:48:00 +000068 list(map(typ, [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])))
Raymond Hettinger2d93e6e2010-12-03 02:33:53 +000069 self.assertEqual(list(accumulate('abc')), ['a', 'ab', 'abc']) # works with non-numeric
Raymond Hettinger482ba772010-12-01 22:48:00 +000070 self.assertEqual(list(accumulate([])), []) # empty iterable
Raymond Hettingerd8ff4652010-12-03 02:09:34 +000071 self.assertEqual(list(accumulate([7])), [7]) # iterable of length one
72 self.assertRaises(TypeError, accumulate, range(10), 5) # too many args
Raymond Hettinger482ba772010-12-01 22:48:00 +000073 self.assertRaises(TypeError, accumulate) # too few args
Raymond Hettingerd8ff4652010-12-03 02:09:34 +000074 self.assertRaises(TypeError, accumulate, x=range(10)) # unexpected kwd arg
Raymond Hettinger482ba772010-12-01 22:48:00 +000075 self.assertRaises(TypeError, list, accumulate([1, []])) # args that don't add
76
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000077 def test_chain(self):
Christian Heimesdd15f6c2008-03-16 00:07:10 +000078
79 def chain2(*iterables):
80 'Pure python version in the docs'
81 for it in iterables:
82 for element in it:
83 yield element
84
85 for c in (chain, chain2):
86 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
87 self.assertEqual(list(c('abc')), list('abc'))
88 self.assertEqual(list(c('')), [])
89 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
90 self.assertRaises(TypeError, list,c(2, 3))
Christian Heimesf16baeb2008-02-29 14:57:44 +000091
92 def test_chain_from_iterable(self):
93 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
94 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
95 self.assertEqual(list(chain.from_iterable([''])), [])
96 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
97 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000098
Christian Heimes380f7f22008-02-28 11:19:05 +000099 def test_combinations(self):
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000100 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Christian Heimes380f7f22008-02-28 11:19:05 +0000101 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Christian Heimes78644762008-03-04 23:39:23 +0000102 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Christian Heimes380f7f22008-02-28 11:19:05 +0000103 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000104 self.assertEqual(list(combinations('abc', 32)), []) # r > n
Christian Heimes380f7f22008-02-28 11:19:05 +0000105 self.assertEqual(list(combinations(range(4), 3)),
106 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
Christian Heimes78644762008-03-04 23:39:23 +0000107
108 def combinations1(iterable, r):
109 'Pure python version shown in the docs'
110 pool = tuple(iterable)
111 n = len(pool)
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000112 if r > n:
113 return
Christian Heimes78644762008-03-04 23:39:23 +0000114 indices = list(range(r))
115 yield tuple(pool[i] for i in indices)
116 while 1:
117 for i in reversed(range(r)):
118 if indices[i] != i + n - r:
119 break
120 else:
121 return
122 indices[i] += 1
123 for j in range(i+1, r):
124 indices[j] = indices[j-1] + 1
125 yield tuple(pool[i] for i in indices)
126
127 def combinations2(iterable, r):
128 'Pure python version shown in the docs'
129 pool = tuple(iterable)
130 n = len(pool)
131 for indices in permutations(range(n), r):
132 if sorted(indices) == list(indices):
133 yield tuple(pool[i] for i in indices)
134
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000135 def combinations3(iterable, r):
136 'Pure python version from cwr()'
137 pool = tuple(iterable)
138 n = len(pool)
139 for indices in combinations_with_replacement(range(n), r):
140 if len(set(indices)) == r:
141 yield tuple(pool[i] for i in indices)
142
Christian Heimes78644762008-03-04 23:39:23 +0000143 for n in range(7):
Christian Heimes380f7f22008-02-28 11:19:05 +0000144 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000145 for r in range(n+2):
Christian Heimes380f7f22008-02-28 11:19:05 +0000146 result = list(combinations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000147 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 +0000148 self.assertEqual(len(result), len(set(result))) # no repeats
149 self.assertEqual(result, sorted(result)) # lexicographic order
150 for c in result:
151 self.assertEqual(len(c), r) # r-length combinations
152 self.assertEqual(len(set(c)), r) # no duplicate elements
153 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000154 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000155 self.assertEqual(list(c),
156 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000157 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000158 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000159 self.assertEqual(result, list(combinations3(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000160
161 # Test implementation detail: tuple re-use
162 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
163 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
164
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000165 def test_combinations_with_replacement(self):
166 cwr = combinations_with_replacement
167 self.assertRaises(TypeError, cwr, 'abc') # missing r argument
168 self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
169 self.assertRaises(TypeError, cwr, None) # pool is not iterable
170 self.assertRaises(ValueError, cwr, 'abc', -2) # r is negative
171 self.assertEqual(list(cwr('ABC', 2)),
172 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
173
174 def cwr1(iterable, r):
175 'Pure python version shown in the docs'
176 # number items returned: (n+r-1)! / r! / (n-1)! when n>0
177 pool = tuple(iterable)
178 n = len(pool)
179 if not n and r:
180 return
181 indices = [0] * r
182 yield tuple(pool[i] for i in indices)
183 while 1:
184 for i in reversed(range(r)):
185 if indices[i] != n - 1:
186 break
187 else:
188 return
189 indices[i:] = [indices[i] + 1] * (r - i)
190 yield tuple(pool[i] for i in indices)
191
192 def cwr2(iterable, r):
193 'Pure python version shown in the docs'
194 pool = tuple(iterable)
195 n = len(pool)
196 for indices in product(range(n), repeat=r):
197 if sorted(indices) == list(indices):
198 yield tuple(pool[i] for i in indices)
199
200 def numcombs(n, r):
201 if not n:
202 return 0 if r else 1
203 return fact(n+r-1) / fact(r)/ fact(n-1)
204
205 for n in range(7):
206 values = [5*x-12 for x in range(n)]
207 for r in range(n+2):
208 result = list(cwr(values, r))
209
210 self.assertEqual(len(result), numcombs(n, r)) # right number of combs
211 self.assertEqual(len(result), len(set(result))) # no repeats
212 self.assertEqual(result, sorted(result)) # lexicographic order
213
214 regular_combs = list(combinations(values, r)) # compare to combs without replacement
215 if n == 0 or r <= 1:
Ezio Melottib3aedd42010-11-20 19:04:17 +0000216 self.assertEqual(result, regular_combs) # cases that should be identical
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000217 else:
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000218 self.assertTrue(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000219
220 for c in result:
221 self.assertEqual(len(c), r) # r-length combinations
222 noruns = [k for k,v in groupby(c)] # combo without consecutive repeats
223 self.assertEqual(len(noruns), len(set(noruns))) # no repeats other than consecutive
224 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000225 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000226 self.assertEqual(noruns,
227 [e for e in values if e in c]) # comb is a subsequence of the input iterable
228 self.assertEqual(result, list(cwr1(values, r))) # matches first pure python version
229 self.assertEqual(result, list(cwr2(values, r))) # matches second pure python version
230
231 # Test implementation detail: tuple re-use
232 self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
233 self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
234
Christian Heimes78644762008-03-04 23:39:23 +0000235 def test_permutations(self):
236 self.assertRaises(TypeError, permutations) # too few arguments
237 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000238 self.assertRaises(TypeError, permutations, None) # pool is not iterable
239 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000240 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000241 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Christian Heimes78644762008-03-04 23:39:23 +0000242 self.assertEqual(list(permutations(range(3), 2)),
243 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
244
245 def permutations1(iterable, r=None):
246 'Pure python version shown in the docs'
247 pool = tuple(iterable)
248 n = len(pool)
249 r = n if r is None else r
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000250 if r > n:
251 return
Christian Heimes78644762008-03-04 23:39:23 +0000252 indices = list(range(n))
253 cycles = list(range(n-r+1, n+1))[::-1]
254 yield tuple(pool[i] for i in indices[:r])
255 while n:
256 for i in reversed(range(r)):
257 cycles[i] -= 1
258 if cycles[i] == 0:
259 indices[i:] = indices[i+1:] + indices[i:i+1]
260 cycles[i] = n - i
261 else:
262 j = cycles[i]
263 indices[i], indices[-j] = indices[-j], indices[i]
264 yield tuple(pool[i] for i in indices[:r])
265 break
266 else:
267 return
268
269 def permutations2(iterable, r=None):
270 'Pure python version shown in the docs'
271 pool = tuple(iterable)
272 n = len(pool)
273 r = n if r is None else r
274 for indices in product(range(n), repeat=r):
275 if len(set(indices)) == r:
276 yield tuple(pool[i] for i in indices)
277
278 for n in range(7):
279 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000280 for r in range(n+2):
Christian Heimes78644762008-03-04 23:39:23 +0000281 result = list(permutations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000282 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 +0000283 self.assertEqual(len(result), len(set(result))) # no repeats
284 self.assertEqual(result, sorted(result)) # lexicographic order
285 for p in result:
286 self.assertEqual(len(p), r) # r-length permutations
287 self.assertEqual(len(set(p)), r) # no duplicate elements
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000288 self.assertTrue(all(e in values for e in p)) # elements taken from input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000289 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000290 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000291 if r == n:
292 self.assertEqual(result, list(permutations(values, None))) # test r as None
293 self.assertEqual(result, list(permutations(values))) # test default r
294
295 # Test implementation detail: tuple re-use
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000296 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Christian Heimes78644762008-03-04 23:39:23 +0000297 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Christian Heimes380f7f22008-02-28 11:19:05 +0000298
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000299 def test_combinatorics(self):
300 # Test relationships between product(), permutations(),
301 # combinations() and combinations_with_replacement().
302
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000303 for n in range(6):
304 s = 'ABCDEFG'[:n]
305 for r in range(8):
306 prod = list(product(s, repeat=r))
307 cwr = list(combinations_with_replacement(s, r))
308 perm = list(permutations(s, r))
309 comb = list(combinations(s, r))
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000310
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000311 # Check size
Ezio Melottib3aedd42010-11-20 19:04:17 +0000312 self.assertEqual(len(prod), n**r)
313 self.assertEqual(len(cwr), (fact(n+r-1) / fact(r)/ fact(n-1)) if n else (not r))
314 self.assertEqual(len(perm), 0 if r>n else fact(n) / fact(n-r))
315 self.assertEqual(len(comb), 0 if r>n else fact(n) / fact(r) / fact(n-r))
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000316
317 # Check lexicographic order without repeated tuples
Ezio Melottib3aedd42010-11-20 19:04:17 +0000318 self.assertEqual(prod, sorted(set(prod)))
319 self.assertEqual(cwr, sorted(set(cwr)))
320 self.assertEqual(perm, sorted(set(perm)))
321 self.assertEqual(comb, sorted(set(comb)))
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000322
323 # Check interrelationships
Ezio Melottib3aedd42010-11-20 19:04:17 +0000324 self.assertEqual(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
325 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 +0000326 self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
327 self.assertEqual(comb, [t for t in cwr if len(set(t))==r]) # comb: cwrs without dups
328 self.assertEqual(comb, list(filter(set(cwr).__contains__, perm))) # comb: perm that is a cwr
329 self.assertEqual(comb, list(filter(set(perm).__contains__, cwr))) # comb: cwr that is a perm
330 self.assertEqual(comb, sorted(set(cwr) & set(perm))) # comb: both a cwr and a perm
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000331
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000332 def test_compress(self):
Raymond Hettinger15a49502009-02-19 02:17:09 +0000333 self.assertEqual(list(compress(data='ABCDEF', selectors=[1,0,1,0,1,1])), list('ACEF'))
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000334 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
335 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
336 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
337 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
338 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
339 n = 10000
340 data = chain.from_iterable(repeat(range(6), n))
341 selectors = chain.from_iterable(repeat((0, 1)))
342 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
343 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
344 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
345 self.assertRaises(TypeError, compress, range(6)) # too few args
346 self.assertRaises(TypeError, compress, range(6), None) # too many args
347
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000348 def test_count(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000349 self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
350 self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
351 self.assertEqual(take(2, lzip('abc',count(3))), [('a', 3), ('b', 4)])
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000352 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
353 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000354 self.assertRaises(TypeError, count, 2, 3, 4)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000355 self.assertRaises(TypeError, count, 'a')
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000356 self.assertEqual(list(islice(count(maxsize-5), 10)),
357 list(range(maxsize-5, maxsize+5)))
358 self.assertEqual(list(islice(count(-maxsize-5), 10)),
359 list(range(-maxsize-5, -maxsize+5)))
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +0000360 self.assertEqual(list(islice(count(10, maxsize+5), 3)),
361 list(range(10, 10+3*(maxsize+5), maxsize+5)))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000362 c = count(3)
363 self.assertEqual(repr(c), 'count(3)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000364 next(c)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000365 self.assertEqual(repr(c), 'count(4)')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000366 c = count(-9)
367 self.assertEqual(repr(c), 'count(-9)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000368 next(c)
Raymond Hettinger30729212009-02-12 06:28:27 +0000369 self.assertEqual(repr(count(10.25)), 'count(10.25)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000370 self.assertEqual(next(c), -8)
Christian Heimesa37d4c62007-12-04 23:02:19 +0000371 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 +0000372 # Test repr (ignoring the L in longs)
373 r1 = repr(count(i)).replace('L', '')
374 r2 = 'count(%r)'.__mod__(i).replace('L', '')
375 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000376
Raymond Hettingera3e1ad22009-11-30 22:02:31 +0000377 # check copy, deepcopy, pickle
378 for value in -3, 3, maxsize-5, maxsize+5:
379 c = count(value)
380 self.assertEqual(next(copy.copy(c)), value)
381 self.assertEqual(next(copy.deepcopy(c)), value)
382 self.assertEqual(next(pickle.loads(pickle.dumps(c))), value)
383
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +0000384 #check proper internal error handling for large "step' sizes
385 count(1, maxsize+5); sys.exc_info()
386
Raymond Hettinger30729212009-02-12 06:28:27 +0000387 def test_count_with_stride(self):
388 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingereb13fdd2009-02-21 22:30:12 +0000389 self.assertEqual(lzip('abc',count(start=2,step=3)),
390 [('a', 2), ('b', 5), ('c', 8)])
391 self.assertEqual(lzip('abc',count(step=-1)),
392 [('a', 0), ('b', -1), ('c', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000393 self.assertEqual(lzip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
394 self.assertEqual(lzip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000395 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000396 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
397 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
398 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettinger8a882a72009-02-12 12:08:18 +0000399 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
400 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerde09acf2009-02-12 12:53:53 +0000401 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
402 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000403 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
404 c = count(3, 5)
405 self.assertEqual(repr(c), 'count(3, 5)')
406 next(c)
407 self.assertEqual(repr(c), 'count(8, 5)')
408 c = count(-9, 0)
409 self.assertEqual(repr(c), 'count(-9, 0)')
410 next(c)
411 self.assertEqual(repr(c), 'count(-9, 0)')
412 c = count(-9, -3)
413 self.assertEqual(repr(c), 'count(-9, -3)')
414 next(c)
415 self.assertEqual(repr(c), 'count(-12, -3)')
416 self.assertEqual(repr(c), 'count(-12, -3)')
417 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
418 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
419 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
420 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
421 for j in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 1, 10, sys.maxsize-5, sys.maxsize+5):
422 # Test repr (ignoring the L in longs)
423 r1 = repr(count(i, j)).replace('L', '')
424 if j == 1:
425 r2 = ('count(%r)' % i).replace('L', '')
426 else:
427 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
428 self.assertEqual(r1, r2)
429
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000430 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000431 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000432 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000433 self.assertRaises(TypeError, cycle)
434 self.assertRaises(TypeError, cycle, 5)
435 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000436
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000437 def test_groupby(self):
438 # Check whether it accepts arguments correctly
439 self.assertEqual([], list(groupby([])))
440 self.assertEqual([], list(groupby([], key=id)))
441 self.assertRaises(TypeError, list, groupby('abc', []))
442 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000443 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000444
445 # Check normal input
446 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
447 (2,15,22), (3,16,23), (3,17,23)]
448 dup = []
449 for k, g in groupby(s, lambda r:r[0]):
450 for elem in g:
451 self.assertEqual(k, elem[0])
452 dup.append(elem)
453 self.assertEqual(s, dup)
454
455 # Check nested case
456 dup = []
457 for k, g in groupby(s, lambda r:r[0]):
458 for ik, ig in groupby(g, lambda r:r[2]):
459 for elem in ig:
460 self.assertEqual(k, elem[0])
461 self.assertEqual(ik, elem[2])
462 dup.append(elem)
463 self.assertEqual(s, dup)
464
465 # Check case where inner iterator is not used
466 keys = [k for k, g in groupby(s, lambda r:r[0])]
467 expectedkeys = set([r[0] for r in s])
468 self.assertEqual(set(keys), expectedkeys)
469 self.assertEqual(len(keys), len(expectedkeys))
470
471 # Exercise pipes and filters style
472 s = 'abracadabra'
473 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000474 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000475 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
476 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000477 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000478 self.assertEqual(r, ['a', 'b', 'r'])
479 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000480 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000481 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
482 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000483 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000484 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
485
Georg Brandla18af4e2007-04-21 15:47:16 +0000486 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000487 class ExpectedError(Exception):
488 pass
489 def delayed_raise(n=0):
490 for i in range(n):
491 yield 'yo'
492 raise ExpectedError
493 def gulp(iterable, keyp=None, func=list):
494 return [func(g) for k, g in groupby(iterable, keyp)]
495
Georg Brandla18af4e2007-04-21 15:47:16 +0000496 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000497 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000498 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000499 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
500
501 # __cmp__ failure
502 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000503 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000504 raise ExpectedError
505 s = [DummyCmp(), DummyCmp(), None]
506
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000507 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000508 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000509 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000510 self.assertRaises(ExpectedError, gulp, s)
511
512 # keyfunc failure
513 def keyfunc(obj):
514 if keyfunc.skip > 0:
515 keyfunc.skip -= 1
516 return obj
517 else:
518 raise ExpectedError
519
520 # keyfunc failure on outer object
521 keyfunc.skip = 0
522 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
523 keyfunc.skip = 1
524 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
525
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000526 def test_filter(self):
527 self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
528 self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])
529 self.assertEqual(list(filter(bool, [0,1,0,2,0])), [1,2])
530 self.assertEqual(take(4, filter(isEven, count())), [0,2,4,6])
531 self.assertRaises(TypeError, filter)
532 self.assertRaises(TypeError, filter, lambda x:x)
533 self.assertRaises(TypeError, filter, lambda x:x, range(6), 7)
534 self.assertRaises(TypeError, filter, isEven, 3)
535 self.assertRaises(TypeError, next, filter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000536
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000537 def test_filterfalse(self):
538 self.assertEqual(list(filterfalse(isEven, range(6))), [1,3,5])
539 self.assertEqual(list(filterfalse(None, [0,1,0,2,0])), [0,0,0])
540 self.assertEqual(list(filterfalse(bool, [0,1,0,2,0])), [0,0,0])
541 self.assertEqual(take(4, filterfalse(isEven, count())), [1,3,5,7])
542 self.assertRaises(TypeError, filterfalse)
543 self.assertRaises(TypeError, filterfalse, lambda x:x)
544 self.assertRaises(TypeError, filterfalse, lambda x:x, range(6), 7)
545 self.assertRaises(TypeError, filterfalse, isEven, 3)
546 self.assertRaises(TypeError, next, filterfalse(range(6), range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000547
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000548 def test_zip(self):
549 # XXX This is rather silly now that builtin zip() calls zip()...
550 ans = [(x,y) for x, y in zip('abc',count())]
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000551 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000552 self.assertEqual(list(zip('abc', range(6))), lzip('abc', range(6)))
553 self.assertEqual(list(zip('abcdef', range(3))), lzip('abcdef', range(3)))
554 self.assertEqual(take(3,zip('abcdef', count())), lzip('abcdef', range(3)))
555 self.assertEqual(list(zip('abcdef')), lzip('abcdef'))
556 self.assertEqual(list(zip()), lzip())
557 self.assertRaises(TypeError, zip, 3)
558 self.assertRaises(TypeError, zip, range(3), 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000559 # Check tuple re-use (implementation detail)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000560 self.assertEqual([tuple(list(pair)) for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000561 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000562 self.assertEqual([pair for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000563 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000564 ids = list(map(id, zip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000565 self.assertEqual(min(ids), max(ids))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000566 ids = list(map(id, list(zip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000567 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000568
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000569 def test_ziplongest(self):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000570 for args in [
571 ['abc', range(6)],
572 [range(6), 'abc'],
573 [range(1000), range(2000,2100), range(3000,3050)],
574 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
575 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
576 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000577 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
578 for i in range(max(map(len, args)))]
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000579 self.assertEqual(list(zip_longest(*args)), target)
580 self.assertEqual(list(zip_longest(*args, **{})), target)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000581 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 +0000582 self.assertEqual(list(zip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000583
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000584 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 +0000585
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000586 self.assertEqual(list(zip_longest()), list(zip()))
587 self.assertEqual(list(zip_longest([])), list(zip([])))
588 self.assertEqual(list(zip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000589
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000590 self.assertEqual(list(zip_longest('abc', 'defg', **{})),
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000591 list(zip(list('abc')+[None], 'defg'))) # empty keyword dict
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000592 self.assertRaises(TypeError, zip_longest, 3)
593 self.assertRaises(TypeError, zip_longest, range(3), 3)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000594
595 for stmt in [
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000596 "zip_longest('abc', fv=1)",
597 "zip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000598 ]:
599 try:
600 eval(stmt, globals(), locals())
601 except TypeError:
602 pass
603 else:
604 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000605
Thomas Wouterscf297e42007-02-23 15:07:44 +0000606 # Check tuple re-use (implementation detail)
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000607 self.assertEqual([tuple(list(pair)) for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000608 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000609 self.assertEqual([pair for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000610 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000611 ids = list(map(id, zip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000612 self.assertEqual(min(ids), max(ids))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000613 ids = list(map(id, list(zip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000614 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
615
Raymond Hettingerfc438512009-11-01 20:55:33 +0000616 def test_bug_7244(self):
617
618 class Repeater:
619 # this class is similar to itertools.repeat
620 def __init__(self, o, t, e):
621 self.o = o
622 self.t = int(t)
623 self.e = e
624 def __iter__(self): # its iterator is itself
625 return self
626 def __next__(self):
627 if self.t > 0:
628 self.t -= 1
629 return self.o
630 else:
631 raise self.e
632
633 # Formerly this code in would fail in debug mode
634 # with Undetected Error and Stop Iteration
635 r1 = Repeater(1, 3, StopIteration)
636 r2 = Repeater(2, 4, StopIteration)
637 def run(r1, r2):
638 result = []
639 for i, j in zip_longest(r1, r2, fillvalue=0):
640 with support.captured_output('stdout'):
641 print((i, j))
642 result.append((i, j))
643 return result
644 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
645
646 # Formerly, the RuntimeError would be lost
647 # and StopIteration would stop as expected
648 r1 = Repeater(1, 3, RuntimeError)
649 r2 = Repeater(2, 4, StopIteration)
650 it = zip_longest(r1, r2, fillvalue=0)
651 self.assertEqual(next(it), (1, 2))
652 self.assertEqual(next(it), (1, 2))
653 self.assertEqual(next(it), (1, 2))
654 self.assertRaises(RuntimeError, next, it)
655
Christian Heimesc3f30c42008-02-22 16:37:40 +0000656 def test_product(self):
657 for args, result in [
Christian Heimes78644762008-03-04 23:39:23 +0000658 ([], [()]), # zero iterables
Christian Heimesc3f30c42008-02-22 16:37:40 +0000659 (['ab'], [('a',), ('b',)]), # one iterable
660 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
661 ([range(0), range(2), range(3)], []), # first iterable with zero length
662 ([range(2), range(0), range(3)], []), # middle iterable with zero length
663 ([range(2), range(3), range(0)], []), # last iterable with zero length
664 ]:
665 self.assertEqual(list(product(*args)), result)
Christian Heimesf16baeb2008-02-29 14:57:44 +0000666 for r in range(4):
667 self.assertEqual(list(product(*(args*r))),
668 list(product(*args, **dict(repeat=r))))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000669 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
670 self.assertRaises(TypeError, product, range(6), None)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000671
672 def product1(*args, **kwds):
673 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
674 n = len(pools)
675 if n == 0:
676 yield ()
677 return
678 if any(len(pool) == 0 for pool in pools):
679 return
680 indices = [0] * n
681 yield tuple(pool[i] for pool, i in zip(pools, indices))
682 while 1:
683 for i in reversed(range(n)): # right to left
684 if indices[i] == len(pools[i]) - 1:
685 continue
686 indices[i] += 1
687 for j in range(i+1, n):
688 indices[j] = 0
689 yield tuple(pool[i] for pool, i in zip(pools, indices))
690 break
691 else:
692 return
693
694 def product2(*args, **kwds):
695 'Pure python version used in docs'
696 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
697 result = [[]]
698 for pool in pools:
699 result = [x+[y] for x in result for y in pool]
700 for prod in result:
701 yield tuple(prod)
702
Christian Heimesc3f30c42008-02-22 16:37:40 +0000703 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
704 set('abcdefg'), range(11), tuple(range(13))]
705 for i in range(100):
706 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Christian Heimes78644762008-03-04 23:39:23 +0000707 expected_len = prod(map(len, args))
708 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000709 self.assertEqual(list(product(*args)), list(product1(*args)))
710 self.assertEqual(list(product(*args)), list(product2(*args)))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000711 args = map(iter, args)
Christian Heimes78644762008-03-04 23:39:23 +0000712 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000713
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000714 # Test implementation detail: tuple re-use
715 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
716 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000717
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000718 def test_repeat(self):
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +0000719 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
Guido van Rossum805365e2007-05-07 22:24:25 +0000720 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000721 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000722 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000723 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000724 self.assertEqual(list(repeat('a', 0)), [])
725 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000726 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000727 self.assertRaises(TypeError, repeat, None, 3, 4)
728 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000729 r = repeat(1+0j)
730 self.assertEqual(repr(r), 'repeat((1+0j))')
731 r = repeat(1+0j, 5)
732 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
733 list(r)
734 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000735
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000736 def test_map(self):
737 self.assertEqual(list(map(operator.pow, range(3), range(1,7))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000738 [0**1, 1**2, 2**3])
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000739 def tupleize(*args):
740 return args
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000741 self.assertEqual(list(map(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000742 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000743 self.assertEqual(list(map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000744 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000745 self.assertEqual(take(2,map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000746 [('a',0),('b',1)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000747 self.assertEqual(list(map(operator.pow, [])), [])
748 self.assertRaises(TypeError, map)
749 self.assertRaises(TypeError, list, map(None, range(3), range(3)))
750 self.assertRaises(TypeError, map, operator.neg)
751 self.assertRaises(TypeError, next, map(10, range(5)))
752 self.assertRaises(ValueError, next, map(errfunc, [4], [5]))
753 self.assertRaises(TypeError, next, map(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000754
755 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000756 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
757 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000758 self.assertEqual(take(3, starmap(operator.pow, zip(count(), count(1)))),
Raymond Hettinger02420702003-06-29 20:36:23 +0000759 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000760 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +0000761 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
762 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000763 self.assertRaises(TypeError, starmap)
764 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000765 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
766 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
767 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000768
769 def test_islice(self):
770 for args in [ # islice(args) should agree with range(args)
771 (10, 20, 3),
772 (10, 3, 20),
773 (10, 20),
774 (10, 3),
775 (20,)
776 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000777 self.assertEqual(list(islice(range(100), *args)),
778 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000779
780 for args, tgtargs in [ # Stop when seqn is exhausted
781 ((10, 110, 3), ((10, 100, 3))),
782 ((10, 110), ((10, 100))),
783 ((110,), (100,))
784 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000785 self.assertEqual(list(islice(range(100), *args)),
786 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000787
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000788 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +0000789 self.assertEqual(list(islice(range(10), None)), list(range(10)))
790 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
791 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
792 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
793 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000794
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000795 # Test number of items consumed SF #1171417
796 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +0000797 self.assertEqual(list(islice(it, 3)), list(range(3)))
798 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000799
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000800 # Test invalid arguments
Guido van Rossum805365e2007-05-07 22:24:25 +0000801 self.assertRaises(TypeError, islice, range(10))
802 self.assertRaises(TypeError, islice, range(10), 1, 2, 3, 4)
803 self.assertRaises(ValueError, islice, range(10), -5, 10, 1)
804 self.assertRaises(ValueError, islice, range(10), 1, -5, -1)
805 self.assertRaises(ValueError, islice, range(10), 1, 10, -1)
806 self.assertRaises(ValueError, islice, range(10), 1, 10, 0)
807 self.assertRaises(ValueError, islice, range(10), 'a')
808 self.assertRaises(ValueError, islice, range(10), 'a', 1)
809 self.assertRaises(ValueError, islice, range(10), 1, 'a')
810 self.assertRaises(ValueError, islice, range(10), 'a', 1, 1)
811 self.assertRaises(ValueError, islice, range(10), 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +0000812 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000813
Raymond Hettinger69b34bf2010-11-30 02:49:29 +0000814 # Issue #10323: Less islice in a predictable state
815 c = count()
816 self.assertEqual(list(islice(c, 1, 3, 50)), [1])
817 self.assertEqual(next(c), 3)
818
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000819 def test_takewhile(self):
820 data = [1, 3, 5, 20, 2, 4, 6, 8]
821 underten = lambda x: x<10
822 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000823 self.assertEqual(list(takewhile(underten, [])), [])
824 self.assertRaises(TypeError, takewhile)
825 self.assertRaises(TypeError, takewhile, operator.pow)
826 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000827 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
828 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000829 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
830 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +0000831 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000832
833 def test_dropwhile(self):
834 data = [1, 3, 5, 20, 2, 4, 6, 8]
835 underten = lambda x: x<10
836 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000837 self.assertEqual(list(dropwhile(underten, [])), [])
838 self.assertRaises(TypeError, dropwhile)
839 self.assertRaises(TypeError, dropwhile, operator.pow)
840 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000841 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
842 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000843
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000844 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000845 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000846 def irange(n):
Guido van Rossum805365e2007-05-07 22:24:25 +0000847 for i in range(n):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000848 yield i
849
850 a, b = tee([]) # test empty iterator
851 self.assertEqual(list(a), [])
852 self.assertEqual(list(b), [])
853
854 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +0000855 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000856
857 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +0000858 self.assertEqual(list(a), list(range(n)))
859 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000860
861 a, b = tee(irange(n)) # test dealloc of leading 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 a
Guido van Rossum805365e2007-05-07 22:24:25 +0000865 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000866
867 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000868 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000869 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000870 del b
Guido van Rossum805365e2007-05-07 22:24:25 +0000871 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000872
Guido van Rossum805365e2007-05-07 22:24:25 +0000873 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000874 order = [0]*n + [1]*n
875 random.shuffle(order)
876 lists = ([], [])
877 its = tee(irange(n))
878 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +0000879 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000880 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +0000881 self.assertEqual(lists[0], list(range(n)))
882 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000883
Raymond Hettingerad983e72003-11-12 14:32:26 +0000884 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000885 self.assertRaises(TypeError, tee)
886 self.assertRaises(TypeError, tee, 3)
887 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000888 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000889
Raymond Hettingerad983e72003-11-12 14:32:26 +0000890 # tee object should be instantiable
891 a, b = tee('abc')
892 c = type(a)('def')
893 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000894
Raymond Hettingerad983e72003-11-12 14:32:26 +0000895 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +0000896 a, b, c = tee(range(2000), 3)
897 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000898 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +0000899 self.assertEqual(list(b), list(range(2000)))
900 self.assertEqual([next(c), next(c)], list(range(2)))
901 self.assertEqual(list(a), list(range(100,2000)))
902 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000903
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000904 # test values of n
905 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000906 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +0000907 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000908 result = tee('abc', n)
909 self.assertEqual(type(result), tuple)
910 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000911 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000912
Raymond Hettingerad983e72003-11-12 14:32:26 +0000913 # tee pass-through to copyable iterator
914 a, b = tee('abc')
915 c, d = tee(a)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000916 self.assertTrue(a is c)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000917
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000918 # test tee_new
919 t1, t2 = tee('abc')
920 tnew = type(t1)
921 self.assertRaises(TypeError, tnew)
922 self.assertRaises(TypeError, tnew, 10)
923 t3 = tnew(t1)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000924 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000925
Raymond Hettingera9f60922004-10-17 16:40:14 +0000926 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +0000927 a, b = tee(range(10))
Raymond Hettingera9f60922004-10-17 16:40:14 +0000928 p = proxy(a)
929 self.assertEqual(getattr(p, '__class__'), type(b))
930 del a
931 self.assertRaises(ReferenceError, getattr, p, '__class__')
932
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200933 # Issue 13454: Crash when deleting backward iterator from tee()
934 def test_tee_del_backward(self):
935 forward, backward = tee(range(20000000))
936 for i in forward:
937 pass
938
939 del backward
940
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000941 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000942 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000943
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000944 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +0000945 self.assertRaises(StopIteration, next, f([]))
946 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000947
Georg Brandla18af4e2007-04-21 15:47:16 +0000948 self.assertRaises(StopIteration, next, islice([], None))
949 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000950
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000951 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +0000952 self.assertRaises(StopIteration, next, p)
953 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000954 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +0000955 self.assertRaises(StopIteration, next, p)
956 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000957
Georg Brandla18af4e2007-04-21 15:47:16 +0000958 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000959
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000960 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +0000961 self.assertRaises(StopIteration, next, f(lambda x:x, []))
962 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000963
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000964class TestExamples(unittest.TestCase):
965
Raymond Hettinger482ba772010-12-01 22:48:00 +0000966 def test_accumlate(self):
967 self.assertEqual(list(accumulate([1,2,3,4,5])), [1, 3, 6, 10, 15])
968
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000969 def test_chain(self):
970 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
971
972 def test_chain_from_iterable(self):
973 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
974
975 def test_combinations(self):
976 self.assertEqual(list(combinations('ABCD', 2)),
977 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
978 self.assertEqual(list(combinations(range(4), 3)),
979 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
980
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000981 def test_combinations_with_replacement(self):
982 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
983 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
984
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000985 def test_compress(self):
986 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
987
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000988 def test_count(self):
989 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
990
991 def test_cycle(self):
992 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
993
994 def test_dropwhile(self):
995 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
996
997 def test_groupby(self):
998 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
999 list('ABCDAB'))
1000 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
1001 [list('AAAA'), list('BBB'), list('CC'), list('D')])
1002
1003 def test_filter(self):
1004 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
1005
1006 def test_filterfalse(self):
1007 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
1008
1009 def test_map(self):
1010 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
1011
1012 def test_islice(self):
1013 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
1014 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
1015 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
1016 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1017
1018 def test_zip(self):
1019 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
1020
1021 def test_zip_longest(self):
1022 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
1023 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
1024
1025 def test_permutations(self):
1026 self.assertEqual(list(permutations('ABCD', 2)),
1027 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
1028 self.assertEqual(list(permutations(range(3))),
1029 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
1030
1031 def test_product(self):
1032 self.assertEqual(list(product('ABCD', 'xy')),
1033 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
1034 self.assertEqual(list(product(range(2), repeat=3)),
1035 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
1036 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
1037
1038 def test_repeat(self):
1039 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
1040
1041 def test_stapmap(self):
1042 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
1043 [32, 9, 1000])
1044
1045 def test_takewhile(self):
1046 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
1047
1048
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001049class TestGC(unittest.TestCase):
1050
1051 def makecycle(self, iterator, container):
1052 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +00001053 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001054 del container, iterator
1055
Raymond Hettinger482ba772010-12-01 22:48:00 +00001056 def test_accumulate(self):
1057 a = []
1058 self.makecycle(accumulate([1,2,a,3]), a)
1059
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001060 def test_chain(self):
1061 a = []
1062 self.makecycle(chain(a), a)
1063
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001064 def test_chain_from_iterable(self):
1065 a = []
1066 self.makecycle(chain.from_iterable([a]), a)
1067
1068 def test_combinations(self):
1069 a = []
1070 self.makecycle(combinations([1,2,a,3], 3), a)
1071
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001072 def test_combinations_with_replacement(self):
1073 a = []
1074 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1075
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001076 def test_compress(self):
1077 a = []
1078 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1079
Raymond Hettingerd6280f42009-02-16 20:50:56 +00001080 def test_count(self):
1081 a = []
1082 Int = type('Int', (int,), dict(x=a))
1083 self.makecycle(count(Int(0), Int(1)), a)
1084
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001085 def test_cycle(self):
1086 a = []
1087 self.makecycle(cycle([a]*2), a)
1088
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001089 def test_dropwhile(self):
1090 a = []
1091 self.makecycle(dropwhile(bool, [0, a, a]), a)
1092
1093 def test_groupby(self):
1094 a = []
1095 self.makecycle(groupby([a]*2, lambda x:x), a)
1096
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001097 def test_issue2246(self):
1098 # Issue 2246 -- the _grouper iterator was not included in GC
1099 n = 10
1100 keyfunc = lambda x: x
1101 for i, j in groupby(range(n), key=keyfunc):
1102 keyfunc.__dict__.setdefault('x',[]).append(j)
1103
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001104 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001105 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001106 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001107
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001108 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001109 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001110 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001111
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001112 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001113 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001114 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001115
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001116 def test_zip_longest(self):
1117 a = []
1118 self.makecycle(zip_longest([a]*2, [a]*3), a)
1119 b = [a, None]
1120 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1121
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001122 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001123 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001124 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001125
1126 def test_islice(self):
1127 a = []
1128 self.makecycle(islice([a]*2, None), a)
1129
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001130 def test_permutations(self):
1131 a = []
1132 self.makecycle(permutations([1,2,a,3], 3), a)
1133
1134 def test_product(self):
1135 a = []
1136 self.makecycle(product([1,2,a,3], repeat=3), a)
1137
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001138 def test_repeat(self):
1139 a = []
1140 self.makecycle(repeat(a), a)
1141
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001142 def test_starmap(self):
1143 a = []
1144 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1145
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001146 def test_takewhile(self):
1147 a = []
1148 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1149
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001150def R(seqn):
1151 'Regular generator'
1152 for i in seqn:
1153 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001154
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001155class G:
1156 'Sequence using __getitem__'
1157 def __init__(self, seqn):
1158 self.seqn = seqn
1159 def __getitem__(self, i):
1160 return self.seqn[i]
1161
1162class I:
1163 'Sequence using iterator protocol'
1164 def __init__(self, seqn):
1165 self.seqn = seqn
1166 self.i = 0
1167 def __iter__(self):
1168 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001169 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001170 if self.i >= len(self.seqn): raise StopIteration
1171 v = self.seqn[self.i]
1172 self.i += 1
1173 return v
1174
1175class Ig:
1176 'Sequence using iterator protocol defined with a generator'
1177 def __init__(self, seqn):
1178 self.seqn = seqn
1179 self.i = 0
1180 def __iter__(self):
1181 for val in self.seqn:
1182 yield val
1183
1184class X:
1185 'Missing __getitem__ and __iter__'
1186 def __init__(self, seqn):
1187 self.seqn = seqn
1188 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001189 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001190 if self.i >= len(self.seqn): raise StopIteration
1191 v = self.seqn[self.i]
1192 self.i += 1
1193 return v
1194
1195class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001196 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001197 def __init__(self, seqn):
1198 self.seqn = seqn
1199 self.i = 0
1200 def __iter__(self):
1201 return self
1202
1203class E:
1204 'Test propagation of exceptions'
1205 def __init__(self, seqn):
1206 self.seqn = seqn
1207 self.i = 0
1208 def __iter__(self):
1209 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001210 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001211 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001212
1213class S:
1214 'Test immediate stop'
1215 def __init__(self, seqn):
1216 pass
1217 def __iter__(self):
1218 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001219 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001220 raise StopIteration
1221
1222def L(seqn):
1223 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001224 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001225
1226
1227class TestVariousIteratorArgs(unittest.TestCase):
1228
Raymond Hettinger482ba772010-12-01 22:48:00 +00001229 def test_accumulate(self):
1230 s = [1,2,3,4,5]
1231 r = [1,3,6,10,15]
1232 n = len(s)
1233 for g in (G, I, Ig, L, R):
1234 self.assertEqual(list(accumulate(g(s))), r)
1235 self.assertEqual(list(accumulate(S(s))), [])
1236 self.assertRaises(TypeError, accumulate, X(s))
1237 self.assertRaises(TypeError, accumulate, N(s))
1238 self.assertRaises(ZeroDivisionError, list, accumulate(E(s)))
1239
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001240 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001241 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001242 for g in (G, I, Ig, S, L, R):
1243 self.assertEqual(list(chain(g(s))), list(g(s)))
1244 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001245 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001246 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001247 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1248
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001249 def test_compress(self):
1250 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1251 n = len(s)
1252 for g in (G, I, Ig, S, L, R):
1253 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1254 self.assertRaises(TypeError, compress, X(s), repeat(1))
1255 self.assertRaises(TypeError, compress, N(s), repeat(1))
1256 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1257
Christian Heimesc3f30c42008-02-22 16:37:40 +00001258 def test_product(self):
1259 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1260 self.assertRaises(TypeError, product, X(s))
1261 self.assertRaises(TypeError, product, N(s))
1262 self.assertRaises(ZeroDivisionError, product, E(s))
1263
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001264 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001265 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001266 for g in (G, I, Ig, S, L, R):
1267 tgtlen = len(s) * 3
1268 expected = list(g(s))*3
1269 actual = list(islice(cycle(g(s)), tgtlen))
1270 self.assertEqual(actual, expected)
1271 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001272 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001273 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1274
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001275 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001276 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001277 for g in (G, I, Ig, S, L, R):
1278 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1279 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001280 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001281 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1282
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001283 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001284 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001285 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001286 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001287 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001288 self.assertRaises(TypeError, filter, isEven, X(s))
1289 self.assertRaises(TypeError, filter, isEven, N(s))
1290 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001291
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001292 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001293 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001294 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001295 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001296 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001297 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1298 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1299 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001300
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001301 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001302 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001303 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001304 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1305 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1306 self.assertRaises(TypeError, zip, X(s))
1307 self.assertRaises(TypeError, zip, N(s))
1308 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001309
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001310 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001311 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001312 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001313 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1314 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1315 self.assertRaises(TypeError, zip_longest, X(s))
1316 self.assertRaises(TypeError, zip_longest, N(s))
1317 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001318
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001319 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001320 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001321 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001322 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001323 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001324 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001325 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001326 self.assertRaises(TypeError, map, onearg, X(s))
1327 self.assertRaises(TypeError, map, onearg, N(s))
1328 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001329
1330 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001331 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001332 for g in (G, I, Ig, S, L, R):
1333 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1334 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001335 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001336 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1337
1338 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001339 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001340 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001341 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001342 self.assertEqual(list(starmap(operator.pow, g(ss))),
1343 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001344 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001345 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001346 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1347
1348 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001349 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001350 for g in (G, I, Ig, S, L, R):
1351 tgt = []
1352 for elem in g(s):
1353 if not isEven(elem): break
1354 tgt.append(elem)
1355 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1356 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001357 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001358 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1359
1360 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001361 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001362 for g in (G, I, Ig, S, L, R):
1363 tgt = []
1364 for elem in g(s):
1365 if not tgt and isOdd(elem): continue
1366 tgt.append(elem)
1367 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1368 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001369 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001370 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1371
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001372 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001373 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001374 for g in (G, I, Ig, S, L, R):
1375 it1, it2 = tee(g(s))
1376 self.assertEqual(list(it1), list(g(s)))
1377 self.assertEqual(list(it2), list(g(s)))
1378 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001379 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001380 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1381
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001382class LengthTransparency(unittest.TestCase):
1383
1384 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001385 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001386 self.assertEqual(len(repeat(None, 50)), 50)
1387 self.assertRaises(TypeError, len, repeat(None))
1388
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001389class RegressionTests(unittest.TestCase):
1390
1391 def test_sf_793826(self):
1392 # Fix Armin Rigo's successful efforts to wreak havoc
1393
1394 def mutatingtuple(tuple1, f, tuple2):
1395 # this builds a tuple t which is a copy of tuple1,
1396 # then calls f(t), then mutates t to be equal to tuple2
1397 # (needs len(tuple1) == len(tuple2)).
1398 def g(value, first=[1]):
1399 if first:
1400 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00001401 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001402 return value
1403 items = list(tuple2)
1404 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001405 gen = map(g, items)
1406 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00001407 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001408
1409 def f(t):
1410 global T
1411 T = t
1412 first[:] = list(T)
1413
1414 first = []
1415 mutatingtuple((1,2,3), f, (4,5,6))
1416 second = list(T)
1417 self.assertEqual(first, second)
1418
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001419
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001420 def test_sf_950057(self):
1421 # Make sure that chain() and cycle() catch exceptions immediately
1422 # rather than when shifting between input sources
1423
1424 def gen1():
1425 hist.append(0)
1426 yield 1
1427 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001428 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001429 hist.append(2)
1430
1431 def gen2(x):
1432 hist.append(3)
1433 yield 2
1434 hist.append(4)
1435 if x:
1436 raise StopIteration
1437
1438 hist = []
1439 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1440 self.assertEqual(hist, [0,1])
1441
1442 hist = []
1443 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1444 self.assertEqual(hist, [0,1])
1445
1446 hist = []
1447 self.assertRaises(AssertionError, list, cycle(gen1()))
1448 self.assertEqual(hist, [0,1])
1449
Thomas Woutersb2137042007-02-01 18:02:27 +00001450class SubclassWithKwargsTest(unittest.TestCase):
1451 def test_keywords_in_subclass(self):
1452 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001453 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001454 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00001455 class Subclass(cls):
1456 def __init__(self, newarg=None, *args):
1457 cls.__init__(self, *args)
1458 try:
1459 Subclass(newarg=1)
1460 except TypeError as err:
1461 # we expect type errors because of wrong argument count
Ezio Melottib58e0bd2010-01-23 15:40:09 +00001462 self.assertNotIn("does not take keyword arguments", err.args[0])
Thomas Woutersb2137042007-02-01 18:02:27 +00001463
1464
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001465libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001466
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001467
1468>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001469>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00001470... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001471...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001472Check 1200 is for $120.15
1473Check 1201 is for $764.05
1474Check 1202 is for $823.14
1475
1476>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001477>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00001478... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00001479...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000014801
14818
148227
1483
1484>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001485>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00001486... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00001487...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001488Alex
1489Laura
1490Martin
1491Walter
1492Samuele
1493
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001494>>> from operator import itemgetter
1495>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00001496>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001497>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001498... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001499...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000015001 ['a', 'c', 'e']
15012 ['b', 'd', 'f']
15023 ['g']
1503
Raymond Hettinger734fb572004-01-20 20:04:40 +00001504# Find runs of consecutive numbers using groupby. The key to the solution
1505# is differencing with a range so that consecutive numbers all appear in
1506# same group.
1507>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00001508>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001509... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001510...
Raymond Hettinger734fb572004-01-20 20:04:40 +00001511[1]
1512[4, 5, 6]
1513[10]
1514[15, 16, 17, 18]
1515[22]
1516[25, 26, 27, 28]
1517
Georg Brandl3dbca812008-07-23 16:10:53 +00001518>>> def take(n, iterable):
1519... "Return first n items of the iterable as a list"
1520... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001521
Georg Brandl3dbca812008-07-23 16:10:53 +00001522>>> def enumerate(iterable, start=0):
1523... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001524
Georg Brandl3dbca812008-07-23 16:10:53 +00001525>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001526... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00001527... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001528
Raymond Hettingercdf8ba32009-02-19 04:45:07 +00001529>>> def nth(iterable, n, default=None):
1530... "Returns the nth item or a default value"
1531... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001532
Georg Brandl3dbca812008-07-23 16:10:53 +00001533>>> def quantify(iterable, pred=bool):
1534... "Count how many times the predicate is true"
1535... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001536
Georg Brandl3dbca812008-07-23 16:10:53 +00001537>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001538... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00001539... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001540
Georg Brandl3dbca812008-07-23 16:10:53 +00001541>>> def ncycles(iterable, n):
Ezio Melotti13925002011-03-16 11:05:33 +02001542... "Returns the sequence elements n times"
Georg Brandl3dbca812008-07-23 16:10:53 +00001543... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001544
1545>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001546... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001547
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001548>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001549... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001550
1551>>> def repeatfunc(func, times=None, *args):
1552... "Repeat calls to func with specified arguments."
1553... " Example: repeatfunc(random.random)"
1554... if times is None:
1555... return starmap(func, repeat(args))
1556... else:
1557... return starmap(func, repeat(args, times))
1558
Raymond Hettingerd591f662003-10-26 15:34:50 +00001559>>> def pairwise(iterable):
1560... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1561... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001562... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00001563... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001564... except StopIteration:
1565... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001566... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001567
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001568>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001569... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001570... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00001571... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001572
1573>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001574... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001575... # Recipe credited to George Sakkis
1576... pending = len(iterables)
1577... nexts = cycle(iter(it).__next__ for it in iterables)
1578... while pending:
1579... try:
1580... for next in nexts:
1581... yield next()
1582... except StopIteration:
1583... pending -= 1
1584... nexts = cycle(islice(nexts, pending))
1585
1586>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00001587... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1588... s = list(iterable)
1589... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001590
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001591>>> def unique_everseen(iterable, key=None):
1592... "List unique elements, preserving order. Remember all elements ever seen."
1593... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1594... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1595... seen = set()
1596... seen_add = seen.add
1597... if key is None:
1598... for element in iterable:
1599... if element not in seen:
1600... seen_add(element)
1601... yield element
1602... else:
1603... for element in iterable:
1604... k = key(element)
1605... if k not in seen:
1606... seen_add(k)
1607... yield element
1608
1609>>> def unique_justseen(iterable, key=None):
1610... "List unique elements, preserving order. Remember only the element just seen."
1611... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1612... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1613... return map(next, map(itemgetter(1), groupby(iterable, key)))
1614
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001615This is not part of the examples but it tests to make sure the definitions
1616perform as purported.
1617
Raymond Hettingera098b332003-09-08 23:58:40 +00001618>>> take(10, count())
1619[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1620
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001621>>> list(enumerate('abc'))
1622[(0, 'a'), (1, 'b'), (2, 'c')]
1623
1624>>> list(islice(tabulate(lambda x: 2*x), 4))
1625[0, 2, 4, 6]
1626
1627>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00001628'd'
1629
1630>>> nth('abcde', 9) is None
1631True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001632
Guido van Rossum805365e2007-05-07 22:24:25 +00001633>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000163450
1635
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001636>>> a = [[1, 2, 3], [4, 5, 6]]
1637>>> flatten(a)
1638[1, 2, 3, 4, 5, 6]
1639
1640>>> list(repeatfunc(pow, 5, 2, 3))
1641[8, 8, 8, 8, 8]
1642
1643>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001644>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001645[0, 0, 0, 0, 0]
1646
Raymond Hettingerd591f662003-10-26 15:34:50 +00001647>>> list(pairwise('abcd'))
1648[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001649
Raymond Hettingerd591f662003-10-26 15:34:50 +00001650>>> list(pairwise([]))
1651[]
1652
1653>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001654[]
1655
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001656>>> list(islice(padnone('abc'), 0, 6))
1657['a', 'b', 'c', None, None, None]
1658
1659>>> list(ncycles('abc', 3))
1660['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1661
1662>>> dotproduct([1,2,3], [4,5,6])
166332
1664
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001665>>> list(grouper(3, 'abcdefg', 'x'))
1666[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1667
1668>>> list(roundrobin('abc', 'd', 'ef'))
1669['a', 'd', 'e', 'b', 'f', 'c']
1670
Raymond Hettingerace67332009-01-26 02:23:50 +00001671>>> list(powerset([1,2,3]))
1672[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001673
Raymond Hettinger191e8502009-01-27 13:29:43 +00001674>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1675True
1676
1677>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1678True
1679
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001680>>> list(unique_everseen('AAAABBBCCDAABBB'))
1681['A', 'B', 'C', 'D']
1682
1683>>> list(unique_everseen('ABBCcAD', str.lower))
1684['A', 'B', 'C', 'D']
1685
1686>>> list(unique_justseen('AAAABBBCCDAABBB'))
1687['A', 'B', 'C', 'D', 'A', 'B']
1688
1689>>> list(unique_justseen('ABBCcAD', str.lower))
1690['A', 'B', 'C', 'A', 'D']
1691
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001692"""
1693
1694__test__ = {'libreftest' : libreftest}
1695
1696def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001697 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00001698 RegressionTests, LengthTransparency,
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001699 SubclassWithKwargsTest, TestExamples)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001700 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001701
1702 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001703 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001704 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001705 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00001706 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001707 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001708 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001709 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001710 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001711
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001712 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001713 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001714
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001715if __name__ == "__main__":
1716 test_main(verbose=True)