blob: f4928caac5c06e6a8a325b796c6a2058112458d1 [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)))
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +0000370 self.assertEqual(list(islice(count(10, maxsize+5), 3)),
371 list(range(10, 10+3*(maxsize+5), maxsize+5)))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000372 c = count(3)
373 self.assertEqual(repr(c), 'count(3)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000374 next(c)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000375 self.assertEqual(repr(c), 'count(4)')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000376 c = count(-9)
377 self.assertEqual(repr(c), 'count(-9)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000378 next(c)
Raymond Hettinger30729212009-02-12 06:28:27 +0000379 self.assertEqual(repr(count(10.25)), 'count(10.25)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000380 self.assertEqual(next(c), -8)
Christian Heimesa37d4c62007-12-04 23:02:19 +0000381 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 +0000382 # Test repr (ignoring the L in longs)
383 r1 = repr(count(i)).replace('L', '')
384 r2 = 'count(%r)'.__mod__(i).replace('L', '')
385 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000386
Raymond Hettingera3e1ad22009-11-30 22:02:31 +0000387 # check copy, deepcopy, pickle
388 for value in -3, 3, maxsize-5, maxsize+5:
389 c = count(value)
390 self.assertEqual(next(copy.copy(c)), value)
391 self.assertEqual(next(copy.deepcopy(c)), value)
392 self.assertEqual(next(pickle.loads(pickle.dumps(c))), value)
393
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +0000394 #check proper internal error handling for large "step' sizes
395 count(1, maxsize+5); sys.exc_info()
396
Raymond Hettinger30729212009-02-12 06:28:27 +0000397 def test_count_with_stride(self):
398 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingereb13fdd2009-02-21 22:30:12 +0000399 self.assertEqual(lzip('abc',count(start=2,step=3)),
400 [('a', 2), ('b', 5), ('c', 8)])
401 self.assertEqual(lzip('abc',count(step=-1)),
402 [('a', 0), ('b', -1), ('c', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000403 self.assertEqual(lzip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
404 self.assertEqual(lzip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000405 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000406 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
407 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
408 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettinger8a882a72009-02-12 12:08:18 +0000409 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
410 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerde09acf2009-02-12 12:53:53 +0000411 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
412 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000413 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
414 c = count(3, 5)
415 self.assertEqual(repr(c), 'count(3, 5)')
416 next(c)
417 self.assertEqual(repr(c), 'count(8, 5)')
418 c = count(-9, 0)
419 self.assertEqual(repr(c), 'count(-9, 0)')
420 next(c)
421 self.assertEqual(repr(c), 'count(-9, 0)')
422 c = count(-9, -3)
423 self.assertEqual(repr(c), 'count(-9, -3)')
424 next(c)
425 self.assertEqual(repr(c), 'count(-12, -3)')
426 self.assertEqual(repr(c), 'count(-12, -3)')
427 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
428 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
429 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
430 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
431 for j in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 1, 10, sys.maxsize-5, sys.maxsize+5):
432 # Test repr (ignoring the L in longs)
433 r1 = repr(count(i, j)).replace('L', '')
434 if j == 1:
435 r2 = ('count(%r)' % i).replace('L', '')
436 else:
437 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
438 self.assertEqual(r1, r2)
439
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000440 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000441 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000442 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000443 self.assertRaises(TypeError, cycle)
444 self.assertRaises(TypeError, cycle, 5)
445 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000446
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000447 def test_groupby(self):
448 # Check whether it accepts arguments correctly
449 self.assertEqual([], list(groupby([])))
450 self.assertEqual([], list(groupby([], key=id)))
451 self.assertRaises(TypeError, list, groupby('abc', []))
452 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000453 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000454
455 # Check normal input
456 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
457 (2,15,22), (3,16,23), (3,17,23)]
458 dup = []
459 for k, g in groupby(s, lambda r:r[0]):
460 for elem in g:
461 self.assertEqual(k, elem[0])
462 dup.append(elem)
463 self.assertEqual(s, dup)
464
465 # Check nested case
466 dup = []
467 for k, g in groupby(s, lambda r:r[0]):
468 for ik, ig in groupby(g, lambda r:r[2]):
469 for elem in ig:
470 self.assertEqual(k, elem[0])
471 self.assertEqual(ik, elem[2])
472 dup.append(elem)
473 self.assertEqual(s, dup)
474
475 # Check case where inner iterator is not used
476 keys = [k for k, g in groupby(s, lambda r:r[0])]
477 expectedkeys = set([r[0] for r in s])
478 self.assertEqual(set(keys), expectedkeys)
479 self.assertEqual(len(keys), len(expectedkeys))
480
481 # Exercise pipes and filters style
482 s = 'abracadabra'
483 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000484 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000485 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
486 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000487 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000488 self.assertEqual(r, ['a', 'b', 'r'])
489 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000490 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000491 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
492 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000493 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000494 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
495
Georg Brandla18af4e2007-04-21 15:47:16 +0000496 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000497 class ExpectedError(Exception):
498 pass
499 def delayed_raise(n=0):
500 for i in range(n):
501 yield 'yo'
502 raise ExpectedError
503 def gulp(iterable, keyp=None, func=list):
504 return [func(g) for k, g in groupby(iterable, keyp)]
505
Georg Brandla18af4e2007-04-21 15:47:16 +0000506 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000507 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000508 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000509 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
510
511 # __cmp__ failure
512 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000513 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000514 raise ExpectedError
515 s = [DummyCmp(), DummyCmp(), None]
516
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000517 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000518 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000519 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000520 self.assertRaises(ExpectedError, gulp, s)
521
522 # keyfunc failure
523 def keyfunc(obj):
524 if keyfunc.skip > 0:
525 keyfunc.skip -= 1
526 return obj
527 else:
528 raise ExpectedError
529
530 # keyfunc failure on outer object
531 keyfunc.skip = 0
532 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
533 keyfunc.skip = 1
534 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
535
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000536 def test_filter(self):
537 self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
538 self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])
539 self.assertEqual(list(filter(bool, [0,1,0,2,0])), [1,2])
540 self.assertEqual(take(4, filter(isEven, count())), [0,2,4,6])
541 self.assertRaises(TypeError, filter)
542 self.assertRaises(TypeError, filter, lambda x:x)
543 self.assertRaises(TypeError, filter, lambda x:x, range(6), 7)
544 self.assertRaises(TypeError, filter, isEven, 3)
545 self.assertRaises(TypeError, next, filter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000546
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000547 def test_filterfalse(self):
548 self.assertEqual(list(filterfalse(isEven, range(6))), [1,3,5])
549 self.assertEqual(list(filterfalse(None, [0,1,0,2,0])), [0,0,0])
550 self.assertEqual(list(filterfalse(bool, [0,1,0,2,0])), [0,0,0])
551 self.assertEqual(take(4, filterfalse(isEven, count())), [1,3,5,7])
552 self.assertRaises(TypeError, filterfalse)
553 self.assertRaises(TypeError, filterfalse, lambda x:x)
554 self.assertRaises(TypeError, filterfalse, lambda x:x, range(6), 7)
555 self.assertRaises(TypeError, filterfalse, isEven, 3)
556 self.assertRaises(TypeError, next, filterfalse(range(6), range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000557
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000558 def test_zip(self):
559 # XXX This is rather silly now that builtin zip() calls zip()...
560 ans = [(x,y) for x, y in zip('abc',count())]
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000561 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000562 self.assertEqual(list(zip('abc', range(6))), lzip('abc', range(6)))
563 self.assertEqual(list(zip('abcdef', range(3))), lzip('abcdef', range(3)))
564 self.assertEqual(take(3,zip('abcdef', count())), lzip('abcdef', range(3)))
565 self.assertEqual(list(zip('abcdef')), lzip('abcdef'))
566 self.assertEqual(list(zip()), lzip())
567 self.assertRaises(TypeError, zip, 3)
568 self.assertRaises(TypeError, zip, range(3), 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000569 # Check tuple re-use (implementation detail)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000570 self.assertEqual([tuple(list(pair)) for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000571 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000572 self.assertEqual([pair for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000573 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000574 ids = list(map(id, zip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000575 self.assertEqual(min(ids), max(ids))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000576 ids = list(map(id, list(zip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000577 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000578
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000579 def test_ziplongest(self):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000580 for args in [
581 ['abc', range(6)],
582 [range(6), 'abc'],
583 [range(1000), range(2000,2100), range(3000,3050)],
584 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
585 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
586 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000587 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
588 for i in range(max(map(len, args)))]
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000589 self.assertEqual(list(zip_longest(*args)), target)
590 self.assertEqual(list(zip_longest(*args, **{})), target)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000591 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 +0000592 self.assertEqual(list(zip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000593
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000594 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 +0000595
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000596 self.assertEqual(list(zip_longest()), list(zip()))
597 self.assertEqual(list(zip_longest([])), list(zip([])))
598 self.assertEqual(list(zip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000599
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000600 self.assertEqual(list(zip_longest('abc', 'defg', **{})),
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000601 list(zip(list('abc')+[None], 'defg'))) # empty keyword dict
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000602 self.assertRaises(TypeError, zip_longest, 3)
603 self.assertRaises(TypeError, zip_longest, range(3), 3)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000604
605 for stmt in [
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000606 "zip_longest('abc', fv=1)",
607 "zip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000608 ]:
609 try:
610 eval(stmt, globals(), locals())
611 except TypeError:
612 pass
613 else:
614 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000615
Thomas Wouterscf297e42007-02-23 15:07:44 +0000616 # Check tuple re-use (implementation detail)
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000617 self.assertEqual([tuple(list(pair)) for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000618 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000619 self.assertEqual([pair for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000620 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000621 ids = list(map(id, zip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000622 self.assertEqual(min(ids), max(ids))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000623 ids = list(map(id, list(zip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000624 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
625
Raymond Hettingerfc438512009-11-01 20:55:33 +0000626 def test_bug_7244(self):
627
628 class Repeater:
629 # this class is similar to itertools.repeat
630 def __init__(self, o, t, e):
631 self.o = o
632 self.t = int(t)
633 self.e = e
634 def __iter__(self): # its iterator is itself
635 return self
636 def __next__(self):
637 if self.t > 0:
638 self.t -= 1
639 return self.o
640 else:
641 raise self.e
642
643 # Formerly this code in would fail in debug mode
644 # with Undetected Error and Stop Iteration
645 r1 = Repeater(1, 3, StopIteration)
646 r2 = Repeater(2, 4, StopIteration)
647 def run(r1, r2):
648 result = []
649 for i, j in zip_longest(r1, r2, fillvalue=0):
650 with support.captured_output('stdout'):
651 print((i, j))
652 result.append((i, j))
653 return result
654 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
655
656 # Formerly, the RuntimeError would be lost
657 # and StopIteration would stop as expected
658 r1 = Repeater(1, 3, RuntimeError)
659 r2 = Repeater(2, 4, StopIteration)
660 it = zip_longest(r1, r2, fillvalue=0)
661 self.assertEqual(next(it), (1, 2))
662 self.assertEqual(next(it), (1, 2))
663 self.assertEqual(next(it), (1, 2))
664 self.assertRaises(RuntimeError, next, it)
665
Christian Heimesc3f30c42008-02-22 16:37:40 +0000666 def test_product(self):
667 for args, result in [
Christian Heimes78644762008-03-04 23:39:23 +0000668 ([], [()]), # zero iterables
Christian Heimesc3f30c42008-02-22 16:37:40 +0000669 (['ab'], [('a',), ('b',)]), # one iterable
670 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
671 ([range(0), range(2), range(3)], []), # first iterable with zero length
672 ([range(2), range(0), range(3)], []), # middle iterable with zero length
673 ([range(2), range(3), range(0)], []), # last iterable with zero length
674 ]:
675 self.assertEqual(list(product(*args)), result)
Christian Heimesf16baeb2008-02-29 14:57:44 +0000676 for r in range(4):
677 self.assertEqual(list(product(*(args*r))),
678 list(product(*args, **dict(repeat=r))))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000679 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
680 self.assertRaises(TypeError, product, range(6), None)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000681
682 def product1(*args, **kwds):
683 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
684 n = len(pools)
685 if n == 0:
686 yield ()
687 return
688 if any(len(pool) == 0 for pool in pools):
689 return
690 indices = [0] * n
691 yield tuple(pool[i] for pool, i in zip(pools, indices))
692 while 1:
693 for i in reversed(range(n)): # right to left
694 if indices[i] == len(pools[i]) - 1:
695 continue
696 indices[i] += 1
697 for j in range(i+1, n):
698 indices[j] = 0
699 yield tuple(pool[i] for pool, i in zip(pools, indices))
700 break
701 else:
702 return
703
704 def product2(*args, **kwds):
705 'Pure python version used in docs'
706 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
707 result = [[]]
708 for pool in pools:
709 result = [x+[y] for x in result for y in pool]
710 for prod in result:
711 yield tuple(prod)
712
Christian Heimesc3f30c42008-02-22 16:37:40 +0000713 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
714 set('abcdefg'), range(11), tuple(range(13))]
715 for i in range(100):
716 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Christian Heimes78644762008-03-04 23:39:23 +0000717 expected_len = prod(map(len, args))
718 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000719 self.assertEqual(list(product(*args)), list(product1(*args)))
720 self.assertEqual(list(product(*args)), list(product2(*args)))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000721 args = map(iter, args)
Christian Heimes78644762008-03-04 23:39:23 +0000722 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000723
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000724 # Test implementation detail: tuple re-use
725 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
726 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000727
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000728 def test_repeat(self):
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +0000729 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
Guido van Rossum805365e2007-05-07 22:24:25 +0000730 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000731 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000732 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000733 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000734 self.assertEqual(list(repeat('a', 0)), [])
735 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000736 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000737 self.assertRaises(TypeError, repeat, None, 3, 4)
738 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000739 r = repeat(1+0j)
740 self.assertEqual(repr(r), 'repeat((1+0j))')
741 r = repeat(1+0j, 5)
742 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
743 list(r)
744 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000745
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000746 def test_map(self):
747 self.assertEqual(list(map(operator.pow, range(3), range(1,7))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000748 [0**1, 1**2, 2**3])
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000749 def tupleize(*args):
750 return args
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000751 self.assertEqual(list(map(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000752 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000753 self.assertEqual(list(map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000754 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000755 self.assertEqual(take(2,map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000756 [('a',0),('b',1)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000757 self.assertEqual(list(map(operator.pow, [])), [])
758 self.assertRaises(TypeError, map)
759 self.assertRaises(TypeError, list, map(None, range(3), range(3)))
760 self.assertRaises(TypeError, map, operator.neg)
761 self.assertRaises(TypeError, next, map(10, range(5)))
762 self.assertRaises(ValueError, next, map(errfunc, [4], [5]))
763 self.assertRaises(TypeError, next, map(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000764
765 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000766 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
767 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000768 self.assertEqual(take(3, starmap(operator.pow, zip(count(), count(1)))),
Raymond Hettinger02420702003-06-29 20:36:23 +0000769 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000770 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +0000771 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
772 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000773 self.assertRaises(TypeError, starmap)
774 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000775 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
776 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
777 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000778
779 def test_islice(self):
780 for args in [ # islice(args) should agree with range(args)
781 (10, 20, 3),
782 (10, 3, 20),
783 (10, 20),
784 (10, 3),
785 (20,)
786 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000787 self.assertEqual(list(islice(range(100), *args)),
788 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000789
790 for args, tgtargs in [ # Stop when seqn is exhausted
791 ((10, 110, 3), ((10, 100, 3))),
792 ((10, 110), ((10, 100))),
793 ((110,), (100,))
794 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000795 self.assertEqual(list(islice(range(100), *args)),
796 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000797
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000798 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +0000799 self.assertEqual(list(islice(range(10), None)), list(range(10)))
800 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
801 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
802 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
803 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000804
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000805 # Test number of items consumed SF #1171417
806 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +0000807 self.assertEqual(list(islice(it, 3)), list(range(3)))
808 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000809
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000810 # Test invalid arguments
Guido van Rossum805365e2007-05-07 22:24:25 +0000811 self.assertRaises(TypeError, islice, range(10))
812 self.assertRaises(TypeError, islice, range(10), 1, 2, 3, 4)
813 self.assertRaises(ValueError, islice, range(10), -5, 10, 1)
814 self.assertRaises(ValueError, islice, range(10), 1, -5, -1)
815 self.assertRaises(ValueError, islice, range(10), 1, 10, -1)
816 self.assertRaises(ValueError, islice, range(10), 1, 10, 0)
817 self.assertRaises(ValueError, islice, range(10), 'a')
818 self.assertRaises(ValueError, islice, range(10), 'a', 1)
819 self.assertRaises(ValueError, islice, range(10), 1, 'a')
820 self.assertRaises(ValueError, islice, range(10), 'a', 1, 1)
821 self.assertRaises(ValueError, islice, range(10), 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +0000822 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000823
Raymond Hettinger69b34bf2010-11-30 02:49:29 +0000824 # Issue #10323: Less islice in a predictable state
825 c = count()
826 self.assertEqual(list(islice(c, 1, 3, 50)), [1])
827 self.assertEqual(next(c), 3)
828
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000829 def test_takewhile(self):
830 data = [1, 3, 5, 20, 2, 4, 6, 8]
831 underten = lambda x: x<10
832 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000833 self.assertEqual(list(takewhile(underten, [])), [])
834 self.assertRaises(TypeError, takewhile)
835 self.assertRaises(TypeError, takewhile, operator.pow)
836 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000837 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
838 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000839 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
840 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +0000841 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000842
843 def test_dropwhile(self):
844 data = [1, 3, 5, 20, 2, 4, 6, 8]
845 underten = lambda x: x<10
846 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000847 self.assertEqual(list(dropwhile(underten, [])), [])
848 self.assertRaises(TypeError, dropwhile)
849 self.assertRaises(TypeError, dropwhile, operator.pow)
850 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000851 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
852 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000853
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000854 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000855 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000856 def irange(n):
Guido van Rossum805365e2007-05-07 22:24:25 +0000857 for i in range(n):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000858 yield i
859
860 a, b = tee([]) # test empty iterator
861 self.assertEqual(list(a), [])
862 self.assertEqual(list(b), [])
863
864 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +0000865 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000866
867 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +0000868 self.assertEqual(list(a), list(range(n)))
869 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000870
871 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000872 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000873 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000874 del a
Guido van Rossum805365e2007-05-07 22:24:25 +0000875 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000876
877 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000878 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000879 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000880 del b
Guido van Rossum805365e2007-05-07 22:24:25 +0000881 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000882
Guido van Rossum805365e2007-05-07 22:24:25 +0000883 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000884 order = [0]*n + [1]*n
885 random.shuffle(order)
886 lists = ([], [])
887 its = tee(irange(n))
888 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +0000889 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000890 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +0000891 self.assertEqual(lists[0], list(range(n)))
892 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000893
Raymond Hettingerad983e72003-11-12 14:32:26 +0000894 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000895 self.assertRaises(TypeError, tee)
896 self.assertRaises(TypeError, tee, 3)
897 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000898 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000899
Raymond Hettingerad983e72003-11-12 14:32:26 +0000900 # tee object should be instantiable
901 a, b = tee('abc')
902 c = type(a)('def')
903 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000904
Raymond Hettingerad983e72003-11-12 14:32:26 +0000905 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +0000906 a, b, c = tee(range(2000), 3)
907 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000908 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +0000909 self.assertEqual(list(b), list(range(2000)))
910 self.assertEqual([next(c), next(c)], list(range(2)))
911 self.assertEqual(list(a), list(range(100,2000)))
912 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000913
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000914 # test values of n
915 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000916 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +0000917 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000918 result = tee('abc', n)
919 self.assertEqual(type(result), tuple)
920 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000921 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000922
Raymond Hettingerad983e72003-11-12 14:32:26 +0000923 # tee pass-through to copyable iterator
924 a, b = tee('abc')
925 c, d = tee(a)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000926 self.assertTrue(a is c)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000927
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000928 # test tee_new
929 t1, t2 = tee('abc')
930 tnew = type(t1)
931 self.assertRaises(TypeError, tnew)
932 self.assertRaises(TypeError, tnew, 10)
933 t3 = tnew(t1)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000934 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000935
Raymond Hettingera9f60922004-10-17 16:40:14 +0000936 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +0000937 a, b = tee(range(10))
Raymond Hettingera9f60922004-10-17 16:40:14 +0000938 p = proxy(a)
939 self.assertEqual(getattr(p, '__class__'), type(b))
940 del a
941 self.assertRaises(ReferenceError, getattr, p, '__class__')
942
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000943 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000944 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000945
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000946 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +0000947 self.assertRaises(StopIteration, next, f([]))
948 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000949
Georg Brandla18af4e2007-04-21 15:47:16 +0000950 self.assertRaises(StopIteration, next, islice([], None))
951 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000952
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000953 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +0000954 self.assertRaises(StopIteration, next, p)
955 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000956 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +0000957 self.assertRaises(StopIteration, next, p)
958 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000959
Georg Brandla18af4e2007-04-21 15:47:16 +0000960 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000961
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000962 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +0000963 self.assertRaises(StopIteration, next, f(lambda x:x, []))
964 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000965
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000966class TestExamples(unittest.TestCase):
967
Raymond Hettinger482ba772010-12-01 22:48:00 +0000968 def test_accumlate(self):
969 self.assertEqual(list(accumulate([1,2,3,4,5])), [1, 3, 6, 10, 15])
970
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000971 def test_chain(self):
972 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
973
974 def test_chain_from_iterable(self):
975 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
976
977 def test_combinations(self):
978 self.assertEqual(list(combinations('ABCD', 2)),
979 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
980 self.assertEqual(list(combinations(range(4), 3)),
981 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
982
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000983 def test_combinations_with_replacement(self):
984 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
985 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
986
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000987 def test_compress(self):
988 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
989
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000990 def test_count(self):
991 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
992
993 def test_cycle(self):
994 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
995
996 def test_dropwhile(self):
997 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
998
999 def test_groupby(self):
1000 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
1001 list('ABCDAB'))
1002 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
1003 [list('AAAA'), list('BBB'), list('CC'), list('D')])
1004
1005 def test_filter(self):
1006 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
1007
1008 def test_filterfalse(self):
1009 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
1010
1011 def test_map(self):
1012 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
1013
1014 def test_islice(self):
1015 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
1016 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
1017 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
1018 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1019
1020 def test_zip(self):
1021 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
1022
1023 def test_zip_longest(self):
1024 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
1025 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
1026
1027 def test_permutations(self):
1028 self.assertEqual(list(permutations('ABCD', 2)),
1029 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
1030 self.assertEqual(list(permutations(range(3))),
1031 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
1032
1033 def test_product(self):
1034 self.assertEqual(list(product('ABCD', 'xy')),
1035 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
1036 self.assertEqual(list(product(range(2), repeat=3)),
1037 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
1038 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
1039
1040 def test_repeat(self):
1041 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
1042
1043 def test_stapmap(self):
1044 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
1045 [32, 9, 1000])
1046
1047 def test_takewhile(self):
1048 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
1049
1050
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001051class TestGC(unittest.TestCase):
1052
1053 def makecycle(self, iterator, container):
1054 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +00001055 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001056 del container, iterator
1057
Raymond Hettinger482ba772010-12-01 22:48:00 +00001058 def test_accumulate(self):
1059 a = []
1060 self.makecycle(accumulate([1,2,a,3]), a)
1061
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001062 def test_chain(self):
1063 a = []
1064 self.makecycle(chain(a), a)
1065
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001066 def test_chain_from_iterable(self):
1067 a = []
1068 self.makecycle(chain.from_iterable([a]), a)
1069
1070 def test_combinations(self):
1071 a = []
1072 self.makecycle(combinations([1,2,a,3], 3), a)
1073
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001074 def test_combinations_with_replacement(self):
1075 a = []
1076 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1077
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001078 def test_compress(self):
1079 a = []
1080 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1081
Raymond Hettingerd6280f42009-02-16 20:50:56 +00001082 def test_count(self):
1083 a = []
1084 Int = type('Int', (int,), dict(x=a))
1085 self.makecycle(count(Int(0), Int(1)), a)
1086
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001087 def test_cycle(self):
1088 a = []
1089 self.makecycle(cycle([a]*2), a)
1090
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001091 def test_dropwhile(self):
1092 a = []
1093 self.makecycle(dropwhile(bool, [0, a, a]), a)
1094
1095 def test_groupby(self):
1096 a = []
1097 self.makecycle(groupby([a]*2, lambda x:x), a)
1098
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001099 def test_issue2246(self):
1100 # Issue 2246 -- the _grouper iterator was not included in GC
1101 n = 10
1102 keyfunc = lambda x: x
1103 for i, j in groupby(range(n), key=keyfunc):
1104 keyfunc.__dict__.setdefault('x',[]).append(j)
1105
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001106 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001107 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001108 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001109
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001110 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001111 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001112 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001113
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001114 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001115 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001116 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001117
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001118 def test_zip_longest(self):
1119 a = []
1120 self.makecycle(zip_longest([a]*2, [a]*3), a)
1121 b = [a, None]
1122 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1123
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001124 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001125 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001126 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001127
1128 def test_islice(self):
1129 a = []
1130 self.makecycle(islice([a]*2, None), a)
1131
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001132 def test_permutations(self):
1133 a = []
1134 self.makecycle(permutations([1,2,a,3], 3), a)
1135
1136 def test_product(self):
1137 a = []
1138 self.makecycle(product([1,2,a,3], repeat=3), a)
1139
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001140 def test_repeat(self):
1141 a = []
1142 self.makecycle(repeat(a), a)
1143
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001144 def test_starmap(self):
1145 a = []
1146 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1147
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001148 def test_takewhile(self):
1149 a = []
1150 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1151
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001152def R(seqn):
1153 'Regular generator'
1154 for i in seqn:
1155 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001156
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001157class G:
1158 'Sequence using __getitem__'
1159 def __init__(self, seqn):
1160 self.seqn = seqn
1161 def __getitem__(self, i):
1162 return self.seqn[i]
1163
1164class I:
1165 'Sequence using iterator protocol'
1166 def __init__(self, seqn):
1167 self.seqn = seqn
1168 self.i = 0
1169 def __iter__(self):
1170 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001171 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001172 if self.i >= len(self.seqn): raise StopIteration
1173 v = self.seqn[self.i]
1174 self.i += 1
1175 return v
1176
1177class Ig:
1178 'Sequence using iterator protocol defined with a generator'
1179 def __init__(self, seqn):
1180 self.seqn = seqn
1181 self.i = 0
1182 def __iter__(self):
1183 for val in self.seqn:
1184 yield val
1185
1186class X:
1187 'Missing __getitem__ and __iter__'
1188 def __init__(self, seqn):
1189 self.seqn = seqn
1190 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001191 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001192 if self.i >= len(self.seqn): raise StopIteration
1193 v = self.seqn[self.i]
1194 self.i += 1
1195 return v
1196
1197class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001198 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001199 def __init__(self, seqn):
1200 self.seqn = seqn
1201 self.i = 0
1202 def __iter__(self):
1203 return self
1204
1205class E:
1206 'Test propagation of exceptions'
1207 def __init__(self, seqn):
1208 self.seqn = seqn
1209 self.i = 0
1210 def __iter__(self):
1211 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001212 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001213 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001214
1215class S:
1216 'Test immediate stop'
1217 def __init__(self, seqn):
1218 pass
1219 def __iter__(self):
1220 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001221 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001222 raise StopIteration
1223
1224def L(seqn):
1225 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001226 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001227
1228
1229class TestVariousIteratorArgs(unittest.TestCase):
1230
Raymond Hettinger482ba772010-12-01 22:48:00 +00001231 def test_accumulate(self):
1232 s = [1,2,3,4,5]
1233 r = [1,3,6,10,15]
1234 n = len(s)
1235 for g in (G, I, Ig, L, R):
1236 self.assertEqual(list(accumulate(g(s))), r)
1237 self.assertEqual(list(accumulate(S(s))), [])
1238 self.assertRaises(TypeError, accumulate, X(s))
1239 self.assertRaises(TypeError, accumulate, N(s))
1240 self.assertRaises(ZeroDivisionError, list, accumulate(E(s)))
1241
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001242 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001243 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001244 for g in (G, I, Ig, S, L, R):
1245 self.assertEqual(list(chain(g(s))), list(g(s)))
1246 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001247 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001248 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001249 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1250
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001251 def test_compress(self):
1252 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1253 n = len(s)
1254 for g in (G, I, Ig, S, L, R):
1255 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1256 self.assertRaises(TypeError, compress, X(s), repeat(1))
1257 self.assertRaises(TypeError, compress, N(s), repeat(1))
1258 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1259
Christian Heimesc3f30c42008-02-22 16:37:40 +00001260 def test_product(self):
1261 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1262 self.assertRaises(TypeError, product, X(s))
1263 self.assertRaises(TypeError, product, N(s))
1264 self.assertRaises(ZeroDivisionError, product, E(s))
1265
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001266 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001267 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001268 for g in (G, I, Ig, S, L, R):
1269 tgtlen = len(s) * 3
1270 expected = list(g(s))*3
1271 actual = list(islice(cycle(g(s)), tgtlen))
1272 self.assertEqual(actual, expected)
1273 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001274 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001275 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1276
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001277 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001278 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001279 for g in (G, I, Ig, S, L, R):
1280 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1281 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001282 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001283 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1284
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001285 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001286 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001287 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001288 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001289 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001290 self.assertRaises(TypeError, filter, isEven, X(s))
1291 self.assertRaises(TypeError, filter, isEven, N(s))
1292 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001293
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001294 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001295 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001296 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001297 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001298 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001299 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1300 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1301 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001302
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001303 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001304 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001305 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001306 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1307 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1308 self.assertRaises(TypeError, zip, X(s))
1309 self.assertRaises(TypeError, zip, N(s))
1310 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001311
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001312 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001313 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001314 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001315 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1316 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1317 self.assertRaises(TypeError, zip_longest, X(s))
1318 self.assertRaises(TypeError, zip_longest, N(s))
1319 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001320
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001321 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001322 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001323 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001324 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001325 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001326 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001327 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001328 self.assertRaises(TypeError, map, onearg, X(s))
1329 self.assertRaises(TypeError, map, onearg, N(s))
1330 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001331
1332 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001333 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001334 for g in (G, I, Ig, S, L, R):
1335 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1336 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001337 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001338 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1339
1340 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001341 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001342 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001343 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001344 self.assertEqual(list(starmap(operator.pow, g(ss))),
1345 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001346 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001347 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001348 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1349
1350 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001351 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001352 for g in (G, I, Ig, S, L, R):
1353 tgt = []
1354 for elem in g(s):
1355 if not isEven(elem): break
1356 tgt.append(elem)
1357 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1358 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001359 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001360 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1361
1362 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001363 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001364 for g in (G, I, Ig, S, L, R):
1365 tgt = []
1366 for elem in g(s):
1367 if not tgt and isOdd(elem): continue
1368 tgt.append(elem)
1369 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1370 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001371 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001372 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1373
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001374 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001375 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001376 for g in (G, I, Ig, S, L, R):
1377 it1, it2 = tee(g(s))
1378 self.assertEqual(list(it1), list(g(s)))
1379 self.assertEqual(list(it2), list(g(s)))
1380 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001381 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001382 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1383
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001384class LengthTransparency(unittest.TestCase):
1385
1386 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001387 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001388 self.assertEqual(len(repeat(None, 50)), 50)
1389 self.assertRaises(TypeError, len, repeat(None))
1390
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001391class RegressionTests(unittest.TestCase):
1392
1393 def test_sf_793826(self):
1394 # Fix Armin Rigo's successful efforts to wreak havoc
1395
1396 def mutatingtuple(tuple1, f, tuple2):
1397 # this builds a tuple t which is a copy of tuple1,
1398 # then calls f(t), then mutates t to be equal to tuple2
1399 # (needs len(tuple1) == len(tuple2)).
1400 def g(value, first=[1]):
1401 if first:
1402 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00001403 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001404 return value
1405 items = list(tuple2)
1406 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001407 gen = map(g, items)
1408 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00001409 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001410
1411 def f(t):
1412 global T
1413 T = t
1414 first[:] = list(T)
1415
1416 first = []
1417 mutatingtuple((1,2,3), f, (4,5,6))
1418 second = list(T)
1419 self.assertEqual(first, second)
1420
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001421
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001422 def test_sf_950057(self):
1423 # Make sure that chain() and cycle() catch exceptions immediately
1424 # rather than when shifting between input sources
1425
1426 def gen1():
1427 hist.append(0)
1428 yield 1
1429 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001430 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001431 hist.append(2)
1432
1433 def gen2(x):
1434 hist.append(3)
1435 yield 2
1436 hist.append(4)
1437 if x:
1438 raise StopIteration
1439
1440 hist = []
1441 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1442 self.assertEqual(hist, [0,1])
1443
1444 hist = []
1445 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1446 self.assertEqual(hist, [0,1])
1447
1448 hist = []
1449 self.assertRaises(AssertionError, list, cycle(gen1()))
1450 self.assertEqual(hist, [0,1])
1451
Thomas Woutersb2137042007-02-01 18:02:27 +00001452class SubclassWithKwargsTest(unittest.TestCase):
1453 def test_keywords_in_subclass(self):
1454 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001455 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001456 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00001457 class Subclass(cls):
1458 def __init__(self, newarg=None, *args):
1459 cls.__init__(self, *args)
1460 try:
1461 Subclass(newarg=1)
1462 except TypeError as err:
1463 # we expect type errors because of wrong argument count
Ezio Melottib58e0bd2010-01-23 15:40:09 +00001464 self.assertNotIn("does not take keyword arguments", err.args[0])
Thomas Woutersb2137042007-02-01 18:02:27 +00001465
1466
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001467libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001468
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001469
1470>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001471>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00001472... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001473...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001474Check 1200 is for $120.15
1475Check 1201 is for $764.05
1476Check 1202 is for $823.14
1477
1478>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001479>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00001480... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00001481...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000014821
14838
148427
1485
1486>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001487>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00001488... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00001489...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001490Alex
1491Laura
1492Martin
1493Walter
1494Samuele
1495
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001496>>> from operator import itemgetter
1497>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00001498>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001499>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001500... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001501...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000015021 ['a', 'c', 'e']
15032 ['b', 'd', 'f']
15043 ['g']
1505
Raymond Hettinger734fb572004-01-20 20:04:40 +00001506# Find runs of consecutive numbers using groupby. The key to the solution
1507# is differencing with a range so that consecutive numbers all appear in
1508# same group.
1509>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00001510>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001511... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001512...
Raymond Hettinger734fb572004-01-20 20:04:40 +00001513[1]
1514[4, 5, 6]
1515[10]
1516[15, 16, 17, 18]
1517[22]
1518[25, 26, 27, 28]
1519
Georg Brandl3dbca812008-07-23 16:10:53 +00001520>>> def take(n, iterable):
1521... "Return first n items of the iterable as a list"
1522... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001523
Georg Brandl3dbca812008-07-23 16:10:53 +00001524>>> def enumerate(iterable, start=0):
1525... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001526
Georg Brandl3dbca812008-07-23 16:10:53 +00001527>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001528... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00001529... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001530
Raymond Hettingercdf8ba32009-02-19 04:45:07 +00001531>>> def nth(iterable, n, default=None):
1532... "Returns the nth item or a default value"
1533... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001534
Georg Brandl3dbca812008-07-23 16:10:53 +00001535>>> def quantify(iterable, pred=bool):
1536... "Count how many times the predicate is true"
1537... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001538
Georg Brandl3dbca812008-07-23 16:10:53 +00001539>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001540... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00001541... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001542
Georg Brandl3dbca812008-07-23 16:10:53 +00001543>>> def ncycles(iterable, n):
Ezio Melotti13925002011-03-16 11:05:33 +02001544... "Returns the sequence elements n times"
Georg Brandl3dbca812008-07-23 16:10:53 +00001545... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001546
1547>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001548... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001549
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001550>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001551... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001552
1553>>> def repeatfunc(func, times=None, *args):
1554... "Repeat calls to func with specified arguments."
1555... " Example: repeatfunc(random.random)"
1556... if times is None:
1557... return starmap(func, repeat(args))
1558... else:
1559... return starmap(func, repeat(args, times))
1560
Raymond Hettingerd591f662003-10-26 15:34:50 +00001561>>> def pairwise(iterable):
1562... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1563... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001564... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00001565... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001566... except StopIteration:
1567... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001568... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001569
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001570>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001571... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001572... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00001573... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001574
1575>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001576... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001577... # Recipe credited to George Sakkis
1578... pending = len(iterables)
1579... nexts = cycle(iter(it).__next__ for it in iterables)
1580... while pending:
1581... try:
1582... for next in nexts:
1583... yield next()
1584... except StopIteration:
1585... pending -= 1
1586... nexts = cycle(islice(nexts, pending))
1587
1588>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00001589... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1590... s = list(iterable)
1591... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001592
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001593>>> def unique_everseen(iterable, key=None):
1594... "List unique elements, preserving order. Remember all elements ever seen."
1595... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1596... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1597... seen = set()
1598... seen_add = seen.add
1599... if key is None:
1600... for element in iterable:
1601... if element not in seen:
1602... seen_add(element)
1603... yield element
1604... else:
1605... for element in iterable:
1606... k = key(element)
1607... if k not in seen:
1608... seen_add(k)
1609... yield element
1610
1611>>> def unique_justseen(iterable, key=None):
1612... "List unique elements, preserving order. Remember only the element just seen."
1613... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1614... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1615... return map(next, map(itemgetter(1), groupby(iterable, key)))
1616
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001617This is not part of the examples but it tests to make sure the definitions
1618perform as purported.
1619
Raymond Hettingera098b332003-09-08 23:58:40 +00001620>>> take(10, count())
1621[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1622
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001623>>> list(enumerate('abc'))
1624[(0, 'a'), (1, 'b'), (2, 'c')]
1625
1626>>> list(islice(tabulate(lambda x: 2*x), 4))
1627[0, 2, 4, 6]
1628
1629>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00001630'd'
1631
1632>>> nth('abcde', 9) is None
1633True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001634
Guido van Rossum805365e2007-05-07 22:24:25 +00001635>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000163650
1637
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001638>>> a = [[1, 2, 3], [4, 5, 6]]
1639>>> flatten(a)
1640[1, 2, 3, 4, 5, 6]
1641
1642>>> list(repeatfunc(pow, 5, 2, 3))
1643[8, 8, 8, 8, 8]
1644
1645>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001646>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001647[0, 0, 0, 0, 0]
1648
Raymond Hettingerd591f662003-10-26 15:34:50 +00001649>>> list(pairwise('abcd'))
1650[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001651
Raymond Hettingerd591f662003-10-26 15:34:50 +00001652>>> list(pairwise([]))
1653[]
1654
1655>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001656[]
1657
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001658>>> list(islice(padnone('abc'), 0, 6))
1659['a', 'b', 'c', None, None, None]
1660
1661>>> list(ncycles('abc', 3))
1662['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1663
1664>>> dotproduct([1,2,3], [4,5,6])
166532
1666
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001667>>> list(grouper(3, 'abcdefg', 'x'))
1668[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1669
1670>>> list(roundrobin('abc', 'd', 'ef'))
1671['a', 'd', 'e', 'b', 'f', 'c']
1672
Raymond Hettingerace67332009-01-26 02:23:50 +00001673>>> list(powerset([1,2,3]))
1674[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001675
Raymond Hettinger191e8502009-01-27 13:29:43 +00001676>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1677True
1678
1679>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1680True
1681
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001682>>> list(unique_everseen('AAAABBBCCDAABBB'))
1683['A', 'B', 'C', 'D']
1684
1685>>> list(unique_everseen('ABBCcAD', str.lower))
1686['A', 'B', 'C', 'D']
1687
1688>>> list(unique_justseen('AAAABBBCCDAABBB'))
1689['A', 'B', 'C', 'D', 'A', 'B']
1690
1691>>> list(unique_justseen('ABBCcAD', str.lower))
1692['A', 'B', 'C', 'A', 'D']
1693
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001694"""
1695
1696__test__ = {'libreftest' : libreftest}
1697
1698def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001699 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00001700 RegressionTests, LengthTransparency,
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001701 SubclassWithKwargsTest, TestExamples)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001702 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001703
1704 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001705 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001706 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001707 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00001708 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001709 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001710 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001711 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001712 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001713
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001714 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001715 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001716
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001717if __name__ == "__main__":
1718 test_main(verbose=True)