blob: acbb00a9ff2f017105c797f5965ece559eba67ea [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
Raymond Hettinger5d446132011-03-27 18:52:10 -070072 self.assertRaises(TypeError, accumulate, range(10), 5, 6) # 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 Hettinger5d446132011-03-27 18:52:10 -070077 s = [2, 8, 9, 5, 7, 0, 3, 4, 1, 6]
78 self.assertEqual(list(accumulate(s, min)),
79 [2, 2, 2, 2, 2, 0, 0, 0, 0, 0])
80 self.assertEqual(list(accumulate(s, max)),
81 [2, 8, 9, 9, 9, 9, 9, 9, 9, 9])
82 self.assertEqual(list(accumulate(s, operator.mul)),
83 [2, 16, 144, 720, 5040, 0, 0, 0, 0, 0])
84 with self.assertRaises(TypeError):
85 list(accumulate(s, chr)) # unary-operation
86
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000087 def test_chain(self):
Christian Heimesdd15f6c2008-03-16 00:07:10 +000088
89 def chain2(*iterables):
90 'Pure python version in the docs'
91 for it in iterables:
92 for element in it:
93 yield element
94
95 for c in (chain, chain2):
96 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
97 self.assertEqual(list(c('abc')), list('abc'))
98 self.assertEqual(list(c('')), [])
99 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
100 self.assertRaises(TypeError, list,c(2, 3))
Christian Heimesf16baeb2008-02-29 14:57:44 +0000101
102 def test_chain_from_iterable(self):
103 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
104 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
105 self.assertEqual(list(chain.from_iterable([''])), [])
106 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
107 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000108
Christian Heimes380f7f22008-02-28 11:19:05 +0000109 def test_combinations(self):
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000110 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Christian Heimes380f7f22008-02-28 11:19:05 +0000111 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Christian Heimes78644762008-03-04 23:39:23 +0000112 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Christian Heimes380f7f22008-02-28 11:19:05 +0000113 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000114 self.assertEqual(list(combinations('abc', 32)), []) # r > n
Christian Heimes380f7f22008-02-28 11:19:05 +0000115 self.assertEqual(list(combinations(range(4), 3)),
116 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
Christian Heimes78644762008-03-04 23:39:23 +0000117
118 def combinations1(iterable, r):
119 'Pure python version shown in the docs'
120 pool = tuple(iterable)
121 n = len(pool)
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000122 if r > n:
123 return
Christian Heimes78644762008-03-04 23:39:23 +0000124 indices = list(range(r))
125 yield tuple(pool[i] for i in indices)
126 while 1:
127 for i in reversed(range(r)):
128 if indices[i] != i + n - r:
129 break
130 else:
131 return
132 indices[i] += 1
133 for j in range(i+1, r):
134 indices[j] = indices[j-1] + 1
135 yield tuple(pool[i] for i in indices)
136
137 def combinations2(iterable, r):
138 'Pure python version shown in the docs'
139 pool = tuple(iterable)
140 n = len(pool)
141 for indices in permutations(range(n), r):
142 if sorted(indices) == list(indices):
143 yield tuple(pool[i] for i in indices)
144
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000145 def combinations3(iterable, r):
146 'Pure python version from cwr()'
147 pool = tuple(iterable)
148 n = len(pool)
149 for indices in combinations_with_replacement(range(n), r):
150 if len(set(indices)) == r:
151 yield tuple(pool[i] for i in indices)
152
Christian Heimes78644762008-03-04 23:39:23 +0000153 for n in range(7):
Christian Heimes380f7f22008-02-28 11:19:05 +0000154 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000155 for r in range(n+2):
Christian Heimes380f7f22008-02-28 11:19:05 +0000156 result = list(combinations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000157 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 +0000158 self.assertEqual(len(result), len(set(result))) # no repeats
159 self.assertEqual(result, sorted(result)) # lexicographic order
160 for c in result:
161 self.assertEqual(len(c), r) # r-length combinations
162 self.assertEqual(len(set(c)), r) # no duplicate elements
163 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000164 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000165 self.assertEqual(list(c),
166 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000167 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000168 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000169 self.assertEqual(result, list(combinations3(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000170
171 # Test implementation detail: tuple re-use
172 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
173 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
174
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000175 def test_combinations_with_replacement(self):
176 cwr = combinations_with_replacement
177 self.assertRaises(TypeError, cwr, 'abc') # missing r argument
178 self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
179 self.assertRaises(TypeError, cwr, None) # pool is not iterable
180 self.assertRaises(ValueError, cwr, 'abc', -2) # r is negative
181 self.assertEqual(list(cwr('ABC', 2)),
182 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
183
184 def cwr1(iterable, r):
185 'Pure python version shown in the docs'
186 # number items returned: (n+r-1)! / r! / (n-1)! when n>0
187 pool = tuple(iterable)
188 n = len(pool)
189 if not n and r:
190 return
191 indices = [0] * r
192 yield tuple(pool[i] for i in indices)
193 while 1:
194 for i in reversed(range(r)):
195 if indices[i] != n - 1:
196 break
197 else:
198 return
199 indices[i:] = [indices[i] + 1] * (r - i)
200 yield tuple(pool[i] for i in indices)
201
202 def cwr2(iterable, r):
203 'Pure python version shown in the docs'
204 pool = tuple(iterable)
205 n = len(pool)
206 for indices in product(range(n), repeat=r):
207 if sorted(indices) == list(indices):
208 yield tuple(pool[i] for i in indices)
209
210 def numcombs(n, r):
211 if not n:
212 return 0 if r else 1
213 return fact(n+r-1) / fact(r)/ fact(n-1)
214
215 for n in range(7):
216 values = [5*x-12 for x in range(n)]
217 for r in range(n+2):
218 result = list(cwr(values, r))
219
220 self.assertEqual(len(result), numcombs(n, r)) # right number of combs
221 self.assertEqual(len(result), len(set(result))) # no repeats
222 self.assertEqual(result, sorted(result)) # lexicographic order
223
224 regular_combs = list(combinations(values, r)) # compare to combs without replacement
225 if n == 0 or r <= 1:
Ezio Melottib3aedd42010-11-20 19:04:17 +0000226 self.assertEqual(result, regular_combs) # cases that should be identical
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000227 else:
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000228 self.assertTrue(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000229
230 for c in result:
231 self.assertEqual(len(c), r) # r-length combinations
232 noruns = [k for k,v in groupby(c)] # combo without consecutive repeats
233 self.assertEqual(len(noruns), len(set(noruns))) # no repeats other than consecutive
234 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000235 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000236 self.assertEqual(noruns,
237 [e for e in values if e in c]) # comb is a subsequence of the input iterable
238 self.assertEqual(result, list(cwr1(values, r))) # matches first pure python version
239 self.assertEqual(result, list(cwr2(values, r))) # matches second pure python version
240
241 # Test implementation detail: tuple re-use
242 self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
243 self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
244
Christian Heimes78644762008-03-04 23:39:23 +0000245 def test_permutations(self):
246 self.assertRaises(TypeError, permutations) # too few arguments
247 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000248 self.assertRaises(TypeError, permutations, None) # pool is not iterable
249 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000250 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000251 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Christian Heimes78644762008-03-04 23:39:23 +0000252 self.assertEqual(list(permutations(range(3), 2)),
253 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
254
255 def permutations1(iterable, r=None):
256 'Pure python version shown in the docs'
257 pool = tuple(iterable)
258 n = len(pool)
259 r = n if r is None else r
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000260 if r > n:
261 return
Christian Heimes78644762008-03-04 23:39:23 +0000262 indices = list(range(n))
263 cycles = list(range(n-r+1, n+1))[::-1]
264 yield tuple(pool[i] for i in indices[:r])
265 while n:
266 for i in reversed(range(r)):
267 cycles[i] -= 1
268 if cycles[i] == 0:
269 indices[i:] = indices[i+1:] + indices[i:i+1]
270 cycles[i] = n - i
271 else:
272 j = cycles[i]
273 indices[i], indices[-j] = indices[-j], indices[i]
274 yield tuple(pool[i] for i in indices[:r])
275 break
276 else:
277 return
278
279 def permutations2(iterable, r=None):
280 'Pure python version shown in the docs'
281 pool = tuple(iterable)
282 n = len(pool)
283 r = n if r is None else r
284 for indices in product(range(n), repeat=r):
285 if len(set(indices)) == r:
286 yield tuple(pool[i] for i in indices)
287
288 for n in range(7):
289 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000290 for r in range(n+2):
Christian Heimes78644762008-03-04 23:39:23 +0000291 result = list(permutations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000292 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 +0000293 self.assertEqual(len(result), len(set(result))) # no repeats
294 self.assertEqual(result, sorted(result)) # lexicographic order
295 for p in result:
296 self.assertEqual(len(p), r) # r-length permutations
297 self.assertEqual(len(set(p)), r) # no duplicate elements
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000298 self.assertTrue(all(e in values for e in p)) # elements taken from input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000299 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000300 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000301 if r == n:
302 self.assertEqual(result, list(permutations(values, None))) # test r as None
303 self.assertEqual(result, list(permutations(values))) # test default r
304
305 # Test implementation detail: tuple re-use
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000306 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Christian Heimes78644762008-03-04 23:39:23 +0000307 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Christian Heimes380f7f22008-02-28 11:19:05 +0000308
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000309 def test_combinatorics(self):
310 # Test relationships between product(), permutations(),
311 # combinations() and combinations_with_replacement().
312
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000313 for n in range(6):
314 s = 'ABCDEFG'[:n]
315 for r in range(8):
316 prod = list(product(s, repeat=r))
317 cwr = list(combinations_with_replacement(s, r))
318 perm = list(permutations(s, r))
319 comb = list(combinations(s, r))
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000320
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000321 # Check size
Ezio Melottib3aedd42010-11-20 19:04:17 +0000322 self.assertEqual(len(prod), n**r)
323 self.assertEqual(len(cwr), (fact(n+r-1) / fact(r)/ fact(n-1)) if n else (not r))
324 self.assertEqual(len(perm), 0 if r>n else fact(n) / fact(n-r))
325 self.assertEqual(len(comb), 0 if r>n else fact(n) / fact(r) / fact(n-r))
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000326
327 # Check lexicographic order without repeated tuples
Ezio Melottib3aedd42010-11-20 19:04:17 +0000328 self.assertEqual(prod, sorted(set(prod)))
329 self.assertEqual(cwr, sorted(set(cwr)))
330 self.assertEqual(perm, sorted(set(perm)))
331 self.assertEqual(comb, sorted(set(comb)))
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000332
333 # Check interrelationships
Ezio Melottib3aedd42010-11-20 19:04:17 +0000334 self.assertEqual(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
335 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 +0000336 self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
337 self.assertEqual(comb, [t for t in cwr if len(set(t))==r]) # comb: cwrs without dups
338 self.assertEqual(comb, list(filter(set(cwr).__contains__, perm))) # comb: perm that is a cwr
339 self.assertEqual(comb, list(filter(set(perm).__contains__, cwr))) # comb: cwr that is a perm
340 self.assertEqual(comb, sorted(set(cwr) & set(perm))) # comb: both a cwr and a perm
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000341
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000342 def test_compress(self):
Raymond Hettinger15a49502009-02-19 02:17:09 +0000343 self.assertEqual(list(compress(data='ABCDEF', selectors=[1,0,1,0,1,1])), list('ACEF'))
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000344 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
345 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
346 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
347 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
348 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
349 n = 10000
350 data = chain.from_iterable(repeat(range(6), n))
351 selectors = chain.from_iterable(repeat((0, 1)))
352 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
353 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
354 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
355 self.assertRaises(TypeError, compress, range(6)) # too few args
356 self.assertRaises(TypeError, compress, range(6), None) # too many args
357
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000358 def test_count(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000359 self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
360 self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
361 self.assertEqual(take(2, lzip('abc',count(3))), [('a', 3), ('b', 4)])
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000362 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
363 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000364 self.assertRaises(TypeError, count, 2, 3, 4)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000365 self.assertRaises(TypeError, count, 'a')
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000366 self.assertEqual(list(islice(count(maxsize-5), 10)),
367 list(range(maxsize-5, maxsize+5)))
368 self.assertEqual(list(islice(count(-maxsize-5), 10)),
369 list(range(-maxsize-5, -maxsize+5)))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000370 c = count(3)
371 self.assertEqual(repr(c), 'count(3)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000372 next(c)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000373 self.assertEqual(repr(c), 'count(4)')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000374 c = count(-9)
375 self.assertEqual(repr(c), 'count(-9)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000376 next(c)
Raymond Hettinger30729212009-02-12 06:28:27 +0000377 self.assertEqual(repr(count(10.25)), 'count(10.25)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000378 self.assertEqual(next(c), -8)
Christian Heimesa37d4c62007-12-04 23:02:19 +0000379 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 +0000380 # Test repr (ignoring the L in longs)
381 r1 = repr(count(i)).replace('L', '')
382 r2 = 'count(%r)'.__mod__(i).replace('L', '')
383 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000384
Raymond Hettingera3e1ad22009-11-30 22:02:31 +0000385 # check copy, deepcopy, pickle
386 for value in -3, 3, maxsize-5, maxsize+5:
387 c = count(value)
388 self.assertEqual(next(copy.copy(c)), value)
389 self.assertEqual(next(copy.deepcopy(c)), value)
390 self.assertEqual(next(pickle.loads(pickle.dumps(c))), value)
391
Raymond Hettinger30729212009-02-12 06:28:27 +0000392 def test_count_with_stride(self):
393 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingereb13fdd2009-02-21 22:30:12 +0000394 self.assertEqual(lzip('abc',count(start=2,step=3)),
395 [('a', 2), ('b', 5), ('c', 8)])
396 self.assertEqual(lzip('abc',count(step=-1)),
397 [('a', 0), ('b', -1), ('c', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000398 self.assertEqual(lzip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
399 self.assertEqual(lzip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000400 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000401 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
402 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
403 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettinger8a882a72009-02-12 12:08:18 +0000404 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
405 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerde09acf2009-02-12 12:53:53 +0000406 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
407 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000408 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
409 c = count(3, 5)
410 self.assertEqual(repr(c), 'count(3, 5)')
411 next(c)
412 self.assertEqual(repr(c), 'count(8, 5)')
413 c = count(-9, 0)
414 self.assertEqual(repr(c), 'count(-9, 0)')
415 next(c)
416 self.assertEqual(repr(c), 'count(-9, 0)')
417 c = count(-9, -3)
418 self.assertEqual(repr(c), 'count(-9, -3)')
419 next(c)
420 self.assertEqual(repr(c), 'count(-12, -3)')
421 self.assertEqual(repr(c), 'count(-12, -3)')
422 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
423 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
424 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
425 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
426 for j in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 1, 10, sys.maxsize-5, sys.maxsize+5):
427 # Test repr (ignoring the L in longs)
428 r1 = repr(count(i, j)).replace('L', '')
429 if j == 1:
430 r2 = ('count(%r)' % i).replace('L', '')
431 else:
432 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
433 self.assertEqual(r1, r2)
434
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000435 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000436 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000437 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000438 self.assertRaises(TypeError, cycle)
439 self.assertRaises(TypeError, cycle, 5)
440 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000441
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000442 def test_groupby(self):
443 # Check whether it accepts arguments correctly
444 self.assertEqual([], list(groupby([])))
445 self.assertEqual([], list(groupby([], key=id)))
446 self.assertRaises(TypeError, list, groupby('abc', []))
447 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000448 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000449
450 # Check normal input
451 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
452 (2,15,22), (3,16,23), (3,17,23)]
453 dup = []
454 for k, g in groupby(s, lambda r:r[0]):
455 for elem in g:
456 self.assertEqual(k, elem[0])
457 dup.append(elem)
458 self.assertEqual(s, dup)
459
460 # Check nested case
461 dup = []
462 for k, g in groupby(s, lambda r:r[0]):
463 for ik, ig in groupby(g, lambda r:r[2]):
464 for elem in ig:
465 self.assertEqual(k, elem[0])
466 self.assertEqual(ik, elem[2])
467 dup.append(elem)
468 self.assertEqual(s, dup)
469
470 # Check case where inner iterator is not used
471 keys = [k for k, g in groupby(s, lambda r:r[0])]
472 expectedkeys = set([r[0] for r in s])
473 self.assertEqual(set(keys), expectedkeys)
474 self.assertEqual(len(keys), len(expectedkeys))
475
476 # Exercise pipes and filters style
477 s = 'abracadabra'
478 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000479 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000480 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
481 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000482 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000483 self.assertEqual(r, ['a', 'b', 'r'])
484 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000485 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000486 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
487 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000488 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000489 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
490
Georg Brandla18af4e2007-04-21 15:47:16 +0000491 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000492 class ExpectedError(Exception):
493 pass
494 def delayed_raise(n=0):
495 for i in range(n):
496 yield 'yo'
497 raise ExpectedError
498 def gulp(iterable, keyp=None, func=list):
499 return [func(g) for k, g in groupby(iterable, keyp)]
500
Georg Brandla18af4e2007-04-21 15:47:16 +0000501 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000502 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000503 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000504 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
505
506 # __cmp__ failure
507 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000508 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000509 raise ExpectedError
510 s = [DummyCmp(), DummyCmp(), None]
511
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000512 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000513 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000514 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000515 self.assertRaises(ExpectedError, gulp, s)
516
517 # keyfunc failure
518 def keyfunc(obj):
519 if keyfunc.skip > 0:
520 keyfunc.skip -= 1
521 return obj
522 else:
523 raise ExpectedError
524
525 # keyfunc failure on outer object
526 keyfunc.skip = 0
527 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
528 keyfunc.skip = 1
529 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
530
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000531 def test_filter(self):
532 self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
533 self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])
534 self.assertEqual(list(filter(bool, [0,1,0,2,0])), [1,2])
535 self.assertEqual(take(4, filter(isEven, count())), [0,2,4,6])
536 self.assertRaises(TypeError, filter)
537 self.assertRaises(TypeError, filter, lambda x:x)
538 self.assertRaises(TypeError, filter, lambda x:x, range(6), 7)
539 self.assertRaises(TypeError, filter, isEven, 3)
540 self.assertRaises(TypeError, next, filter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000541
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000542 def test_filterfalse(self):
543 self.assertEqual(list(filterfalse(isEven, range(6))), [1,3,5])
544 self.assertEqual(list(filterfalse(None, [0,1,0,2,0])), [0,0,0])
545 self.assertEqual(list(filterfalse(bool, [0,1,0,2,0])), [0,0,0])
546 self.assertEqual(take(4, filterfalse(isEven, count())), [1,3,5,7])
547 self.assertRaises(TypeError, filterfalse)
548 self.assertRaises(TypeError, filterfalse, lambda x:x)
549 self.assertRaises(TypeError, filterfalse, lambda x:x, range(6), 7)
550 self.assertRaises(TypeError, filterfalse, isEven, 3)
551 self.assertRaises(TypeError, next, filterfalse(range(6), range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000552
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000553 def test_zip(self):
554 # XXX This is rather silly now that builtin zip() calls zip()...
555 ans = [(x,y) for x, y in zip('abc',count())]
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000556 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000557 self.assertEqual(list(zip('abc', range(6))), lzip('abc', range(6)))
558 self.assertEqual(list(zip('abcdef', range(3))), lzip('abcdef', range(3)))
559 self.assertEqual(take(3,zip('abcdef', count())), lzip('abcdef', range(3)))
560 self.assertEqual(list(zip('abcdef')), lzip('abcdef'))
561 self.assertEqual(list(zip()), lzip())
562 self.assertRaises(TypeError, zip, 3)
563 self.assertRaises(TypeError, zip, range(3), 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000564 # Check tuple re-use (implementation detail)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000565 self.assertEqual([tuple(list(pair)) for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000566 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000567 self.assertEqual([pair for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000568 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000569 ids = list(map(id, zip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000570 self.assertEqual(min(ids), max(ids))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000571 ids = list(map(id, list(zip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000572 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000573
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000574 def test_ziplongest(self):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000575 for args in [
576 ['abc', range(6)],
577 [range(6), 'abc'],
578 [range(1000), range(2000,2100), range(3000,3050)],
579 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
580 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
581 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000582 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
583 for i in range(max(map(len, args)))]
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000584 self.assertEqual(list(zip_longest(*args)), target)
585 self.assertEqual(list(zip_longest(*args, **{})), target)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000586 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 +0000587 self.assertEqual(list(zip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000588
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000589 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 +0000590
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000591 self.assertEqual(list(zip_longest()), list(zip()))
592 self.assertEqual(list(zip_longest([])), list(zip([])))
593 self.assertEqual(list(zip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000594
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000595 self.assertEqual(list(zip_longest('abc', 'defg', **{})),
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000596 list(zip(list('abc')+[None], 'defg'))) # empty keyword dict
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000597 self.assertRaises(TypeError, zip_longest, 3)
598 self.assertRaises(TypeError, zip_longest, range(3), 3)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000599
600 for stmt in [
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000601 "zip_longest('abc', fv=1)",
602 "zip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000603 ]:
604 try:
605 eval(stmt, globals(), locals())
606 except TypeError:
607 pass
608 else:
609 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000610
Thomas Wouterscf297e42007-02-23 15:07:44 +0000611 # Check tuple re-use (implementation detail)
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000612 self.assertEqual([tuple(list(pair)) for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000613 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000614 self.assertEqual([pair for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000615 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000616 ids = list(map(id, zip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000617 self.assertEqual(min(ids), max(ids))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000618 ids = list(map(id, list(zip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000619 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
620
Raymond Hettingerfc438512009-11-01 20:55:33 +0000621 def test_bug_7244(self):
622
623 class Repeater:
624 # this class is similar to itertools.repeat
625 def __init__(self, o, t, e):
626 self.o = o
627 self.t = int(t)
628 self.e = e
629 def __iter__(self): # its iterator is itself
630 return self
631 def __next__(self):
632 if self.t > 0:
633 self.t -= 1
634 return self.o
635 else:
636 raise self.e
637
638 # Formerly this code in would fail in debug mode
639 # with Undetected Error and Stop Iteration
640 r1 = Repeater(1, 3, StopIteration)
641 r2 = Repeater(2, 4, StopIteration)
642 def run(r1, r2):
643 result = []
644 for i, j in zip_longest(r1, r2, fillvalue=0):
645 with support.captured_output('stdout'):
646 print((i, j))
647 result.append((i, j))
648 return result
649 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
650
651 # Formerly, the RuntimeError would be lost
652 # and StopIteration would stop as expected
653 r1 = Repeater(1, 3, RuntimeError)
654 r2 = Repeater(2, 4, StopIteration)
655 it = zip_longest(r1, r2, fillvalue=0)
656 self.assertEqual(next(it), (1, 2))
657 self.assertEqual(next(it), (1, 2))
658 self.assertEqual(next(it), (1, 2))
659 self.assertRaises(RuntimeError, next, it)
660
Christian Heimesc3f30c42008-02-22 16:37:40 +0000661 def test_product(self):
662 for args, result in [
Christian Heimes78644762008-03-04 23:39:23 +0000663 ([], [()]), # zero iterables
Christian Heimesc3f30c42008-02-22 16:37:40 +0000664 (['ab'], [('a',), ('b',)]), # one iterable
665 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
666 ([range(0), range(2), range(3)], []), # first iterable with zero length
667 ([range(2), range(0), range(3)], []), # middle iterable with zero length
668 ([range(2), range(3), range(0)], []), # last iterable with zero length
669 ]:
670 self.assertEqual(list(product(*args)), result)
Christian Heimesf16baeb2008-02-29 14:57:44 +0000671 for r in range(4):
672 self.assertEqual(list(product(*(args*r))),
673 list(product(*args, **dict(repeat=r))))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000674 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
675 self.assertRaises(TypeError, product, range(6), None)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000676
677 def product1(*args, **kwds):
678 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
679 n = len(pools)
680 if n == 0:
681 yield ()
682 return
683 if any(len(pool) == 0 for pool in pools):
684 return
685 indices = [0] * n
686 yield tuple(pool[i] for pool, i in zip(pools, indices))
687 while 1:
688 for i in reversed(range(n)): # right to left
689 if indices[i] == len(pools[i]) - 1:
690 continue
691 indices[i] += 1
692 for j in range(i+1, n):
693 indices[j] = 0
694 yield tuple(pool[i] for pool, i in zip(pools, indices))
695 break
696 else:
697 return
698
699 def product2(*args, **kwds):
700 'Pure python version used in docs'
701 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
702 result = [[]]
703 for pool in pools:
704 result = [x+[y] for x in result for y in pool]
705 for prod in result:
706 yield tuple(prod)
707
Christian Heimesc3f30c42008-02-22 16:37:40 +0000708 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
709 set('abcdefg'), range(11), tuple(range(13))]
710 for i in range(100):
711 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Christian Heimes78644762008-03-04 23:39:23 +0000712 expected_len = prod(map(len, args))
713 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000714 self.assertEqual(list(product(*args)), list(product1(*args)))
715 self.assertEqual(list(product(*args)), list(product2(*args)))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000716 args = map(iter, args)
Christian Heimes78644762008-03-04 23:39:23 +0000717 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000718
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000719 # Test implementation detail: tuple re-use
720 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
721 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000722
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000723 def test_repeat(self):
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +0000724 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
Guido van Rossum805365e2007-05-07 22:24:25 +0000725 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000726 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000727 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000728 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000729 self.assertEqual(list(repeat('a', 0)), [])
730 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000731 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000732 self.assertRaises(TypeError, repeat, None, 3, 4)
733 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000734 r = repeat(1+0j)
735 self.assertEqual(repr(r), 'repeat((1+0j))')
736 r = repeat(1+0j, 5)
737 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
738 list(r)
739 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000740
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000741 def test_map(self):
742 self.assertEqual(list(map(operator.pow, range(3), range(1,7))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000743 [0**1, 1**2, 2**3])
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000744 def tupleize(*args):
745 return args
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000746 self.assertEqual(list(map(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000747 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000748 self.assertEqual(list(map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000749 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000750 self.assertEqual(take(2,map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000751 [('a',0),('b',1)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000752 self.assertEqual(list(map(operator.pow, [])), [])
753 self.assertRaises(TypeError, map)
754 self.assertRaises(TypeError, list, map(None, range(3), range(3)))
755 self.assertRaises(TypeError, map, operator.neg)
756 self.assertRaises(TypeError, next, map(10, range(5)))
757 self.assertRaises(ValueError, next, map(errfunc, [4], [5]))
758 self.assertRaises(TypeError, next, map(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000759
760 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000761 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
762 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000763 self.assertEqual(take(3, starmap(operator.pow, zip(count(), count(1)))),
Raymond Hettinger02420702003-06-29 20:36:23 +0000764 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000765 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +0000766 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
767 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000768 self.assertRaises(TypeError, starmap)
769 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000770 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
771 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
772 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000773
774 def test_islice(self):
775 for args in [ # islice(args) should agree with range(args)
776 (10, 20, 3),
777 (10, 3, 20),
778 (10, 20),
779 (10, 3),
780 (20,)
781 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000782 self.assertEqual(list(islice(range(100), *args)),
783 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000784
785 for args, tgtargs in [ # Stop when seqn is exhausted
786 ((10, 110, 3), ((10, 100, 3))),
787 ((10, 110), ((10, 100))),
788 ((110,), (100,))
789 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000790 self.assertEqual(list(islice(range(100), *args)),
791 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000792
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000793 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +0000794 self.assertEqual(list(islice(range(10), None)), list(range(10)))
795 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
796 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
797 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
798 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000799
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000800 # Test number of items consumed SF #1171417
801 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +0000802 self.assertEqual(list(islice(it, 3)), list(range(3)))
803 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000804
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000805 # Test invalid arguments
Guido van Rossum805365e2007-05-07 22:24:25 +0000806 self.assertRaises(TypeError, islice, range(10))
807 self.assertRaises(TypeError, islice, range(10), 1, 2, 3, 4)
808 self.assertRaises(ValueError, islice, range(10), -5, 10, 1)
809 self.assertRaises(ValueError, islice, range(10), 1, -5, -1)
810 self.assertRaises(ValueError, islice, range(10), 1, 10, -1)
811 self.assertRaises(ValueError, islice, range(10), 1, 10, 0)
812 self.assertRaises(ValueError, islice, range(10), 'a')
813 self.assertRaises(ValueError, islice, range(10), 'a', 1)
814 self.assertRaises(ValueError, islice, range(10), 1, 'a')
815 self.assertRaises(ValueError, islice, range(10), 'a', 1, 1)
816 self.assertRaises(ValueError, islice, range(10), 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +0000817 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000818
Raymond Hettinger69b34bf2010-11-30 02:49:29 +0000819 # Issue #10323: Less islice in a predictable state
820 c = count()
821 self.assertEqual(list(islice(c, 1, 3, 50)), [1])
822 self.assertEqual(next(c), 3)
823
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000824 def test_takewhile(self):
825 data = [1, 3, 5, 20, 2, 4, 6, 8]
826 underten = lambda x: x<10
827 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000828 self.assertEqual(list(takewhile(underten, [])), [])
829 self.assertRaises(TypeError, takewhile)
830 self.assertRaises(TypeError, takewhile, operator.pow)
831 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000832 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
833 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000834 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
835 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +0000836 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000837
838 def test_dropwhile(self):
839 data = [1, 3, 5, 20, 2, 4, 6, 8]
840 underten = lambda x: x<10
841 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000842 self.assertEqual(list(dropwhile(underten, [])), [])
843 self.assertRaises(TypeError, dropwhile)
844 self.assertRaises(TypeError, dropwhile, operator.pow)
845 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000846 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
847 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000848
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000849 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000850 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000851 def irange(n):
Guido van Rossum805365e2007-05-07 22:24:25 +0000852 for i in range(n):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000853 yield i
854
855 a, b = tee([]) # test empty iterator
856 self.assertEqual(list(a), [])
857 self.assertEqual(list(b), [])
858
859 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +0000860 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000861
862 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +0000863 self.assertEqual(list(a), list(range(n)))
864 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000865
866 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000867 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000868 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000869 del a
Guido van Rossum805365e2007-05-07 22:24:25 +0000870 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000871
872 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000873 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000874 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000875 del b
Guido van Rossum805365e2007-05-07 22:24:25 +0000876 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000877
Guido van Rossum805365e2007-05-07 22:24:25 +0000878 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000879 order = [0]*n + [1]*n
880 random.shuffle(order)
881 lists = ([], [])
882 its = tee(irange(n))
883 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +0000884 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000885 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +0000886 self.assertEqual(lists[0], list(range(n)))
887 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000888
Raymond Hettingerad983e72003-11-12 14:32:26 +0000889 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000890 self.assertRaises(TypeError, tee)
891 self.assertRaises(TypeError, tee, 3)
892 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000893 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000894
Raymond Hettingerad983e72003-11-12 14:32:26 +0000895 # tee object should be instantiable
896 a, b = tee('abc')
897 c = type(a)('def')
898 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000899
Raymond Hettingerad983e72003-11-12 14:32:26 +0000900 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +0000901 a, b, c = tee(range(2000), 3)
902 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000903 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +0000904 self.assertEqual(list(b), list(range(2000)))
905 self.assertEqual([next(c), next(c)], list(range(2)))
906 self.assertEqual(list(a), list(range(100,2000)))
907 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000908
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000909 # test values of n
910 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000911 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +0000912 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000913 result = tee('abc', n)
914 self.assertEqual(type(result), tuple)
915 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000916 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000917
Raymond Hettingerad983e72003-11-12 14:32:26 +0000918 # tee pass-through to copyable iterator
919 a, b = tee('abc')
920 c, d = tee(a)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000921 self.assertTrue(a is c)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000922
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000923 # test tee_new
924 t1, t2 = tee('abc')
925 tnew = type(t1)
926 self.assertRaises(TypeError, tnew)
927 self.assertRaises(TypeError, tnew, 10)
928 t3 = tnew(t1)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000929 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000930
Raymond Hettingera9f60922004-10-17 16:40:14 +0000931 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +0000932 a, b = tee(range(10))
Raymond Hettingera9f60922004-10-17 16:40:14 +0000933 p = proxy(a)
934 self.assertEqual(getattr(p, '__class__'), type(b))
935 del a
936 self.assertRaises(ReferenceError, getattr, p, '__class__')
937
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000938 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000939 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000940
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000941 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +0000942 self.assertRaises(StopIteration, next, f([]))
943 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000944
Georg Brandla18af4e2007-04-21 15:47:16 +0000945 self.assertRaises(StopIteration, next, islice([], None))
946 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000947
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000948 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +0000949 self.assertRaises(StopIteration, next, p)
950 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000951 p, q = tee(StopNow())
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
Georg Brandla18af4e2007-04-21 15:47:16 +0000955 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000956
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000957 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +0000958 self.assertRaises(StopIteration, next, f(lambda x:x, []))
959 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000960
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000961class TestExamples(unittest.TestCase):
962
Raymond Hettinger482ba772010-12-01 22:48:00 +0000963 def test_accumlate(self):
964 self.assertEqual(list(accumulate([1,2,3,4,5])), [1, 3, 6, 10, 15])
965
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000966 def test_chain(self):
967 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
968
969 def test_chain_from_iterable(self):
970 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
971
972 def test_combinations(self):
973 self.assertEqual(list(combinations('ABCD', 2)),
974 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
975 self.assertEqual(list(combinations(range(4), 3)),
976 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
977
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000978 def test_combinations_with_replacement(self):
979 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
980 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
981
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000982 def test_compress(self):
983 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
984
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000985 def test_count(self):
986 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
987
988 def test_cycle(self):
989 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
990
991 def test_dropwhile(self):
992 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
993
994 def test_groupby(self):
995 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
996 list('ABCDAB'))
997 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
998 [list('AAAA'), list('BBB'), list('CC'), list('D')])
999
1000 def test_filter(self):
1001 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
1002
1003 def test_filterfalse(self):
1004 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
1005
1006 def test_map(self):
1007 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
1008
1009 def test_islice(self):
1010 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
1011 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
1012 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
1013 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1014
1015 def test_zip(self):
1016 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
1017
1018 def test_zip_longest(self):
1019 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
1020 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
1021
1022 def test_permutations(self):
1023 self.assertEqual(list(permutations('ABCD', 2)),
1024 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
1025 self.assertEqual(list(permutations(range(3))),
1026 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
1027
1028 def test_product(self):
1029 self.assertEqual(list(product('ABCD', 'xy')),
1030 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
1031 self.assertEqual(list(product(range(2), repeat=3)),
1032 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
1033 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
1034
1035 def test_repeat(self):
1036 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
1037
1038 def test_stapmap(self):
1039 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
1040 [32, 9, 1000])
1041
1042 def test_takewhile(self):
1043 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
1044
1045
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001046class TestGC(unittest.TestCase):
1047
1048 def makecycle(self, iterator, container):
1049 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +00001050 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001051 del container, iterator
1052
Raymond Hettinger482ba772010-12-01 22:48:00 +00001053 def test_accumulate(self):
1054 a = []
1055 self.makecycle(accumulate([1,2,a,3]), a)
1056
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001057 def test_chain(self):
1058 a = []
1059 self.makecycle(chain(a), a)
1060
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001061 def test_chain_from_iterable(self):
1062 a = []
1063 self.makecycle(chain.from_iterable([a]), a)
1064
1065 def test_combinations(self):
1066 a = []
1067 self.makecycle(combinations([1,2,a,3], 3), a)
1068
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001069 def test_combinations_with_replacement(self):
1070 a = []
1071 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1072
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001073 def test_compress(self):
1074 a = []
1075 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1076
Raymond Hettingerd6280f42009-02-16 20:50:56 +00001077 def test_count(self):
1078 a = []
1079 Int = type('Int', (int,), dict(x=a))
1080 self.makecycle(count(Int(0), Int(1)), a)
1081
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001082 def test_cycle(self):
1083 a = []
1084 self.makecycle(cycle([a]*2), a)
1085
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001086 def test_dropwhile(self):
1087 a = []
1088 self.makecycle(dropwhile(bool, [0, a, a]), a)
1089
1090 def test_groupby(self):
1091 a = []
1092 self.makecycle(groupby([a]*2, lambda x:x), a)
1093
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001094 def test_issue2246(self):
1095 # Issue 2246 -- the _grouper iterator was not included in GC
1096 n = 10
1097 keyfunc = lambda x: x
1098 for i, j in groupby(range(n), key=keyfunc):
1099 keyfunc.__dict__.setdefault('x',[]).append(j)
1100
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001101 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001102 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001103 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001104
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001105 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001106 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001107 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001108
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001109 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001110 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001111 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001112
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001113 def test_zip_longest(self):
1114 a = []
1115 self.makecycle(zip_longest([a]*2, [a]*3), a)
1116 b = [a, None]
1117 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1118
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001119 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001120 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001121 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001122
1123 def test_islice(self):
1124 a = []
1125 self.makecycle(islice([a]*2, None), a)
1126
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001127 def test_permutations(self):
1128 a = []
1129 self.makecycle(permutations([1,2,a,3], 3), a)
1130
1131 def test_product(self):
1132 a = []
1133 self.makecycle(product([1,2,a,3], repeat=3), a)
1134
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001135 def test_repeat(self):
1136 a = []
1137 self.makecycle(repeat(a), a)
1138
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001139 def test_starmap(self):
1140 a = []
1141 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1142
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001143 def test_takewhile(self):
1144 a = []
1145 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1146
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001147def R(seqn):
1148 'Regular generator'
1149 for i in seqn:
1150 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001151
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001152class G:
1153 'Sequence using __getitem__'
1154 def __init__(self, seqn):
1155 self.seqn = seqn
1156 def __getitem__(self, i):
1157 return self.seqn[i]
1158
1159class I:
1160 'Sequence using iterator protocol'
1161 def __init__(self, seqn):
1162 self.seqn = seqn
1163 self.i = 0
1164 def __iter__(self):
1165 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001166 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001167 if self.i >= len(self.seqn): raise StopIteration
1168 v = self.seqn[self.i]
1169 self.i += 1
1170 return v
1171
1172class Ig:
1173 'Sequence using iterator protocol defined with a generator'
1174 def __init__(self, seqn):
1175 self.seqn = seqn
1176 self.i = 0
1177 def __iter__(self):
1178 for val in self.seqn:
1179 yield val
1180
1181class X:
1182 'Missing __getitem__ and __iter__'
1183 def __init__(self, seqn):
1184 self.seqn = seqn
1185 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001186 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001187 if self.i >= len(self.seqn): raise StopIteration
1188 v = self.seqn[self.i]
1189 self.i += 1
1190 return v
1191
1192class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001193 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001194 def __init__(self, seqn):
1195 self.seqn = seqn
1196 self.i = 0
1197 def __iter__(self):
1198 return self
1199
1200class E:
1201 'Test propagation of exceptions'
1202 def __init__(self, seqn):
1203 self.seqn = seqn
1204 self.i = 0
1205 def __iter__(self):
1206 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001207 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001208 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001209
1210class S:
1211 'Test immediate stop'
1212 def __init__(self, seqn):
1213 pass
1214 def __iter__(self):
1215 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001216 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001217 raise StopIteration
1218
1219def L(seqn):
1220 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001221 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001222
1223
1224class TestVariousIteratorArgs(unittest.TestCase):
1225
Raymond Hettinger482ba772010-12-01 22:48:00 +00001226 def test_accumulate(self):
1227 s = [1,2,3,4,5]
1228 r = [1,3,6,10,15]
1229 n = len(s)
1230 for g in (G, I, Ig, L, R):
1231 self.assertEqual(list(accumulate(g(s))), r)
1232 self.assertEqual(list(accumulate(S(s))), [])
1233 self.assertRaises(TypeError, accumulate, X(s))
1234 self.assertRaises(TypeError, accumulate, N(s))
1235 self.assertRaises(ZeroDivisionError, list, accumulate(E(s)))
1236
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001237 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001238 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001239 for g in (G, I, Ig, S, L, R):
1240 self.assertEqual(list(chain(g(s))), list(g(s)))
1241 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001242 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001243 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001244 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1245
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001246 def test_compress(self):
1247 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1248 n = len(s)
1249 for g in (G, I, Ig, S, L, R):
1250 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1251 self.assertRaises(TypeError, compress, X(s), repeat(1))
1252 self.assertRaises(TypeError, compress, N(s), repeat(1))
1253 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1254
Christian Heimesc3f30c42008-02-22 16:37:40 +00001255 def test_product(self):
1256 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1257 self.assertRaises(TypeError, product, X(s))
1258 self.assertRaises(TypeError, product, N(s))
1259 self.assertRaises(ZeroDivisionError, product, E(s))
1260
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001261 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001262 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001263 for g in (G, I, Ig, S, L, R):
1264 tgtlen = len(s) * 3
1265 expected = list(g(s))*3
1266 actual = list(islice(cycle(g(s)), tgtlen))
1267 self.assertEqual(actual, expected)
1268 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001269 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001270 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1271
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001272 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001273 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001274 for g in (G, I, Ig, S, L, R):
1275 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1276 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001277 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001278 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1279
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001280 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001281 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001282 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001283 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001284 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001285 self.assertRaises(TypeError, filter, isEven, X(s))
1286 self.assertRaises(TypeError, filter, isEven, N(s))
1287 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001288
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001289 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001290 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001291 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001292 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001293 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001294 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1295 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1296 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001297
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001298 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001299 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001300 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001301 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1302 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1303 self.assertRaises(TypeError, zip, X(s))
1304 self.assertRaises(TypeError, zip, N(s))
1305 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001306
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001307 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001308 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001309 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001310 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1311 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1312 self.assertRaises(TypeError, zip_longest, X(s))
1313 self.assertRaises(TypeError, zip_longest, N(s))
1314 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001315
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001316 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001317 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001318 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001319 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001320 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001321 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001322 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001323 self.assertRaises(TypeError, map, onearg, X(s))
1324 self.assertRaises(TypeError, map, onearg, N(s))
1325 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001326
1327 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001328 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001329 for g in (G, I, Ig, S, L, R):
1330 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1331 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001332 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001333 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1334
1335 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001336 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001337 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001338 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001339 self.assertEqual(list(starmap(operator.pow, g(ss))),
1340 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001341 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001342 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001343 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1344
1345 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001346 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001347 for g in (G, I, Ig, S, L, R):
1348 tgt = []
1349 for elem in g(s):
1350 if not isEven(elem): break
1351 tgt.append(elem)
1352 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1353 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001354 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001355 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1356
1357 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001358 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001359 for g in (G, I, Ig, S, L, R):
1360 tgt = []
1361 for elem in g(s):
1362 if not tgt and isOdd(elem): continue
1363 tgt.append(elem)
1364 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1365 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001366 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001367 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1368
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001369 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001370 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001371 for g in (G, I, Ig, S, L, R):
1372 it1, it2 = tee(g(s))
1373 self.assertEqual(list(it1), list(g(s)))
1374 self.assertEqual(list(it2), list(g(s)))
1375 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001376 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001377 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1378
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001379class LengthTransparency(unittest.TestCase):
1380
1381 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001382 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001383 self.assertEqual(len(repeat(None, 50)), 50)
1384 self.assertRaises(TypeError, len, repeat(None))
1385
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001386class RegressionTests(unittest.TestCase):
1387
1388 def test_sf_793826(self):
1389 # Fix Armin Rigo's successful efforts to wreak havoc
1390
1391 def mutatingtuple(tuple1, f, tuple2):
1392 # this builds a tuple t which is a copy of tuple1,
1393 # then calls f(t), then mutates t to be equal to tuple2
1394 # (needs len(tuple1) == len(tuple2)).
1395 def g(value, first=[1]):
1396 if first:
1397 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00001398 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001399 return value
1400 items = list(tuple2)
1401 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001402 gen = map(g, items)
1403 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00001404 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001405
1406 def f(t):
1407 global T
1408 T = t
1409 first[:] = list(T)
1410
1411 first = []
1412 mutatingtuple((1,2,3), f, (4,5,6))
1413 second = list(T)
1414 self.assertEqual(first, second)
1415
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001416
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001417 def test_sf_950057(self):
1418 # Make sure that chain() and cycle() catch exceptions immediately
1419 # rather than when shifting between input sources
1420
1421 def gen1():
1422 hist.append(0)
1423 yield 1
1424 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001425 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001426 hist.append(2)
1427
1428 def gen2(x):
1429 hist.append(3)
1430 yield 2
1431 hist.append(4)
1432 if x:
1433 raise StopIteration
1434
1435 hist = []
1436 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1437 self.assertEqual(hist, [0,1])
1438
1439 hist = []
1440 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1441 self.assertEqual(hist, [0,1])
1442
1443 hist = []
1444 self.assertRaises(AssertionError, list, cycle(gen1()))
1445 self.assertEqual(hist, [0,1])
1446
Thomas Woutersb2137042007-02-01 18:02:27 +00001447class SubclassWithKwargsTest(unittest.TestCase):
1448 def test_keywords_in_subclass(self):
1449 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001450 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001451 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00001452 class Subclass(cls):
1453 def __init__(self, newarg=None, *args):
1454 cls.__init__(self, *args)
1455 try:
1456 Subclass(newarg=1)
1457 except TypeError as err:
1458 # we expect type errors because of wrong argument count
Ezio Melottib58e0bd2010-01-23 15:40:09 +00001459 self.assertNotIn("does not take keyword arguments", err.args[0])
Thomas Woutersb2137042007-02-01 18:02:27 +00001460
1461
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001462libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001463
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001464
1465>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001466>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00001467... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001468...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001469Check 1200 is for $120.15
1470Check 1201 is for $764.05
1471Check 1202 is for $823.14
1472
1473>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001474>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00001475... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00001476...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000014771
14788
147927
1480
1481>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001482>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00001483... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00001484...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001485Alex
1486Laura
1487Martin
1488Walter
1489Samuele
1490
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001491>>> from operator import itemgetter
1492>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00001493>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001494>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001495... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001496...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000014971 ['a', 'c', 'e']
14982 ['b', 'd', 'f']
14993 ['g']
1500
Raymond Hettinger734fb572004-01-20 20:04:40 +00001501# Find runs of consecutive numbers using groupby. The key to the solution
1502# is differencing with a range so that consecutive numbers all appear in
1503# same group.
1504>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00001505>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001506... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001507...
Raymond Hettinger734fb572004-01-20 20:04:40 +00001508[1]
1509[4, 5, 6]
1510[10]
1511[15, 16, 17, 18]
1512[22]
1513[25, 26, 27, 28]
1514
Georg Brandl3dbca812008-07-23 16:10:53 +00001515>>> def take(n, iterable):
1516... "Return first n items of the iterable as a list"
1517... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001518
Georg Brandl3dbca812008-07-23 16:10:53 +00001519>>> def enumerate(iterable, start=0):
1520... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001521
Georg Brandl3dbca812008-07-23 16:10:53 +00001522>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001523... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00001524... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001525
Raymond Hettingercdf8ba32009-02-19 04:45:07 +00001526>>> def nth(iterable, n, default=None):
1527... "Returns the nth item or a default value"
1528... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001529
Georg Brandl3dbca812008-07-23 16:10:53 +00001530>>> def quantify(iterable, pred=bool):
1531... "Count how many times the predicate is true"
1532... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001533
Georg Brandl3dbca812008-07-23 16:10:53 +00001534>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001535... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00001536... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001537
Georg Brandl3dbca812008-07-23 16:10:53 +00001538>>> def ncycles(iterable, n):
Ezio Melotti13925002011-03-16 11:05:33 +02001539... "Returns the sequence elements n times"
Georg Brandl3dbca812008-07-23 16:10:53 +00001540... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001541
1542>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001543... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001544
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001545>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001546... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001547
1548>>> def repeatfunc(func, times=None, *args):
1549... "Repeat calls to func with specified arguments."
1550... " Example: repeatfunc(random.random)"
1551... if times is None:
1552... return starmap(func, repeat(args))
1553... else:
1554... return starmap(func, repeat(args, times))
1555
Raymond Hettingerd591f662003-10-26 15:34:50 +00001556>>> def pairwise(iterable):
1557... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1558... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001559... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00001560... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001561... except StopIteration:
1562... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001563... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001564
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001565>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001566... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001567... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00001568... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001569
1570>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001571... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001572... # Recipe credited to George Sakkis
1573... pending = len(iterables)
1574... nexts = cycle(iter(it).__next__ for it in iterables)
1575... while pending:
1576... try:
1577... for next in nexts:
1578... yield next()
1579... except StopIteration:
1580... pending -= 1
1581... nexts = cycle(islice(nexts, pending))
1582
1583>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00001584... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1585... s = list(iterable)
1586... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001587
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001588>>> def unique_everseen(iterable, key=None):
1589... "List unique elements, preserving order. Remember all elements ever seen."
1590... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1591... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1592... seen = set()
1593... seen_add = seen.add
1594... if key is None:
1595... for element in iterable:
1596... if element not in seen:
1597... seen_add(element)
1598... yield element
1599... else:
1600... for element in iterable:
1601... k = key(element)
1602... if k not in seen:
1603... seen_add(k)
1604... yield element
1605
1606>>> def unique_justseen(iterable, key=None):
1607... "List unique elements, preserving order. Remember only the element just seen."
1608... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1609... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1610... return map(next, map(itemgetter(1), groupby(iterable, key)))
1611
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001612This is not part of the examples but it tests to make sure the definitions
1613perform as purported.
1614
Raymond Hettingera098b332003-09-08 23:58:40 +00001615>>> take(10, count())
1616[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1617
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001618>>> list(enumerate('abc'))
1619[(0, 'a'), (1, 'b'), (2, 'c')]
1620
1621>>> list(islice(tabulate(lambda x: 2*x), 4))
1622[0, 2, 4, 6]
1623
1624>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00001625'd'
1626
1627>>> nth('abcde', 9) is None
1628True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001629
Guido van Rossum805365e2007-05-07 22:24:25 +00001630>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000163150
1632
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001633>>> a = [[1, 2, 3], [4, 5, 6]]
1634>>> flatten(a)
1635[1, 2, 3, 4, 5, 6]
1636
1637>>> list(repeatfunc(pow, 5, 2, 3))
1638[8, 8, 8, 8, 8]
1639
1640>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001641>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001642[0, 0, 0, 0, 0]
1643
Raymond Hettingerd591f662003-10-26 15:34:50 +00001644>>> list(pairwise('abcd'))
1645[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001646
Raymond Hettingerd591f662003-10-26 15:34:50 +00001647>>> list(pairwise([]))
1648[]
1649
1650>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001651[]
1652
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001653>>> list(islice(padnone('abc'), 0, 6))
1654['a', 'b', 'c', None, None, None]
1655
1656>>> list(ncycles('abc', 3))
1657['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1658
1659>>> dotproduct([1,2,3], [4,5,6])
166032
1661
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001662>>> list(grouper(3, 'abcdefg', 'x'))
1663[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1664
1665>>> list(roundrobin('abc', 'd', 'ef'))
1666['a', 'd', 'e', 'b', 'f', 'c']
1667
Raymond Hettingerace67332009-01-26 02:23:50 +00001668>>> list(powerset([1,2,3]))
1669[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001670
Raymond Hettinger191e8502009-01-27 13:29:43 +00001671>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1672True
1673
1674>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1675True
1676
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001677>>> list(unique_everseen('AAAABBBCCDAABBB'))
1678['A', 'B', 'C', 'D']
1679
1680>>> list(unique_everseen('ABBCcAD', str.lower))
1681['A', 'B', 'C', 'D']
1682
1683>>> list(unique_justseen('AAAABBBCCDAABBB'))
1684['A', 'B', 'C', 'D', 'A', 'B']
1685
1686>>> list(unique_justseen('ABBCcAD', str.lower))
1687['A', 'B', 'C', 'A', 'D']
1688
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001689"""
1690
1691__test__ = {'libreftest' : libreftest}
1692
1693def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001694 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00001695 RegressionTests, LengthTransparency,
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001696 SubclassWithKwargsTest, TestExamples)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001697 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001698
1699 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001700 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001701 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001702 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00001703 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001704 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001705 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001706 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001707 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001708
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001709 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001710 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001711
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001712if __name__ == "__main__":
1713 test_main(verbose=True)