blob: dfa371e2b8665e8f98a6955d7b209c51f2e8a9df [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):
Serhiy Storchaka5bb893c2013-01-26 11:52:06 +0200935 forward, backward = tee(repeat(None, 20000000))
936 any(forward) # exhaust the iterator
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200937 del backward
938
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000939 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000940 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000941
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000942 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +0000943 self.assertRaises(StopIteration, next, f([]))
944 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000945
Georg Brandla18af4e2007-04-21 15:47:16 +0000946 self.assertRaises(StopIteration, next, islice([], None))
947 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000948
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000949 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +0000950 self.assertRaises(StopIteration, next, p)
951 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000952 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +0000953 self.assertRaises(StopIteration, next, p)
954 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000955
Georg Brandla18af4e2007-04-21 15:47:16 +0000956 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000957
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000958 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +0000959 self.assertRaises(StopIteration, next, f(lambda x:x, []))
960 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000961
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000962class TestExamples(unittest.TestCase):
963
Raymond Hettinger482ba772010-12-01 22:48:00 +0000964 def test_accumlate(self):
965 self.assertEqual(list(accumulate([1,2,3,4,5])), [1, 3, 6, 10, 15])
966
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000967 def test_chain(self):
968 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
969
970 def test_chain_from_iterable(self):
971 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
972
973 def test_combinations(self):
974 self.assertEqual(list(combinations('ABCD', 2)),
975 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
976 self.assertEqual(list(combinations(range(4), 3)),
977 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
978
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000979 def test_combinations_with_replacement(self):
980 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
981 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
982
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000983 def test_compress(self):
984 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
985
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000986 def test_count(self):
987 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
988
989 def test_cycle(self):
990 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
991
992 def test_dropwhile(self):
993 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
994
995 def test_groupby(self):
996 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
997 list('ABCDAB'))
998 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
999 [list('AAAA'), list('BBB'), list('CC'), list('D')])
1000
1001 def test_filter(self):
1002 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
1003
1004 def test_filterfalse(self):
1005 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
1006
1007 def test_map(self):
1008 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
1009
1010 def test_islice(self):
1011 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
1012 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
1013 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
1014 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1015
1016 def test_zip(self):
1017 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
1018
1019 def test_zip_longest(self):
1020 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
1021 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
1022
1023 def test_permutations(self):
1024 self.assertEqual(list(permutations('ABCD', 2)),
1025 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
1026 self.assertEqual(list(permutations(range(3))),
1027 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
1028
1029 def test_product(self):
1030 self.assertEqual(list(product('ABCD', 'xy')),
1031 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
1032 self.assertEqual(list(product(range(2), repeat=3)),
1033 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
1034 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
1035
1036 def test_repeat(self):
1037 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
1038
1039 def test_stapmap(self):
1040 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
1041 [32, 9, 1000])
1042
1043 def test_takewhile(self):
1044 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
1045
1046
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001047class TestGC(unittest.TestCase):
1048
1049 def makecycle(self, iterator, container):
1050 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +00001051 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001052 del container, iterator
1053
Raymond Hettinger482ba772010-12-01 22:48:00 +00001054 def test_accumulate(self):
1055 a = []
1056 self.makecycle(accumulate([1,2,a,3]), a)
1057
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001058 def test_chain(self):
1059 a = []
1060 self.makecycle(chain(a), a)
1061
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001062 def test_chain_from_iterable(self):
1063 a = []
1064 self.makecycle(chain.from_iterable([a]), a)
1065
1066 def test_combinations(self):
1067 a = []
1068 self.makecycle(combinations([1,2,a,3], 3), a)
1069
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001070 def test_combinations_with_replacement(self):
1071 a = []
1072 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1073
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001074 def test_compress(self):
1075 a = []
1076 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1077
Raymond Hettingerd6280f42009-02-16 20:50:56 +00001078 def test_count(self):
1079 a = []
1080 Int = type('Int', (int,), dict(x=a))
1081 self.makecycle(count(Int(0), Int(1)), a)
1082
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001083 def test_cycle(self):
1084 a = []
1085 self.makecycle(cycle([a]*2), a)
1086
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001087 def test_dropwhile(self):
1088 a = []
1089 self.makecycle(dropwhile(bool, [0, a, a]), a)
1090
1091 def test_groupby(self):
1092 a = []
1093 self.makecycle(groupby([a]*2, lambda x:x), a)
1094
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001095 def test_issue2246(self):
1096 # Issue 2246 -- the _grouper iterator was not included in GC
1097 n = 10
1098 keyfunc = lambda x: x
1099 for i, j in groupby(range(n), key=keyfunc):
1100 keyfunc.__dict__.setdefault('x',[]).append(j)
1101
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001102 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001103 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001104 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001105
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001106 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001107 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001108 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001109
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001110 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001111 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001112 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001113
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001114 def test_zip_longest(self):
1115 a = []
1116 self.makecycle(zip_longest([a]*2, [a]*3), a)
1117 b = [a, None]
1118 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1119
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001120 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001121 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001122 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001123
1124 def test_islice(self):
1125 a = []
1126 self.makecycle(islice([a]*2, None), a)
1127
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001128 def test_permutations(self):
1129 a = []
1130 self.makecycle(permutations([1,2,a,3], 3), a)
1131
1132 def test_product(self):
1133 a = []
1134 self.makecycle(product([1,2,a,3], repeat=3), a)
1135
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001136 def test_repeat(self):
1137 a = []
1138 self.makecycle(repeat(a), a)
1139
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001140 def test_starmap(self):
1141 a = []
1142 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1143
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001144 def test_takewhile(self):
1145 a = []
1146 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1147
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001148def R(seqn):
1149 'Regular generator'
1150 for i in seqn:
1151 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001152
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001153class G:
1154 'Sequence using __getitem__'
1155 def __init__(self, seqn):
1156 self.seqn = seqn
1157 def __getitem__(self, i):
1158 return self.seqn[i]
1159
1160class I:
1161 'Sequence using iterator protocol'
1162 def __init__(self, seqn):
1163 self.seqn = seqn
1164 self.i = 0
1165 def __iter__(self):
1166 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001167 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001168 if self.i >= len(self.seqn): raise StopIteration
1169 v = self.seqn[self.i]
1170 self.i += 1
1171 return v
1172
1173class Ig:
1174 'Sequence using iterator protocol defined with a generator'
1175 def __init__(self, seqn):
1176 self.seqn = seqn
1177 self.i = 0
1178 def __iter__(self):
1179 for val in self.seqn:
1180 yield val
1181
1182class X:
1183 'Missing __getitem__ and __iter__'
1184 def __init__(self, seqn):
1185 self.seqn = seqn
1186 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001187 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001188 if self.i >= len(self.seqn): raise StopIteration
1189 v = self.seqn[self.i]
1190 self.i += 1
1191 return v
1192
1193class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001194 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001195 def __init__(self, seqn):
1196 self.seqn = seqn
1197 self.i = 0
1198 def __iter__(self):
1199 return self
1200
1201class E:
1202 'Test propagation of exceptions'
1203 def __init__(self, seqn):
1204 self.seqn = seqn
1205 self.i = 0
1206 def __iter__(self):
1207 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001208 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001209 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001210
1211class S:
1212 'Test immediate stop'
1213 def __init__(self, seqn):
1214 pass
1215 def __iter__(self):
1216 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001217 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001218 raise StopIteration
1219
1220def L(seqn):
1221 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001222 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001223
1224
1225class TestVariousIteratorArgs(unittest.TestCase):
1226
Raymond Hettinger482ba772010-12-01 22:48:00 +00001227 def test_accumulate(self):
1228 s = [1,2,3,4,5]
1229 r = [1,3,6,10,15]
1230 n = len(s)
1231 for g in (G, I, Ig, L, R):
1232 self.assertEqual(list(accumulate(g(s))), r)
1233 self.assertEqual(list(accumulate(S(s))), [])
1234 self.assertRaises(TypeError, accumulate, X(s))
1235 self.assertRaises(TypeError, accumulate, N(s))
1236 self.assertRaises(ZeroDivisionError, list, accumulate(E(s)))
1237
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001238 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001239 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001240 for g in (G, I, Ig, S, L, R):
1241 self.assertEqual(list(chain(g(s))), list(g(s)))
1242 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001243 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001244 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001245 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1246
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001247 def test_compress(self):
1248 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1249 n = len(s)
1250 for g in (G, I, Ig, S, L, R):
1251 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1252 self.assertRaises(TypeError, compress, X(s), repeat(1))
1253 self.assertRaises(TypeError, compress, N(s), repeat(1))
1254 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1255
Christian Heimesc3f30c42008-02-22 16:37:40 +00001256 def test_product(self):
1257 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1258 self.assertRaises(TypeError, product, X(s))
1259 self.assertRaises(TypeError, product, N(s))
1260 self.assertRaises(ZeroDivisionError, product, E(s))
1261
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001262 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001263 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001264 for g in (G, I, Ig, S, L, R):
1265 tgtlen = len(s) * 3
1266 expected = list(g(s))*3
1267 actual = list(islice(cycle(g(s)), tgtlen))
1268 self.assertEqual(actual, expected)
1269 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001270 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001271 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1272
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001273 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001274 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001275 for g in (G, I, Ig, S, L, R):
1276 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1277 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001278 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001279 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1280
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001281 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001282 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001283 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001284 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001285 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001286 self.assertRaises(TypeError, filter, isEven, X(s))
1287 self.assertRaises(TypeError, filter, isEven, N(s))
1288 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001289
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001290 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001291 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001292 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001293 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001294 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001295 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1296 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1297 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001298
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001299 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001300 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001301 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001302 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1303 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1304 self.assertRaises(TypeError, zip, X(s))
1305 self.assertRaises(TypeError, zip, N(s))
1306 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001307
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001308 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001309 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001310 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001311 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1312 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1313 self.assertRaises(TypeError, zip_longest, X(s))
1314 self.assertRaises(TypeError, zip_longest, N(s))
1315 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001316
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001317 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001318 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001319 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001320 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001321 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001322 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001323 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001324 self.assertRaises(TypeError, map, onearg, X(s))
1325 self.assertRaises(TypeError, map, onearg, N(s))
1326 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001327
1328 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001329 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001330 for g in (G, I, Ig, S, L, R):
1331 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1332 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001333 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001334 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1335
1336 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001337 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001338 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001339 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001340 self.assertEqual(list(starmap(operator.pow, g(ss))),
1341 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001342 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001343 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001344 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1345
1346 def test_takewhile(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 isEven(elem): break
1352 tgt.append(elem)
1353 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1354 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001355 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001356 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1357
1358 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001359 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001360 for g in (G, I, Ig, S, L, R):
1361 tgt = []
1362 for elem in g(s):
1363 if not tgt and isOdd(elem): continue
1364 tgt.append(elem)
1365 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1366 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001367 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001368 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1369
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001370 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001371 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001372 for g in (G, I, Ig, S, L, R):
1373 it1, it2 = tee(g(s))
1374 self.assertEqual(list(it1), list(g(s)))
1375 self.assertEqual(list(it2), list(g(s)))
1376 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001377 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001378 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1379
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001380class LengthTransparency(unittest.TestCase):
1381
1382 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001383 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001384 self.assertEqual(len(repeat(None, 50)), 50)
1385 self.assertRaises(TypeError, len, repeat(None))
1386
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001387class RegressionTests(unittest.TestCase):
1388
1389 def test_sf_793826(self):
1390 # Fix Armin Rigo's successful efforts to wreak havoc
1391
1392 def mutatingtuple(tuple1, f, tuple2):
1393 # this builds a tuple t which is a copy of tuple1,
1394 # then calls f(t), then mutates t to be equal to tuple2
1395 # (needs len(tuple1) == len(tuple2)).
1396 def g(value, first=[1]):
1397 if first:
1398 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00001399 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001400 return value
1401 items = list(tuple2)
1402 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001403 gen = map(g, items)
1404 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00001405 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001406
1407 def f(t):
1408 global T
1409 T = t
1410 first[:] = list(T)
1411
1412 first = []
1413 mutatingtuple((1,2,3), f, (4,5,6))
1414 second = list(T)
1415 self.assertEqual(first, second)
1416
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001417
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001418 def test_sf_950057(self):
1419 # Make sure that chain() and cycle() catch exceptions immediately
1420 # rather than when shifting between input sources
1421
1422 def gen1():
1423 hist.append(0)
1424 yield 1
1425 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001426 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001427 hist.append(2)
1428
1429 def gen2(x):
1430 hist.append(3)
1431 yield 2
1432 hist.append(4)
1433 if x:
1434 raise StopIteration
1435
1436 hist = []
1437 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1438 self.assertEqual(hist, [0,1])
1439
1440 hist = []
1441 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1442 self.assertEqual(hist, [0,1])
1443
1444 hist = []
1445 self.assertRaises(AssertionError, list, cycle(gen1()))
1446 self.assertEqual(hist, [0,1])
1447
Thomas Woutersb2137042007-02-01 18:02:27 +00001448class SubclassWithKwargsTest(unittest.TestCase):
1449 def test_keywords_in_subclass(self):
1450 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001451 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001452 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00001453 class Subclass(cls):
1454 def __init__(self, newarg=None, *args):
1455 cls.__init__(self, *args)
1456 try:
1457 Subclass(newarg=1)
1458 except TypeError as err:
1459 # we expect type errors because of wrong argument count
Ezio Melottib58e0bd2010-01-23 15:40:09 +00001460 self.assertNotIn("does not take keyword arguments", err.args[0])
Thomas Woutersb2137042007-02-01 18:02:27 +00001461
1462
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001463libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001464
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001465
1466>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001467>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00001468... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001469...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001470Check 1200 is for $120.15
1471Check 1201 is for $764.05
1472Check 1202 is for $823.14
1473
1474>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001475>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00001476... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00001477...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000014781
14798
148027
1481
1482>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001483>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00001484... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00001485...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001486Alex
1487Laura
1488Martin
1489Walter
1490Samuele
1491
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001492>>> from operator import itemgetter
1493>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00001494>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001495>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001496... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001497...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000014981 ['a', 'c', 'e']
14992 ['b', 'd', 'f']
15003 ['g']
1501
Raymond Hettinger734fb572004-01-20 20:04:40 +00001502# Find runs of consecutive numbers using groupby. The key to the solution
1503# is differencing with a range so that consecutive numbers all appear in
1504# same group.
1505>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00001506>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001507... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001508...
Raymond Hettinger734fb572004-01-20 20:04:40 +00001509[1]
1510[4, 5, 6]
1511[10]
1512[15, 16, 17, 18]
1513[22]
1514[25, 26, 27, 28]
1515
Georg Brandl3dbca812008-07-23 16:10:53 +00001516>>> def take(n, iterable):
1517... "Return first n items of the iterable as a list"
1518... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001519
Georg Brandl3dbca812008-07-23 16:10:53 +00001520>>> def enumerate(iterable, start=0):
1521... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001522
Georg Brandl3dbca812008-07-23 16:10:53 +00001523>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001524... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00001525... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001526
Raymond Hettingercdf8ba32009-02-19 04:45:07 +00001527>>> def nth(iterable, n, default=None):
1528... "Returns the nth item or a default value"
1529... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001530
Georg Brandl3dbca812008-07-23 16:10:53 +00001531>>> def quantify(iterable, pred=bool):
1532... "Count how many times the predicate is true"
1533... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001534
Georg Brandl3dbca812008-07-23 16:10:53 +00001535>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001536... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00001537... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001538
Georg Brandl3dbca812008-07-23 16:10:53 +00001539>>> def ncycles(iterable, n):
Ezio Melotti13925002011-03-16 11:05:33 +02001540... "Returns the sequence elements n times"
Georg Brandl3dbca812008-07-23 16:10:53 +00001541... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001542
1543>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001544... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001545
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001546>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001547... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001548
1549>>> def repeatfunc(func, times=None, *args):
1550... "Repeat calls to func with specified arguments."
1551... " Example: repeatfunc(random.random)"
1552... if times is None:
1553... return starmap(func, repeat(args))
1554... else:
1555... return starmap(func, repeat(args, times))
1556
Raymond Hettingerd591f662003-10-26 15:34:50 +00001557>>> def pairwise(iterable):
1558... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1559... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001560... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00001561... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001562... except StopIteration:
1563... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001564... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001565
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001566>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001567... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001568... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00001569... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001570
1571>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001572... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001573... # Recipe credited to George Sakkis
1574... pending = len(iterables)
1575... nexts = cycle(iter(it).__next__ for it in iterables)
1576... while pending:
1577... try:
1578... for next in nexts:
1579... yield next()
1580... except StopIteration:
1581... pending -= 1
1582... nexts = cycle(islice(nexts, pending))
1583
1584>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00001585... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1586... s = list(iterable)
1587... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001588
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001589>>> def unique_everseen(iterable, key=None):
1590... "List unique elements, preserving order. Remember all elements ever seen."
1591... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1592... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1593... seen = set()
1594... seen_add = seen.add
1595... if key is None:
1596... for element in iterable:
1597... if element not in seen:
1598... seen_add(element)
1599... yield element
1600... else:
1601... for element in iterable:
1602... k = key(element)
1603... if k not in seen:
1604... seen_add(k)
1605... yield element
1606
1607>>> def unique_justseen(iterable, key=None):
1608... "List unique elements, preserving order. Remember only the element just seen."
1609... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1610... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1611... return map(next, map(itemgetter(1), groupby(iterable, key)))
1612
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001613This is not part of the examples but it tests to make sure the definitions
1614perform as purported.
1615
Raymond Hettingera098b332003-09-08 23:58:40 +00001616>>> take(10, count())
1617[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1618
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001619>>> list(enumerate('abc'))
1620[(0, 'a'), (1, 'b'), (2, 'c')]
1621
1622>>> list(islice(tabulate(lambda x: 2*x), 4))
1623[0, 2, 4, 6]
1624
1625>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00001626'd'
1627
1628>>> nth('abcde', 9) is None
1629True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001630
Guido van Rossum805365e2007-05-07 22:24:25 +00001631>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000163250
1633
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001634>>> a = [[1, 2, 3], [4, 5, 6]]
1635>>> flatten(a)
1636[1, 2, 3, 4, 5, 6]
1637
1638>>> list(repeatfunc(pow, 5, 2, 3))
1639[8, 8, 8, 8, 8]
1640
1641>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001642>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001643[0, 0, 0, 0, 0]
1644
Raymond Hettingerd591f662003-10-26 15:34:50 +00001645>>> list(pairwise('abcd'))
1646[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001647
Raymond Hettingerd591f662003-10-26 15:34:50 +00001648>>> list(pairwise([]))
1649[]
1650
1651>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001652[]
1653
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001654>>> list(islice(padnone('abc'), 0, 6))
1655['a', 'b', 'c', None, None, None]
1656
1657>>> list(ncycles('abc', 3))
1658['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1659
1660>>> dotproduct([1,2,3], [4,5,6])
166132
1662
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001663>>> list(grouper(3, 'abcdefg', 'x'))
1664[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1665
1666>>> list(roundrobin('abc', 'd', 'ef'))
1667['a', 'd', 'e', 'b', 'f', 'c']
1668
Raymond Hettingerace67332009-01-26 02:23:50 +00001669>>> list(powerset([1,2,3]))
1670[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001671
Raymond Hettinger191e8502009-01-27 13:29:43 +00001672>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1673True
1674
1675>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1676True
1677
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001678>>> list(unique_everseen('AAAABBBCCDAABBB'))
1679['A', 'B', 'C', 'D']
1680
1681>>> list(unique_everseen('ABBCcAD', str.lower))
1682['A', 'B', 'C', 'D']
1683
1684>>> list(unique_justseen('AAAABBBCCDAABBB'))
1685['A', 'B', 'C', 'D', 'A', 'B']
1686
1687>>> list(unique_justseen('ABBCcAD', str.lower))
1688['A', 'B', 'C', 'A', 'D']
1689
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001690"""
1691
1692__test__ = {'libreftest' : libreftest}
1693
1694def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001695 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00001696 RegressionTests, LengthTransparency,
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001697 SubclassWithKwargsTest, TestExamples)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001698 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001699
1700 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001701 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001702 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001703 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00001704 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001705 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001706 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001707 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001708 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001709
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001710 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001711 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001712
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001713if __name__ == "__main__":
1714 test_main(verbose=True)