blob: 8cdc597e58bc113d73794042e2e6a50f17e8ffde [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001import unittest
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002from test import support
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003from itertools import *
Raymond Hettingera9f60922004-10-17 16:40:14 +00004from weakref import proxy
Raymond Hettinger8a882a72009-02-12 12:08:18 +00005from decimal import Decimal
Raymond Hettingerde09acf2009-02-12 12:53:53 +00006from fractions import Fraction
Raymond Hettinger2012f172003-02-07 05:32:58 +00007import sys
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00008import operator
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00009import random
Raymond Hettingera3e1ad22009-11-30 22:02:31 +000010import copy
11import pickle
Christian Heimesc3f30c42008-02-22 16:37:40 +000012from functools import reduce
Benjamin Petersonee8712c2008-05-20 21:35:26 +000013maxsize = support.MAX_Py_ssize_t
Guido van Rossum360e4b82007-05-14 22:51:27 +000014minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000015
Guido van Rossum801f0d72006-08-24 19:48:10 +000016def lzip(*args):
17 return list(zip(*args))
18
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000019def onearg(x):
20 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000021 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000022
23def errfunc(*args):
24 'Test function that raises an error'
25 raise ValueError
26
27def gen3():
28 'Non-restartable source sequence'
29 for i in (0, 1, 2):
30 yield i
31
32def isEven(x):
33 'Test predicate'
34 return x%2==0
35
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000036def isOdd(x):
37 'Test predicate'
38 return x%2==1
39
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000040class StopNow:
41 'Class emulating an empty iterable.'
42 def __iter__(self):
43 return self
Georg Brandla18af4e2007-04-21 15:47:16 +000044 def __next__(self):
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000045 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000046
Raymond Hettinger02420702003-06-29 20:36:23 +000047def take(n, seq):
48 'Convenience function for partially consuming a long of infinite iterable'
49 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000050
Christian Heimes78644762008-03-04 23:39:23 +000051def prod(iterable):
52 return reduce(operator.mul, iterable, 1)
53
Christian Heimes380f7f22008-02-28 11:19:05 +000054def fact(n):
55 'Factorial'
Christian Heimes78644762008-03-04 23:39:23 +000056 return prod(range(1, n+1))
57
Raymond Hettinger96ef8112003-02-01 00:10:11 +000058class TestBasicOps(unittest.TestCase):
Raymond Hettinger482ba772010-12-01 22:48:00 +000059
60 def test_accumulate(self):
61 self.assertEqual(list(accumulate(range(10))), # one positional arg
Raymond Hettingerd8ff4652010-12-03 02:09:34 +000062 [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])
63 self.assertEqual(list(accumulate(iterable=range(10))), # kw arg
64 [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])
Raymond Hettinger482ba772010-12-01 22:48:00 +000065 for typ in int, complex, Decimal, Fraction: # multiple types
Raymond Hettingerd8ff4652010-12-03 02:09:34 +000066 self.assertEqual(
67 list(accumulate(map(typ, range(10)))),
Raymond Hettinger482ba772010-12-01 22:48:00 +000068 list(map(typ, [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])))
Raymond Hettinger2d93e6e2010-12-03 02:33:53 +000069 self.assertEqual(list(accumulate('abc')), ['a', 'ab', 'abc']) # works with non-numeric
Raymond Hettinger482ba772010-12-01 22:48:00 +000070 self.assertEqual(list(accumulate([])), []) # empty iterable
Raymond Hettingerd8ff4652010-12-03 02:09:34 +000071 self.assertEqual(list(accumulate([7])), [7]) # iterable of length one
72 self.assertRaises(TypeError, accumulate, range(10), 5) # too many args
Raymond Hettinger482ba772010-12-01 22:48:00 +000073 self.assertRaises(TypeError, accumulate) # too few args
Raymond Hettingerd8ff4652010-12-03 02:09:34 +000074 self.assertRaises(TypeError, accumulate, x=range(10)) # unexpected kwd arg
Raymond Hettinger482ba772010-12-01 22:48:00 +000075 self.assertRaises(TypeError, list, accumulate([1, []])) # args that don't add
76
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000077 def test_chain(self):
Christian Heimesdd15f6c2008-03-16 00:07:10 +000078
79 def chain2(*iterables):
80 'Pure python version in the docs'
81 for it in iterables:
82 for element in it:
83 yield element
84
85 for c in (chain, chain2):
86 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
87 self.assertEqual(list(c('abc')), list('abc'))
88 self.assertEqual(list(c('')), [])
89 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
90 self.assertRaises(TypeError, list,c(2, 3))
Christian Heimesf16baeb2008-02-29 14:57:44 +000091
92 def test_chain_from_iterable(self):
93 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
94 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
95 self.assertEqual(list(chain.from_iterable([''])), [])
96 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
97 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000098
Christian Heimes380f7f22008-02-28 11:19:05 +000099 def test_combinations(self):
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000100 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Christian Heimes380f7f22008-02-28 11:19:05 +0000101 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Christian Heimes78644762008-03-04 23:39:23 +0000102 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Christian Heimes380f7f22008-02-28 11:19:05 +0000103 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000104 self.assertEqual(list(combinations('abc', 32)), []) # r > n
Christian Heimes380f7f22008-02-28 11:19:05 +0000105 self.assertEqual(list(combinations(range(4), 3)),
106 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
Christian Heimes78644762008-03-04 23:39:23 +0000107
108 def combinations1(iterable, r):
109 'Pure python version shown in the docs'
110 pool = tuple(iterable)
111 n = len(pool)
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000112 if r > n:
113 return
Christian Heimes78644762008-03-04 23:39:23 +0000114 indices = list(range(r))
115 yield tuple(pool[i] for i in indices)
116 while 1:
117 for i in reversed(range(r)):
118 if indices[i] != i + n - r:
119 break
120 else:
121 return
122 indices[i] += 1
123 for j in range(i+1, r):
124 indices[j] = indices[j-1] + 1
125 yield tuple(pool[i] for i in indices)
126
127 def combinations2(iterable, r):
128 'Pure python version shown in the docs'
129 pool = tuple(iterable)
130 n = len(pool)
131 for indices in permutations(range(n), r):
132 if sorted(indices) == list(indices):
133 yield tuple(pool[i] for i in indices)
134
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000135 def combinations3(iterable, r):
136 'Pure python version from cwr()'
137 pool = tuple(iterable)
138 n = len(pool)
139 for indices in combinations_with_replacement(range(n), r):
140 if len(set(indices)) == r:
141 yield tuple(pool[i] for i in indices)
142
Christian Heimes78644762008-03-04 23:39:23 +0000143 for n in range(7):
Christian Heimes380f7f22008-02-28 11:19:05 +0000144 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000145 for r in range(n+2):
Christian Heimes380f7f22008-02-28 11:19:05 +0000146 result = list(combinations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000147 self.assertEqual(len(result), 0 if r>n else fact(n) / fact(r) / fact(n-r)) # right number of combs
Christian Heimes380f7f22008-02-28 11:19:05 +0000148 self.assertEqual(len(result), len(set(result))) # no repeats
149 self.assertEqual(result, sorted(result)) # lexicographic order
150 for c in result:
151 self.assertEqual(len(c), r) # r-length combinations
152 self.assertEqual(len(set(c)), r) # no duplicate elements
153 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000154 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000155 self.assertEqual(list(c),
156 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000157 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000158 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000159 self.assertEqual(result, list(combinations3(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000160
161 # Test implementation detail: tuple re-use
162 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
163 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
164
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000165 def test_combinations_with_replacement(self):
166 cwr = combinations_with_replacement
167 self.assertRaises(TypeError, cwr, 'abc') # missing r argument
168 self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
169 self.assertRaises(TypeError, cwr, None) # pool is not iterable
170 self.assertRaises(ValueError, cwr, 'abc', -2) # r is negative
171 self.assertEqual(list(cwr('ABC', 2)),
172 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
173
174 def cwr1(iterable, r):
175 'Pure python version shown in the docs'
176 # number items returned: (n+r-1)! / r! / (n-1)! when n>0
177 pool = tuple(iterable)
178 n = len(pool)
179 if not n and r:
180 return
181 indices = [0] * r
182 yield tuple(pool[i] for i in indices)
183 while 1:
184 for i in reversed(range(r)):
185 if indices[i] != n - 1:
186 break
187 else:
188 return
189 indices[i:] = [indices[i] + 1] * (r - i)
190 yield tuple(pool[i] for i in indices)
191
192 def cwr2(iterable, r):
193 'Pure python version shown in the docs'
194 pool = tuple(iterable)
195 n = len(pool)
196 for indices in product(range(n), repeat=r):
197 if sorted(indices) == list(indices):
198 yield tuple(pool[i] for i in indices)
199
200 def numcombs(n, r):
201 if not n:
202 return 0 if r else 1
203 return fact(n+r-1) / fact(r)/ fact(n-1)
204
205 for n in range(7):
206 values = [5*x-12 for x in range(n)]
207 for r in range(n+2):
208 result = list(cwr(values, r))
209
210 self.assertEqual(len(result), numcombs(n, r)) # right number of combs
211 self.assertEqual(len(result), len(set(result))) # no repeats
212 self.assertEqual(result, sorted(result)) # lexicographic order
213
214 regular_combs = list(combinations(values, r)) # compare to combs without replacement
215 if n == 0 or r <= 1:
Ezio Melottib3aedd42010-11-20 19:04:17 +0000216 self.assertEqual(result, regular_combs) # cases that should be identical
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000217 else:
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000218 self.assertTrue(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000219
220 for c in result:
221 self.assertEqual(len(c), r) # r-length combinations
222 noruns = [k for k,v in groupby(c)] # combo without consecutive repeats
223 self.assertEqual(len(noruns), len(set(noruns))) # no repeats other than consecutive
224 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000225 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000226 self.assertEqual(noruns,
227 [e for e in values if e in c]) # comb is a subsequence of the input iterable
228 self.assertEqual(result, list(cwr1(values, r))) # matches first pure python version
229 self.assertEqual(result, list(cwr2(values, r))) # matches second pure python version
230
231 # Test implementation detail: tuple re-use
232 self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
233 self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
234
Christian Heimes78644762008-03-04 23:39:23 +0000235 def test_permutations(self):
236 self.assertRaises(TypeError, permutations) # too few arguments
237 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000238 self.assertRaises(TypeError, permutations, None) # pool is not iterable
239 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000240 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000241 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Christian Heimes78644762008-03-04 23:39:23 +0000242 self.assertEqual(list(permutations(range(3), 2)),
243 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
244
245 def permutations1(iterable, r=None):
246 'Pure python version shown in the docs'
247 pool = tuple(iterable)
248 n = len(pool)
249 r = n if r is None else r
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000250 if r > n:
251 return
Christian Heimes78644762008-03-04 23:39:23 +0000252 indices = list(range(n))
253 cycles = list(range(n-r+1, n+1))[::-1]
254 yield tuple(pool[i] for i in indices[:r])
255 while n:
256 for i in reversed(range(r)):
257 cycles[i] -= 1
258 if cycles[i] == 0:
259 indices[i:] = indices[i+1:] + indices[i:i+1]
260 cycles[i] = n - i
261 else:
262 j = cycles[i]
263 indices[i], indices[-j] = indices[-j], indices[i]
264 yield tuple(pool[i] for i in indices[:r])
265 break
266 else:
267 return
268
269 def permutations2(iterable, r=None):
270 'Pure python version shown in the docs'
271 pool = tuple(iterable)
272 n = len(pool)
273 r = n if r is None else r
274 for indices in product(range(n), repeat=r):
275 if len(set(indices)) == r:
276 yield tuple(pool[i] for i in indices)
277
278 for n in range(7):
279 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000280 for r in range(n+2):
Christian Heimes78644762008-03-04 23:39:23 +0000281 result = list(permutations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000282 self.assertEqual(len(result), 0 if r>n else fact(n) / fact(n-r)) # right number of perms
Christian Heimes78644762008-03-04 23:39:23 +0000283 self.assertEqual(len(result), len(set(result))) # no repeats
284 self.assertEqual(result, sorted(result)) # lexicographic order
285 for p in result:
286 self.assertEqual(len(p), r) # r-length permutations
287 self.assertEqual(len(set(p)), r) # no duplicate elements
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000288 self.assertTrue(all(e in values for e in p)) # elements taken from input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000289 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000290 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000291 if r == n:
292 self.assertEqual(result, list(permutations(values, None))) # test r as None
293 self.assertEqual(result, list(permutations(values))) # test default r
294
295 # Test implementation detail: tuple re-use
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000296 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Christian Heimes78644762008-03-04 23:39:23 +0000297 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Christian Heimes380f7f22008-02-28 11:19:05 +0000298
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000299 def test_combinatorics(self):
300 # Test relationships between product(), permutations(),
301 # combinations() and combinations_with_replacement().
302
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000303 for n in range(6):
304 s = 'ABCDEFG'[:n]
305 for r in range(8):
306 prod = list(product(s, repeat=r))
307 cwr = list(combinations_with_replacement(s, r))
308 perm = list(permutations(s, r))
309 comb = list(combinations(s, r))
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000310
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000311 # Check size
Ezio Melottib3aedd42010-11-20 19:04:17 +0000312 self.assertEqual(len(prod), n**r)
313 self.assertEqual(len(cwr), (fact(n+r-1) / fact(r)/ fact(n-1)) if n else (not r))
314 self.assertEqual(len(perm), 0 if r>n else fact(n) / fact(n-r))
315 self.assertEqual(len(comb), 0 if r>n else fact(n) / fact(r) / fact(n-r))
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000316
317 # Check lexicographic order without repeated tuples
Ezio Melottib3aedd42010-11-20 19:04:17 +0000318 self.assertEqual(prod, sorted(set(prod)))
319 self.assertEqual(cwr, sorted(set(cwr)))
320 self.assertEqual(perm, sorted(set(perm)))
321 self.assertEqual(comb, sorted(set(comb)))
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000322
323 # Check interrelationships
Ezio Melottib3aedd42010-11-20 19:04:17 +0000324 self.assertEqual(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
325 self.assertEqual(perm, [t for t in prod if len(set(t))==r]) # perm: prods with no dups
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000326 self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
327 self.assertEqual(comb, [t for t in cwr if len(set(t))==r]) # comb: cwrs without dups
328 self.assertEqual(comb, list(filter(set(cwr).__contains__, perm))) # comb: perm that is a cwr
329 self.assertEqual(comb, list(filter(set(perm).__contains__, cwr))) # comb: cwr that is a perm
330 self.assertEqual(comb, sorted(set(cwr) & set(perm))) # comb: both a cwr and a perm
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000331
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000332 def test_compress(self):
Raymond Hettinger15a49502009-02-19 02:17:09 +0000333 self.assertEqual(list(compress(data='ABCDEF', selectors=[1,0,1,0,1,1])), list('ACEF'))
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000334 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
335 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
336 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
337 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
338 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
339 n = 10000
340 data = chain.from_iterable(repeat(range(6), n))
341 selectors = chain.from_iterable(repeat((0, 1)))
342 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
343 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
344 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
345 self.assertRaises(TypeError, compress, range(6)) # too few args
346 self.assertRaises(TypeError, compress, range(6), None) # too many args
347
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000348 def test_count(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000349 self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
350 self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
351 self.assertEqual(take(2, lzip('abc',count(3))), [('a', 3), ('b', 4)])
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000352 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
353 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000354 self.assertRaises(TypeError, count, 2, 3, 4)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000355 self.assertRaises(TypeError, count, 'a')
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000356 self.assertEqual(list(islice(count(maxsize-5), 10)),
357 list(range(maxsize-5, maxsize+5)))
358 self.assertEqual(list(islice(count(-maxsize-5), 10)),
359 list(range(-maxsize-5, -maxsize+5)))
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +0000360 self.assertEqual(list(islice(count(10, maxsize+5), 3)),
361 list(range(10, 10+3*(maxsize+5), maxsize+5)))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000362 c = count(3)
363 self.assertEqual(repr(c), 'count(3)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000364 next(c)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000365 self.assertEqual(repr(c), 'count(4)')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000366 c = count(-9)
367 self.assertEqual(repr(c), 'count(-9)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000368 next(c)
Raymond Hettinger30729212009-02-12 06:28:27 +0000369 self.assertEqual(repr(count(10.25)), 'count(10.25)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000370 self.assertEqual(next(c), -8)
Christian Heimesa37d4c62007-12-04 23:02:19 +0000371 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000372 # Test repr (ignoring the L in longs)
373 r1 = repr(count(i)).replace('L', '')
374 r2 = 'count(%r)'.__mod__(i).replace('L', '')
375 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000376
Raymond Hettingera3e1ad22009-11-30 22:02:31 +0000377 # check copy, deepcopy, pickle
378 for value in -3, 3, maxsize-5, maxsize+5:
379 c = count(value)
380 self.assertEqual(next(copy.copy(c)), value)
381 self.assertEqual(next(copy.deepcopy(c)), value)
382 self.assertEqual(next(pickle.loads(pickle.dumps(c))), value)
383
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +0000384 #check proper internal error handling for large "step' sizes
385 count(1, maxsize+5); sys.exc_info()
386
Raymond Hettinger30729212009-02-12 06:28:27 +0000387 def test_count_with_stride(self):
388 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingereb13fdd2009-02-21 22:30:12 +0000389 self.assertEqual(lzip('abc',count(start=2,step=3)),
390 [('a', 2), ('b', 5), ('c', 8)])
391 self.assertEqual(lzip('abc',count(step=-1)),
392 [('a', 0), ('b', -1), ('c', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000393 self.assertEqual(lzip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
394 self.assertEqual(lzip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000395 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000396 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
397 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
398 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettinger8a882a72009-02-12 12:08:18 +0000399 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
400 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerde09acf2009-02-12 12:53:53 +0000401 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
402 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000403 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
404 c = count(3, 5)
405 self.assertEqual(repr(c), 'count(3, 5)')
406 next(c)
407 self.assertEqual(repr(c), 'count(8, 5)')
408 c = count(-9, 0)
409 self.assertEqual(repr(c), 'count(-9, 0)')
410 next(c)
411 self.assertEqual(repr(c), 'count(-9, 0)')
412 c = count(-9, -3)
413 self.assertEqual(repr(c), 'count(-9, -3)')
414 next(c)
415 self.assertEqual(repr(c), 'count(-12, -3)')
416 self.assertEqual(repr(c), 'count(-12, -3)')
417 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
418 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
419 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
420 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
421 for j in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 1, 10, sys.maxsize-5, sys.maxsize+5):
422 # Test repr (ignoring the L in longs)
423 r1 = repr(count(i, j)).replace('L', '')
424 if j == 1:
425 r2 = ('count(%r)' % i).replace('L', '')
426 else:
427 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
428 self.assertEqual(r1, r2)
429
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000430 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000431 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000432 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000433 self.assertRaises(TypeError, cycle)
434 self.assertRaises(TypeError, cycle, 5)
435 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000436
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000437 def test_groupby(self):
438 # Check whether it accepts arguments correctly
439 self.assertEqual([], list(groupby([])))
440 self.assertEqual([], list(groupby([], key=id)))
441 self.assertRaises(TypeError, list, groupby('abc', []))
442 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000443 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000444
445 # Check normal input
446 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
447 (2,15,22), (3,16,23), (3,17,23)]
448 dup = []
449 for k, g in groupby(s, lambda r:r[0]):
450 for elem in g:
451 self.assertEqual(k, elem[0])
452 dup.append(elem)
453 self.assertEqual(s, dup)
454
455 # Check nested case
456 dup = []
457 for k, g in groupby(s, lambda r:r[0]):
458 for ik, ig in groupby(g, lambda r:r[2]):
459 for elem in ig:
460 self.assertEqual(k, elem[0])
461 self.assertEqual(ik, elem[2])
462 dup.append(elem)
463 self.assertEqual(s, dup)
464
465 # Check case where inner iterator is not used
466 keys = [k for k, g in groupby(s, lambda r:r[0])]
467 expectedkeys = set([r[0] for r in s])
468 self.assertEqual(set(keys), expectedkeys)
469 self.assertEqual(len(keys), len(expectedkeys))
470
471 # Exercise pipes and filters style
472 s = 'abracadabra'
473 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000474 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000475 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
476 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000477 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000478 self.assertEqual(r, ['a', 'b', 'r'])
479 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000480 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000481 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
482 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000483 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000484 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
485
Georg Brandla18af4e2007-04-21 15:47:16 +0000486 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000487 class ExpectedError(Exception):
488 pass
489 def delayed_raise(n=0):
490 for i in range(n):
491 yield 'yo'
492 raise ExpectedError
493 def gulp(iterable, keyp=None, func=list):
494 return [func(g) for k, g in groupby(iterable, keyp)]
495
Georg Brandla18af4e2007-04-21 15:47:16 +0000496 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000497 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000498 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000499 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
500
501 # __cmp__ failure
502 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000503 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000504 raise ExpectedError
505 s = [DummyCmp(), DummyCmp(), None]
506
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000507 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000508 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000509 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000510 self.assertRaises(ExpectedError, gulp, s)
511
512 # keyfunc failure
513 def keyfunc(obj):
514 if keyfunc.skip > 0:
515 keyfunc.skip -= 1
516 return obj
517 else:
518 raise ExpectedError
519
520 # keyfunc failure on outer object
521 keyfunc.skip = 0
522 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
523 keyfunc.skip = 1
524 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
525
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000526 def test_filter(self):
527 self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
528 self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])
529 self.assertEqual(list(filter(bool, [0,1,0,2,0])), [1,2])
530 self.assertEqual(take(4, filter(isEven, count())), [0,2,4,6])
531 self.assertRaises(TypeError, filter)
532 self.assertRaises(TypeError, filter, lambda x:x)
533 self.assertRaises(TypeError, filter, lambda x:x, range(6), 7)
534 self.assertRaises(TypeError, filter, isEven, 3)
535 self.assertRaises(TypeError, next, filter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000536
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000537 def test_filterfalse(self):
538 self.assertEqual(list(filterfalse(isEven, range(6))), [1,3,5])
539 self.assertEqual(list(filterfalse(None, [0,1,0,2,0])), [0,0,0])
540 self.assertEqual(list(filterfalse(bool, [0,1,0,2,0])), [0,0,0])
541 self.assertEqual(take(4, filterfalse(isEven, count())), [1,3,5,7])
542 self.assertRaises(TypeError, filterfalse)
543 self.assertRaises(TypeError, filterfalse, lambda x:x)
544 self.assertRaises(TypeError, filterfalse, lambda x:x, range(6), 7)
545 self.assertRaises(TypeError, filterfalse, isEven, 3)
546 self.assertRaises(TypeError, next, filterfalse(range(6), range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000547
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000548 def test_zip(self):
549 # XXX This is rather silly now that builtin zip() calls zip()...
550 ans = [(x,y) for x, y in zip('abc',count())]
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000551 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000552 self.assertEqual(list(zip('abc', range(6))), lzip('abc', range(6)))
553 self.assertEqual(list(zip('abcdef', range(3))), lzip('abcdef', range(3)))
554 self.assertEqual(take(3,zip('abcdef', count())), lzip('abcdef', range(3)))
555 self.assertEqual(list(zip('abcdef')), lzip('abcdef'))
556 self.assertEqual(list(zip()), lzip())
557 self.assertRaises(TypeError, zip, 3)
558 self.assertRaises(TypeError, zip, range(3), 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000559 # Check tuple re-use (implementation detail)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000560 self.assertEqual([tuple(list(pair)) for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000561 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000562 self.assertEqual([pair for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000563 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000564 ids = list(map(id, zip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000565 self.assertEqual(min(ids), max(ids))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000566 ids = list(map(id, list(zip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000567 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000568
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000569 def test_ziplongest(self):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000570 for args in [
571 ['abc', range(6)],
572 [range(6), 'abc'],
573 [range(1000), range(2000,2100), range(3000,3050)],
574 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
575 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
576 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000577 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
578 for i in range(max(map(len, args)))]
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000579 self.assertEqual(list(zip_longest(*args)), target)
580 self.assertEqual(list(zip_longest(*args, **{})), target)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000581 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000582 self.assertEqual(list(zip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000583
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000584 self.assertEqual(take(3,zip_longest('abcdef', count())), list(zip('abcdef', range(3)))) # take 3 from infinite input
Thomas Wouterscf297e42007-02-23 15:07:44 +0000585
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000586 self.assertEqual(list(zip_longest()), list(zip()))
587 self.assertEqual(list(zip_longest([])), list(zip([])))
588 self.assertEqual(list(zip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000589
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000590 self.assertEqual(list(zip_longest('abc', 'defg', **{})),
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000591 list(zip(list('abc')+[None], 'defg'))) # empty keyword dict
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000592 self.assertRaises(TypeError, zip_longest, 3)
593 self.assertRaises(TypeError, zip_longest, range(3), 3)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000594
595 for stmt in [
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000596 "zip_longest('abc', fv=1)",
597 "zip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000598 ]:
599 try:
600 eval(stmt, globals(), locals())
601 except TypeError:
602 pass
603 else:
604 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000605
Thomas Wouterscf297e42007-02-23 15:07:44 +0000606 # Check tuple re-use (implementation detail)
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000607 self.assertEqual([tuple(list(pair)) for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000608 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000609 self.assertEqual([pair for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000610 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000611 ids = list(map(id, zip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000612 self.assertEqual(min(ids), max(ids))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000613 ids = list(map(id, list(zip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000614 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
615
Raymond Hettingerfc438512009-11-01 20:55:33 +0000616 def test_bug_7244(self):
617
618 class Repeater:
619 # this class is similar to itertools.repeat
620 def __init__(self, o, t, e):
621 self.o = o
622 self.t = int(t)
623 self.e = e
624 def __iter__(self): # its iterator is itself
625 return self
626 def __next__(self):
627 if self.t > 0:
628 self.t -= 1
629 return self.o
630 else:
631 raise self.e
632
633 # Formerly this code in would fail in debug mode
634 # with Undetected Error and Stop Iteration
635 r1 = Repeater(1, 3, StopIteration)
636 r2 = Repeater(2, 4, StopIteration)
637 def run(r1, r2):
638 result = []
639 for i, j in zip_longest(r1, r2, fillvalue=0):
640 with support.captured_output('stdout'):
641 print((i, j))
642 result.append((i, j))
643 return result
644 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
645
646 # Formerly, the RuntimeError would be lost
647 # and StopIteration would stop as expected
648 r1 = Repeater(1, 3, RuntimeError)
649 r2 = Repeater(2, 4, StopIteration)
650 it = zip_longest(r1, r2, fillvalue=0)
651 self.assertEqual(next(it), (1, 2))
652 self.assertEqual(next(it), (1, 2))
653 self.assertEqual(next(it), (1, 2))
654 self.assertRaises(RuntimeError, next, it)
655
Christian Heimesc3f30c42008-02-22 16:37:40 +0000656 def test_product(self):
657 for args, result in [
Christian Heimes78644762008-03-04 23:39:23 +0000658 ([], [()]), # zero iterables
Christian Heimesc3f30c42008-02-22 16:37:40 +0000659 (['ab'], [('a',), ('b',)]), # one iterable
660 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
661 ([range(0), range(2), range(3)], []), # first iterable with zero length
662 ([range(2), range(0), range(3)], []), # middle iterable with zero length
663 ([range(2), range(3), range(0)], []), # last iterable with zero length
664 ]:
665 self.assertEqual(list(product(*args)), result)
Christian Heimesf16baeb2008-02-29 14:57:44 +0000666 for r in range(4):
667 self.assertEqual(list(product(*(args*r))),
668 list(product(*args, **dict(repeat=r))))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000669 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
670 self.assertRaises(TypeError, product, range(6), None)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000671
672 def product1(*args, **kwds):
673 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
674 n = len(pools)
675 if n == 0:
676 yield ()
677 return
678 if any(len(pool) == 0 for pool in pools):
679 return
680 indices = [0] * n
681 yield tuple(pool[i] for pool, i in zip(pools, indices))
682 while 1:
683 for i in reversed(range(n)): # right to left
684 if indices[i] == len(pools[i]) - 1:
685 continue
686 indices[i] += 1
687 for j in range(i+1, n):
688 indices[j] = 0
689 yield tuple(pool[i] for pool, i in zip(pools, indices))
690 break
691 else:
692 return
693
694 def product2(*args, **kwds):
695 'Pure python version used in docs'
696 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
697 result = [[]]
698 for pool in pools:
699 result = [x+[y] for x in result for y in pool]
700 for prod in result:
701 yield tuple(prod)
702
Christian Heimesc3f30c42008-02-22 16:37:40 +0000703 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
704 set('abcdefg'), range(11), tuple(range(13))]
705 for i in range(100):
706 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Christian Heimes78644762008-03-04 23:39:23 +0000707 expected_len = prod(map(len, args))
708 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000709 self.assertEqual(list(product(*args)), list(product1(*args)))
710 self.assertEqual(list(product(*args)), list(product2(*args)))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000711 args = map(iter, args)
Christian Heimes78644762008-03-04 23:39:23 +0000712 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000713
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000714 # Test implementation detail: tuple re-use
715 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
716 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000717
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000718 def test_repeat(self):
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +0000719 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
Guido van Rossum805365e2007-05-07 22:24:25 +0000720 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000721 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000722 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000723 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000724 self.assertEqual(list(repeat('a', 0)), [])
725 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000726 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000727 self.assertRaises(TypeError, repeat, None, 3, 4)
728 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000729 r = repeat(1+0j)
730 self.assertEqual(repr(r), 'repeat((1+0j))')
731 r = repeat(1+0j, 5)
732 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
733 list(r)
734 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000735
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000736 def test_map(self):
737 self.assertEqual(list(map(operator.pow, range(3), range(1,7))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000738 [0**1, 1**2, 2**3])
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000739 def tupleize(*args):
740 return args
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000741 self.assertEqual(list(map(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000742 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000743 self.assertEqual(list(map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000744 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000745 self.assertEqual(take(2,map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000746 [('a',0),('b',1)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000747 self.assertEqual(list(map(operator.pow, [])), [])
748 self.assertRaises(TypeError, map)
749 self.assertRaises(TypeError, list, map(None, range(3), range(3)))
750 self.assertRaises(TypeError, map, operator.neg)
751 self.assertRaises(TypeError, next, map(10, range(5)))
752 self.assertRaises(ValueError, next, map(errfunc, [4], [5]))
753 self.assertRaises(TypeError, next, map(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000754
755 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000756 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
757 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000758 self.assertEqual(take(3, starmap(operator.pow, zip(count(), count(1)))),
Raymond Hettinger02420702003-06-29 20:36:23 +0000759 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000760 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +0000761 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
762 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000763 self.assertRaises(TypeError, starmap)
764 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000765 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
766 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
767 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000768
769 def test_islice(self):
770 for args in [ # islice(args) should agree with range(args)
771 (10, 20, 3),
772 (10, 3, 20),
773 (10, 20),
774 (10, 3),
775 (20,)
776 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000777 self.assertEqual(list(islice(range(100), *args)),
778 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000779
780 for args, tgtargs in [ # Stop when seqn is exhausted
781 ((10, 110, 3), ((10, 100, 3))),
782 ((10, 110), ((10, 100))),
783 ((110,), (100,))
784 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000785 self.assertEqual(list(islice(range(100), *args)),
786 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000787
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000788 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +0000789 self.assertEqual(list(islice(range(10), None)), list(range(10)))
790 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
791 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
792 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
793 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000794
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000795 # Test number of items consumed SF #1171417
796 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +0000797 self.assertEqual(list(islice(it, 3)), list(range(3)))
798 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000799
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000800 # Test invalid arguments
Guido van Rossum805365e2007-05-07 22:24:25 +0000801 self.assertRaises(TypeError, islice, range(10))
802 self.assertRaises(TypeError, islice, range(10), 1, 2, 3, 4)
803 self.assertRaises(ValueError, islice, range(10), -5, 10, 1)
804 self.assertRaises(ValueError, islice, range(10), 1, -5, -1)
805 self.assertRaises(ValueError, islice, range(10), 1, 10, -1)
806 self.assertRaises(ValueError, islice, range(10), 1, 10, 0)
807 self.assertRaises(ValueError, islice, range(10), 'a')
808 self.assertRaises(ValueError, islice, range(10), 'a', 1)
809 self.assertRaises(ValueError, islice, range(10), 1, 'a')
810 self.assertRaises(ValueError, islice, range(10), 'a', 1, 1)
811 self.assertRaises(ValueError, islice, range(10), 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +0000812 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000813
Raymond Hettinger69b34bf2010-11-30 02:49:29 +0000814 # Issue #10323: Less islice in a predictable state
815 c = count()
816 self.assertEqual(list(islice(c, 1, 3, 50)), [1])
817 self.assertEqual(next(c), 3)
818
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000819 def test_takewhile(self):
820 data = [1, 3, 5, 20, 2, 4, 6, 8]
821 underten = lambda x: x<10
822 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000823 self.assertEqual(list(takewhile(underten, [])), [])
824 self.assertRaises(TypeError, takewhile)
825 self.assertRaises(TypeError, takewhile, operator.pow)
826 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000827 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
828 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000829 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
830 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +0000831 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000832
833 def test_dropwhile(self):
834 data = [1, 3, 5, 20, 2, 4, 6, 8]
835 underten = lambda x: x<10
836 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000837 self.assertEqual(list(dropwhile(underten, [])), [])
838 self.assertRaises(TypeError, dropwhile)
839 self.assertRaises(TypeError, dropwhile, operator.pow)
840 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000841 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
842 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000843
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000844 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000845 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000846 def irange(n):
Guido van Rossum805365e2007-05-07 22:24:25 +0000847 for i in range(n):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000848 yield i
849
850 a, b = tee([]) # test empty iterator
851 self.assertEqual(list(a), [])
852 self.assertEqual(list(b), [])
853
854 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +0000855 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000856
857 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +0000858 self.assertEqual(list(a), list(range(n)))
859 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000860
861 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000862 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000863 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000864 del a
Guido van Rossum805365e2007-05-07 22:24:25 +0000865 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000866
867 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000868 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000869 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000870 del b
Guido van Rossum805365e2007-05-07 22:24:25 +0000871 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000872
Guido van Rossum805365e2007-05-07 22:24:25 +0000873 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000874 order = [0]*n + [1]*n
875 random.shuffle(order)
876 lists = ([], [])
877 its = tee(irange(n))
878 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +0000879 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000880 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +0000881 self.assertEqual(lists[0], list(range(n)))
882 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000883
Raymond Hettingerad983e72003-11-12 14:32:26 +0000884 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000885 self.assertRaises(TypeError, tee)
886 self.assertRaises(TypeError, tee, 3)
887 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000888 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000889
Raymond Hettingerad983e72003-11-12 14:32:26 +0000890 # tee object should be instantiable
891 a, b = tee('abc')
892 c = type(a)('def')
893 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000894
Raymond Hettingerad983e72003-11-12 14:32:26 +0000895 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +0000896 a, b, c = tee(range(2000), 3)
897 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000898 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +0000899 self.assertEqual(list(b), list(range(2000)))
900 self.assertEqual([next(c), next(c)], list(range(2)))
901 self.assertEqual(list(a), list(range(100,2000)))
902 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000903
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000904 # test values of n
905 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000906 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +0000907 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000908 result = tee('abc', n)
909 self.assertEqual(type(result), tuple)
910 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000911 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000912
Raymond Hettingerad983e72003-11-12 14:32:26 +0000913 # tee pass-through to copyable iterator
914 a, b = tee('abc')
915 c, d = tee(a)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000916 self.assertTrue(a is c)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000917
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000918 # test tee_new
919 t1, t2 = tee('abc')
920 tnew = type(t1)
921 self.assertRaises(TypeError, tnew)
922 self.assertRaises(TypeError, tnew, 10)
923 t3 = tnew(t1)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000924 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000925
Raymond Hettingera9f60922004-10-17 16:40:14 +0000926 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +0000927 a, b = tee(range(10))
Raymond Hettingera9f60922004-10-17 16:40:14 +0000928 p = proxy(a)
929 self.assertEqual(getattr(p, '__class__'), type(b))
930 del a
931 self.assertRaises(ReferenceError, getattr, p, '__class__')
932
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000933 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000934 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000935
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000936 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +0000937 self.assertRaises(StopIteration, next, f([]))
938 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000939
Georg Brandla18af4e2007-04-21 15:47:16 +0000940 self.assertRaises(StopIteration, next, islice([], None))
941 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000942
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000943 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +0000944 self.assertRaises(StopIteration, next, p)
945 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000946 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +0000947 self.assertRaises(StopIteration, next, p)
948 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000949
Georg Brandla18af4e2007-04-21 15:47:16 +0000950 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000951
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000952 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +0000953 self.assertRaises(StopIteration, next, f(lambda x:x, []))
954 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000955
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000956class TestExamples(unittest.TestCase):
957
Raymond Hettinger482ba772010-12-01 22:48:00 +0000958 def test_accumlate(self):
959 self.assertEqual(list(accumulate([1,2,3,4,5])), [1, 3, 6, 10, 15])
960
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000961 def test_chain(self):
962 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
963
964 def test_chain_from_iterable(self):
965 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
966
967 def test_combinations(self):
968 self.assertEqual(list(combinations('ABCD', 2)),
969 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
970 self.assertEqual(list(combinations(range(4), 3)),
971 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
972
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000973 def test_combinations_with_replacement(self):
974 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
975 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
976
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000977 def test_compress(self):
978 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
979
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000980 def test_count(self):
981 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
982
983 def test_cycle(self):
984 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
985
986 def test_dropwhile(self):
987 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
988
989 def test_groupby(self):
990 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
991 list('ABCDAB'))
992 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
993 [list('AAAA'), list('BBB'), list('CC'), list('D')])
994
995 def test_filter(self):
996 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
997
998 def test_filterfalse(self):
999 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
1000
1001 def test_map(self):
1002 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
1003
1004 def test_islice(self):
1005 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
1006 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
1007 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
1008 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1009
1010 def test_zip(self):
1011 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
1012
1013 def test_zip_longest(self):
1014 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
1015 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
1016
1017 def test_permutations(self):
1018 self.assertEqual(list(permutations('ABCD', 2)),
1019 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
1020 self.assertEqual(list(permutations(range(3))),
1021 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
1022
1023 def test_product(self):
1024 self.assertEqual(list(product('ABCD', 'xy')),
1025 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
1026 self.assertEqual(list(product(range(2), repeat=3)),
1027 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
1028 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
1029
1030 def test_repeat(self):
1031 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
1032
1033 def test_stapmap(self):
1034 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
1035 [32, 9, 1000])
1036
1037 def test_takewhile(self):
1038 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
1039
1040
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001041class TestGC(unittest.TestCase):
1042
1043 def makecycle(self, iterator, container):
1044 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +00001045 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001046 del container, iterator
1047
Raymond Hettinger482ba772010-12-01 22:48:00 +00001048 def test_accumulate(self):
1049 a = []
1050 self.makecycle(accumulate([1,2,a,3]), a)
1051
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001052 def test_chain(self):
1053 a = []
1054 self.makecycle(chain(a), a)
1055
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001056 def test_chain_from_iterable(self):
1057 a = []
1058 self.makecycle(chain.from_iterable([a]), a)
1059
1060 def test_combinations(self):
1061 a = []
1062 self.makecycle(combinations([1,2,a,3], 3), a)
1063
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001064 def test_combinations_with_replacement(self):
1065 a = []
1066 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1067
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001068 def test_compress(self):
1069 a = []
1070 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1071
Raymond Hettingerd6280f42009-02-16 20:50:56 +00001072 def test_count(self):
1073 a = []
1074 Int = type('Int', (int,), dict(x=a))
1075 self.makecycle(count(Int(0), Int(1)), a)
1076
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001077 def test_cycle(self):
1078 a = []
1079 self.makecycle(cycle([a]*2), a)
1080
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001081 def test_dropwhile(self):
1082 a = []
1083 self.makecycle(dropwhile(bool, [0, a, a]), a)
1084
1085 def test_groupby(self):
1086 a = []
1087 self.makecycle(groupby([a]*2, lambda x:x), a)
1088
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001089 def test_issue2246(self):
1090 # Issue 2246 -- the _grouper iterator was not included in GC
1091 n = 10
1092 keyfunc = lambda x: x
1093 for i, j in groupby(range(n), key=keyfunc):
1094 keyfunc.__dict__.setdefault('x',[]).append(j)
1095
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001096 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001097 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001098 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001099
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001100 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001101 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001102 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001103
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001104 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001105 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001106 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001107
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001108 def test_zip_longest(self):
1109 a = []
1110 self.makecycle(zip_longest([a]*2, [a]*3), a)
1111 b = [a, None]
1112 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1113
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001114 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001115 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001116 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001117
1118 def test_islice(self):
1119 a = []
1120 self.makecycle(islice([a]*2, None), a)
1121
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001122 def test_permutations(self):
1123 a = []
1124 self.makecycle(permutations([1,2,a,3], 3), a)
1125
1126 def test_product(self):
1127 a = []
1128 self.makecycle(product([1,2,a,3], repeat=3), a)
1129
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001130 def test_repeat(self):
1131 a = []
1132 self.makecycle(repeat(a), a)
1133
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001134 def test_starmap(self):
1135 a = []
1136 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1137
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001138 def test_takewhile(self):
1139 a = []
1140 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1141
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001142def R(seqn):
1143 'Regular generator'
1144 for i in seqn:
1145 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001146
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001147class G:
1148 'Sequence using __getitem__'
1149 def __init__(self, seqn):
1150 self.seqn = seqn
1151 def __getitem__(self, i):
1152 return self.seqn[i]
1153
1154class I:
1155 'Sequence using iterator protocol'
1156 def __init__(self, seqn):
1157 self.seqn = seqn
1158 self.i = 0
1159 def __iter__(self):
1160 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001161 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001162 if self.i >= len(self.seqn): raise StopIteration
1163 v = self.seqn[self.i]
1164 self.i += 1
1165 return v
1166
1167class Ig:
1168 'Sequence using iterator protocol defined with a generator'
1169 def __init__(self, seqn):
1170 self.seqn = seqn
1171 self.i = 0
1172 def __iter__(self):
1173 for val in self.seqn:
1174 yield val
1175
1176class X:
1177 'Missing __getitem__ and __iter__'
1178 def __init__(self, seqn):
1179 self.seqn = seqn
1180 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001181 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001182 if self.i >= len(self.seqn): raise StopIteration
1183 v = self.seqn[self.i]
1184 self.i += 1
1185 return v
1186
1187class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001188 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001189 def __init__(self, seqn):
1190 self.seqn = seqn
1191 self.i = 0
1192 def __iter__(self):
1193 return self
1194
1195class E:
1196 'Test propagation of exceptions'
1197 def __init__(self, seqn):
1198 self.seqn = seqn
1199 self.i = 0
1200 def __iter__(self):
1201 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001202 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001203 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001204
1205class S:
1206 'Test immediate stop'
1207 def __init__(self, seqn):
1208 pass
1209 def __iter__(self):
1210 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001211 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001212 raise StopIteration
1213
1214def L(seqn):
1215 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001216 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001217
1218
1219class TestVariousIteratorArgs(unittest.TestCase):
1220
Raymond Hettinger482ba772010-12-01 22:48:00 +00001221 def test_accumulate(self):
1222 s = [1,2,3,4,5]
1223 r = [1,3,6,10,15]
1224 n = len(s)
1225 for g in (G, I, Ig, L, R):
1226 self.assertEqual(list(accumulate(g(s))), r)
1227 self.assertEqual(list(accumulate(S(s))), [])
1228 self.assertRaises(TypeError, accumulate, X(s))
1229 self.assertRaises(TypeError, accumulate, N(s))
1230 self.assertRaises(ZeroDivisionError, list, accumulate(E(s)))
1231
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001232 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001233 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001234 for g in (G, I, Ig, S, L, R):
1235 self.assertEqual(list(chain(g(s))), list(g(s)))
1236 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001237 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001238 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001239 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1240
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001241 def test_compress(self):
1242 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1243 n = len(s)
1244 for g in (G, I, Ig, S, L, R):
1245 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1246 self.assertRaises(TypeError, compress, X(s), repeat(1))
1247 self.assertRaises(TypeError, compress, N(s), repeat(1))
1248 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1249
Christian Heimesc3f30c42008-02-22 16:37:40 +00001250 def test_product(self):
1251 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1252 self.assertRaises(TypeError, product, X(s))
1253 self.assertRaises(TypeError, product, N(s))
1254 self.assertRaises(ZeroDivisionError, product, E(s))
1255
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001256 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001257 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001258 for g in (G, I, Ig, S, L, R):
1259 tgtlen = len(s) * 3
1260 expected = list(g(s))*3
1261 actual = list(islice(cycle(g(s)), tgtlen))
1262 self.assertEqual(actual, expected)
1263 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001264 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001265 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1266
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001267 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001268 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001269 for g in (G, I, Ig, S, L, R):
1270 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1271 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001272 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001273 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1274
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001275 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001276 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001277 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001278 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001279 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001280 self.assertRaises(TypeError, filter, isEven, X(s))
1281 self.assertRaises(TypeError, filter, isEven, N(s))
1282 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001283
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001284 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001285 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001286 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001287 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001288 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001289 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1290 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1291 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001292
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001293 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001294 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001295 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001296 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1297 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1298 self.assertRaises(TypeError, zip, X(s))
1299 self.assertRaises(TypeError, zip, N(s))
1300 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001301
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001302 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001303 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001304 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001305 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1306 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1307 self.assertRaises(TypeError, zip_longest, X(s))
1308 self.assertRaises(TypeError, zip_longest, N(s))
1309 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001310
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001311 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001312 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001313 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001314 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001315 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001316 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001317 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001318 self.assertRaises(TypeError, map, onearg, X(s))
1319 self.assertRaises(TypeError, map, onearg, N(s))
1320 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001321
1322 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001323 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001324 for g in (G, I, Ig, S, L, R):
1325 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1326 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001327 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001328 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1329
1330 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001331 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001332 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001333 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001334 self.assertEqual(list(starmap(operator.pow, g(ss))),
1335 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001336 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001337 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001338 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1339
1340 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001341 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001342 for g in (G, I, Ig, S, L, R):
1343 tgt = []
1344 for elem in g(s):
1345 if not isEven(elem): break
1346 tgt.append(elem)
1347 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1348 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001349 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001350 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1351
1352 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001353 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001354 for g in (G, I, Ig, S, L, R):
1355 tgt = []
1356 for elem in g(s):
1357 if not tgt and isOdd(elem): continue
1358 tgt.append(elem)
1359 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1360 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001361 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001362 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1363
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001364 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001365 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001366 for g in (G, I, Ig, S, L, R):
1367 it1, it2 = tee(g(s))
1368 self.assertEqual(list(it1), list(g(s)))
1369 self.assertEqual(list(it2), list(g(s)))
1370 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001371 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001372 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1373
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001374class LengthTransparency(unittest.TestCase):
1375
1376 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001377 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001378 self.assertEqual(len(repeat(None, 50)), 50)
1379 self.assertRaises(TypeError, len, repeat(None))
1380
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001381class RegressionTests(unittest.TestCase):
1382
1383 def test_sf_793826(self):
1384 # Fix Armin Rigo's successful efforts to wreak havoc
1385
1386 def mutatingtuple(tuple1, f, tuple2):
1387 # this builds a tuple t which is a copy of tuple1,
1388 # then calls f(t), then mutates t to be equal to tuple2
1389 # (needs len(tuple1) == len(tuple2)).
1390 def g(value, first=[1]):
1391 if first:
1392 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00001393 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001394 return value
1395 items = list(tuple2)
1396 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001397 gen = map(g, items)
1398 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00001399 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001400
1401 def f(t):
1402 global T
1403 T = t
1404 first[:] = list(T)
1405
1406 first = []
1407 mutatingtuple((1,2,3), f, (4,5,6))
1408 second = list(T)
1409 self.assertEqual(first, second)
1410
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001411
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001412 def test_sf_950057(self):
1413 # Make sure that chain() and cycle() catch exceptions immediately
1414 # rather than when shifting between input sources
1415
1416 def gen1():
1417 hist.append(0)
1418 yield 1
1419 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001420 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001421 hist.append(2)
1422
1423 def gen2(x):
1424 hist.append(3)
1425 yield 2
1426 hist.append(4)
1427 if x:
1428 raise StopIteration
1429
1430 hist = []
1431 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1432 self.assertEqual(hist, [0,1])
1433
1434 hist = []
1435 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1436 self.assertEqual(hist, [0,1])
1437
1438 hist = []
1439 self.assertRaises(AssertionError, list, cycle(gen1()))
1440 self.assertEqual(hist, [0,1])
1441
Thomas Woutersb2137042007-02-01 18:02:27 +00001442class SubclassWithKwargsTest(unittest.TestCase):
1443 def test_keywords_in_subclass(self):
1444 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001445 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001446 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00001447 class Subclass(cls):
1448 def __init__(self, newarg=None, *args):
1449 cls.__init__(self, *args)
1450 try:
1451 Subclass(newarg=1)
1452 except TypeError as err:
1453 # we expect type errors because of wrong argument count
Ezio Melottib58e0bd2010-01-23 15:40:09 +00001454 self.assertNotIn("does not take keyword arguments", err.args[0])
Thomas Woutersb2137042007-02-01 18:02:27 +00001455
1456
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001457libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001458
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001459
1460>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001461>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00001462... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001463...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001464Check 1200 is for $120.15
1465Check 1201 is for $764.05
1466Check 1202 is for $823.14
1467
1468>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001469>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00001470... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00001471...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000014721
14738
147427
1475
1476>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001477>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00001478... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00001479...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001480Alex
1481Laura
1482Martin
1483Walter
1484Samuele
1485
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001486>>> from operator import itemgetter
1487>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00001488>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001489>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001490... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001491...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000014921 ['a', 'c', 'e']
14932 ['b', 'd', 'f']
14943 ['g']
1495
Raymond Hettinger734fb572004-01-20 20:04:40 +00001496# Find runs of consecutive numbers using groupby. The key to the solution
1497# is differencing with a range so that consecutive numbers all appear in
1498# same group.
1499>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00001500>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001501... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001502...
Raymond Hettinger734fb572004-01-20 20:04:40 +00001503[1]
1504[4, 5, 6]
1505[10]
1506[15, 16, 17, 18]
1507[22]
1508[25, 26, 27, 28]
1509
Georg Brandl3dbca812008-07-23 16:10:53 +00001510>>> def take(n, iterable):
1511... "Return first n items of the iterable as a list"
1512... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001513
Georg Brandl3dbca812008-07-23 16:10:53 +00001514>>> def enumerate(iterable, start=0):
1515... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001516
Georg Brandl3dbca812008-07-23 16:10:53 +00001517>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001518... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00001519... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001520
Raymond Hettingercdf8ba32009-02-19 04:45:07 +00001521>>> def nth(iterable, n, default=None):
1522... "Returns the nth item or a default value"
1523... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001524
Georg Brandl3dbca812008-07-23 16:10:53 +00001525>>> def quantify(iterable, pred=bool):
1526... "Count how many times the predicate is true"
1527... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001528
Georg Brandl3dbca812008-07-23 16:10:53 +00001529>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001530... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00001531... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001532
Georg Brandl3dbca812008-07-23 16:10:53 +00001533>>> def ncycles(iterable, n):
Ezio Melotti13925002011-03-16 11:05:33 +02001534... "Returns the sequence elements n times"
Georg Brandl3dbca812008-07-23 16:10:53 +00001535... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001536
1537>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001538... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001539
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001540>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001541... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001542
1543>>> def repeatfunc(func, times=None, *args):
1544... "Repeat calls to func with specified arguments."
1545... " Example: repeatfunc(random.random)"
1546... if times is None:
1547... return starmap(func, repeat(args))
1548... else:
1549... return starmap(func, repeat(args, times))
1550
Raymond Hettingerd591f662003-10-26 15:34:50 +00001551>>> def pairwise(iterable):
1552... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1553... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001554... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00001555... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001556... except StopIteration:
1557... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001558... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001559
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001560>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001561... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001562... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00001563... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001564
1565>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001566... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001567... # Recipe credited to George Sakkis
1568... pending = len(iterables)
1569... nexts = cycle(iter(it).__next__ for it in iterables)
1570... while pending:
1571... try:
1572... for next in nexts:
1573... yield next()
1574... except StopIteration:
1575... pending -= 1
1576... nexts = cycle(islice(nexts, pending))
1577
1578>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00001579... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1580... s = list(iterable)
1581... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001582
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001583>>> def unique_everseen(iterable, key=None):
1584... "List unique elements, preserving order. Remember all elements ever seen."
1585... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1586... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1587... seen = set()
1588... seen_add = seen.add
1589... if key is None:
1590... for element in iterable:
1591... if element not in seen:
1592... seen_add(element)
1593... yield element
1594... else:
1595... for element in iterable:
1596... k = key(element)
1597... if k not in seen:
1598... seen_add(k)
1599... yield element
1600
1601>>> def unique_justseen(iterable, key=None):
1602... "List unique elements, preserving order. Remember only the element just seen."
1603... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1604... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1605... return map(next, map(itemgetter(1), groupby(iterable, key)))
1606
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001607This is not part of the examples but it tests to make sure the definitions
1608perform as purported.
1609
Raymond Hettingera098b332003-09-08 23:58:40 +00001610>>> take(10, count())
1611[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1612
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001613>>> list(enumerate('abc'))
1614[(0, 'a'), (1, 'b'), (2, 'c')]
1615
1616>>> list(islice(tabulate(lambda x: 2*x), 4))
1617[0, 2, 4, 6]
1618
1619>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00001620'd'
1621
1622>>> nth('abcde', 9) is None
1623True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001624
Guido van Rossum805365e2007-05-07 22:24:25 +00001625>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000162650
1627
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001628>>> a = [[1, 2, 3], [4, 5, 6]]
1629>>> flatten(a)
1630[1, 2, 3, 4, 5, 6]
1631
1632>>> list(repeatfunc(pow, 5, 2, 3))
1633[8, 8, 8, 8, 8]
1634
1635>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001636>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001637[0, 0, 0, 0, 0]
1638
Raymond Hettingerd591f662003-10-26 15:34:50 +00001639>>> list(pairwise('abcd'))
1640[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001641
Raymond Hettingerd591f662003-10-26 15:34:50 +00001642>>> list(pairwise([]))
1643[]
1644
1645>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001646[]
1647
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001648>>> list(islice(padnone('abc'), 0, 6))
1649['a', 'b', 'c', None, None, None]
1650
1651>>> list(ncycles('abc', 3))
1652['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1653
1654>>> dotproduct([1,2,3], [4,5,6])
165532
1656
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001657>>> list(grouper(3, 'abcdefg', 'x'))
1658[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1659
1660>>> list(roundrobin('abc', 'd', 'ef'))
1661['a', 'd', 'e', 'b', 'f', 'c']
1662
Raymond Hettingerace67332009-01-26 02:23:50 +00001663>>> list(powerset([1,2,3]))
1664[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001665
Raymond Hettinger191e8502009-01-27 13:29:43 +00001666>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1667True
1668
1669>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1670True
1671
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001672>>> list(unique_everseen('AAAABBBCCDAABBB'))
1673['A', 'B', 'C', 'D']
1674
1675>>> list(unique_everseen('ABBCcAD', str.lower))
1676['A', 'B', 'C', 'D']
1677
1678>>> list(unique_justseen('AAAABBBCCDAABBB'))
1679['A', 'B', 'C', 'D', 'A', 'B']
1680
1681>>> list(unique_justseen('ABBCcAD', str.lower))
1682['A', 'B', 'C', 'A', 'D']
1683
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001684"""
1685
1686__test__ = {'libreftest' : libreftest}
1687
1688def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001689 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00001690 RegressionTests, LengthTransparency,
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001691 SubclassWithKwargsTest, TestExamples)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001692 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001693
1694 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001695 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001696 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001697 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00001698 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001699 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001700 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001701 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001702 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001703
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001704 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001705 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001706
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001707if __name__ == "__main__":
1708 test_main(verbose=True)