blob: b744636551ff52a7bd6d273ad5cd973b8c6ed0e0 [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)))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000360 c = count(3)
361 self.assertEqual(repr(c), 'count(3)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000362 next(c)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000363 self.assertEqual(repr(c), 'count(4)')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000364 c = count(-9)
365 self.assertEqual(repr(c), 'count(-9)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000366 next(c)
Raymond Hettinger30729212009-02-12 06:28:27 +0000367 self.assertEqual(repr(count(10.25)), 'count(10.25)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000368 self.assertEqual(next(c), -8)
Christian Heimesa37d4c62007-12-04 23:02:19 +0000369 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 +0000370 # Test repr (ignoring the L in longs)
371 r1 = repr(count(i)).replace('L', '')
372 r2 = 'count(%r)'.__mod__(i).replace('L', '')
373 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000374
Raymond Hettingera3e1ad22009-11-30 22:02:31 +0000375 # check copy, deepcopy, pickle
376 for value in -3, 3, maxsize-5, maxsize+5:
377 c = count(value)
378 self.assertEqual(next(copy.copy(c)), value)
379 self.assertEqual(next(copy.deepcopy(c)), value)
380 self.assertEqual(next(pickle.loads(pickle.dumps(c))), value)
381
Raymond Hettinger30729212009-02-12 06:28:27 +0000382 def test_count_with_stride(self):
383 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingereb13fdd2009-02-21 22:30:12 +0000384 self.assertEqual(lzip('abc',count(start=2,step=3)),
385 [('a', 2), ('b', 5), ('c', 8)])
386 self.assertEqual(lzip('abc',count(step=-1)),
387 [('a', 0), ('b', -1), ('c', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000388 self.assertEqual(lzip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
389 self.assertEqual(lzip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000390 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000391 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
392 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
393 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettinger8a882a72009-02-12 12:08:18 +0000394 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
395 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerde09acf2009-02-12 12:53:53 +0000396 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
397 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000398 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
399 c = count(3, 5)
400 self.assertEqual(repr(c), 'count(3, 5)')
401 next(c)
402 self.assertEqual(repr(c), 'count(8, 5)')
403 c = count(-9, 0)
404 self.assertEqual(repr(c), 'count(-9, 0)')
405 next(c)
406 self.assertEqual(repr(c), 'count(-9, 0)')
407 c = count(-9, -3)
408 self.assertEqual(repr(c), 'count(-9, -3)')
409 next(c)
410 self.assertEqual(repr(c), 'count(-12, -3)')
411 self.assertEqual(repr(c), 'count(-12, -3)')
412 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
413 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
414 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
415 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
416 for j in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 1, 10, sys.maxsize-5, sys.maxsize+5):
417 # Test repr (ignoring the L in longs)
418 r1 = repr(count(i, j)).replace('L', '')
419 if j == 1:
420 r2 = ('count(%r)' % i).replace('L', '')
421 else:
422 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
423 self.assertEqual(r1, r2)
424
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000425 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000426 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000427 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000428 self.assertRaises(TypeError, cycle)
429 self.assertRaises(TypeError, cycle, 5)
430 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000431
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000432 def test_groupby(self):
433 # Check whether it accepts arguments correctly
434 self.assertEqual([], list(groupby([])))
435 self.assertEqual([], list(groupby([], key=id)))
436 self.assertRaises(TypeError, list, groupby('abc', []))
437 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000438 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000439
440 # Check normal input
441 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
442 (2,15,22), (3,16,23), (3,17,23)]
443 dup = []
444 for k, g in groupby(s, lambda r:r[0]):
445 for elem in g:
446 self.assertEqual(k, elem[0])
447 dup.append(elem)
448 self.assertEqual(s, dup)
449
450 # Check nested case
451 dup = []
452 for k, g in groupby(s, lambda r:r[0]):
453 for ik, ig in groupby(g, lambda r:r[2]):
454 for elem in ig:
455 self.assertEqual(k, elem[0])
456 self.assertEqual(ik, elem[2])
457 dup.append(elem)
458 self.assertEqual(s, dup)
459
460 # Check case where inner iterator is not used
461 keys = [k for k, g in groupby(s, lambda r:r[0])]
462 expectedkeys = set([r[0] for r in s])
463 self.assertEqual(set(keys), expectedkeys)
464 self.assertEqual(len(keys), len(expectedkeys))
465
466 # Exercise pipes and filters style
467 s = 'abracadabra'
468 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000469 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000470 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
471 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000472 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000473 self.assertEqual(r, ['a', 'b', 'r'])
474 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000475 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000476 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
477 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000478 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000479 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
480
Georg Brandla18af4e2007-04-21 15:47:16 +0000481 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000482 class ExpectedError(Exception):
483 pass
484 def delayed_raise(n=0):
485 for i in range(n):
486 yield 'yo'
487 raise ExpectedError
488 def gulp(iterable, keyp=None, func=list):
489 return [func(g) for k, g in groupby(iterable, keyp)]
490
Georg Brandla18af4e2007-04-21 15:47:16 +0000491 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000492 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000493 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000494 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
495
496 # __cmp__ failure
497 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000498 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000499 raise ExpectedError
500 s = [DummyCmp(), DummyCmp(), None]
501
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000502 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000503 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000504 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000505 self.assertRaises(ExpectedError, gulp, s)
506
507 # keyfunc failure
508 def keyfunc(obj):
509 if keyfunc.skip > 0:
510 keyfunc.skip -= 1
511 return obj
512 else:
513 raise ExpectedError
514
515 # keyfunc failure on outer object
516 keyfunc.skip = 0
517 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
518 keyfunc.skip = 1
519 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
520
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000521 def test_filter(self):
522 self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
523 self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])
524 self.assertEqual(list(filter(bool, [0,1,0,2,0])), [1,2])
525 self.assertEqual(take(4, filter(isEven, count())), [0,2,4,6])
526 self.assertRaises(TypeError, filter)
527 self.assertRaises(TypeError, filter, lambda x:x)
528 self.assertRaises(TypeError, filter, lambda x:x, range(6), 7)
529 self.assertRaises(TypeError, filter, isEven, 3)
530 self.assertRaises(TypeError, next, filter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000531
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000532 def test_filterfalse(self):
533 self.assertEqual(list(filterfalse(isEven, range(6))), [1,3,5])
534 self.assertEqual(list(filterfalse(None, [0,1,0,2,0])), [0,0,0])
535 self.assertEqual(list(filterfalse(bool, [0,1,0,2,0])), [0,0,0])
536 self.assertEqual(take(4, filterfalse(isEven, count())), [1,3,5,7])
537 self.assertRaises(TypeError, filterfalse)
538 self.assertRaises(TypeError, filterfalse, lambda x:x)
539 self.assertRaises(TypeError, filterfalse, lambda x:x, range(6), 7)
540 self.assertRaises(TypeError, filterfalse, isEven, 3)
541 self.assertRaises(TypeError, next, filterfalse(range(6), range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000542
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000543 def test_zip(self):
544 # XXX This is rather silly now that builtin zip() calls zip()...
545 ans = [(x,y) for x, y in zip('abc',count())]
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000546 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000547 self.assertEqual(list(zip('abc', range(6))), lzip('abc', range(6)))
548 self.assertEqual(list(zip('abcdef', range(3))), lzip('abcdef', range(3)))
549 self.assertEqual(take(3,zip('abcdef', count())), lzip('abcdef', range(3)))
550 self.assertEqual(list(zip('abcdef')), lzip('abcdef'))
551 self.assertEqual(list(zip()), lzip())
552 self.assertRaises(TypeError, zip, 3)
553 self.assertRaises(TypeError, zip, range(3), 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000554 # Check tuple re-use (implementation detail)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000555 self.assertEqual([tuple(list(pair)) for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000556 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000557 self.assertEqual([pair for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000558 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000559 ids = list(map(id, zip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000560 self.assertEqual(min(ids), max(ids))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000561 ids = list(map(id, list(zip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000562 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000563
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000564 def test_ziplongest(self):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000565 for args in [
566 ['abc', range(6)],
567 [range(6), 'abc'],
568 [range(1000), range(2000,2100), range(3000,3050)],
569 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
570 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
571 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000572 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
573 for i in range(max(map(len, args)))]
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000574 self.assertEqual(list(zip_longest(*args)), target)
575 self.assertEqual(list(zip_longest(*args, **{})), target)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000576 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 +0000577 self.assertEqual(list(zip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000578
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000579 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 +0000580
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000581 self.assertEqual(list(zip_longest()), list(zip()))
582 self.assertEqual(list(zip_longest([])), list(zip([])))
583 self.assertEqual(list(zip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000584
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000585 self.assertEqual(list(zip_longest('abc', 'defg', **{})),
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000586 list(zip(list('abc')+[None], 'defg'))) # empty keyword dict
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000587 self.assertRaises(TypeError, zip_longest, 3)
588 self.assertRaises(TypeError, zip_longest, range(3), 3)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000589
590 for stmt in [
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000591 "zip_longest('abc', fv=1)",
592 "zip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000593 ]:
594 try:
595 eval(stmt, globals(), locals())
596 except TypeError:
597 pass
598 else:
599 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000600
Thomas Wouterscf297e42007-02-23 15:07:44 +0000601 # Check tuple re-use (implementation detail)
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000602 self.assertEqual([tuple(list(pair)) for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000603 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000604 self.assertEqual([pair for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000605 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000606 ids = list(map(id, zip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000607 self.assertEqual(min(ids), max(ids))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000608 ids = list(map(id, list(zip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000609 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
610
Raymond Hettingerfc438512009-11-01 20:55:33 +0000611 def test_bug_7244(self):
612
613 class Repeater:
614 # this class is similar to itertools.repeat
615 def __init__(self, o, t, e):
616 self.o = o
617 self.t = int(t)
618 self.e = e
619 def __iter__(self): # its iterator is itself
620 return self
621 def __next__(self):
622 if self.t > 0:
623 self.t -= 1
624 return self.o
625 else:
626 raise self.e
627
628 # Formerly this code in would fail in debug mode
629 # with Undetected Error and Stop Iteration
630 r1 = Repeater(1, 3, StopIteration)
631 r2 = Repeater(2, 4, StopIteration)
632 def run(r1, r2):
633 result = []
634 for i, j in zip_longest(r1, r2, fillvalue=0):
635 with support.captured_output('stdout'):
636 print((i, j))
637 result.append((i, j))
638 return result
639 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
640
641 # Formerly, the RuntimeError would be lost
642 # and StopIteration would stop as expected
643 r1 = Repeater(1, 3, RuntimeError)
644 r2 = Repeater(2, 4, StopIteration)
645 it = zip_longest(r1, r2, fillvalue=0)
646 self.assertEqual(next(it), (1, 2))
647 self.assertEqual(next(it), (1, 2))
648 self.assertEqual(next(it), (1, 2))
649 self.assertRaises(RuntimeError, next, it)
650
Christian Heimesc3f30c42008-02-22 16:37:40 +0000651 def test_product(self):
652 for args, result in [
Christian Heimes78644762008-03-04 23:39:23 +0000653 ([], [()]), # zero iterables
Christian Heimesc3f30c42008-02-22 16:37:40 +0000654 (['ab'], [('a',), ('b',)]), # one iterable
655 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
656 ([range(0), range(2), range(3)], []), # first iterable with zero length
657 ([range(2), range(0), range(3)], []), # middle iterable with zero length
658 ([range(2), range(3), range(0)], []), # last iterable with zero length
659 ]:
660 self.assertEqual(list(product(*args)), result)
Christian Heimesf16baeb2008-02-29 14:57:44 +0000661 for r in range(4):
662 self.assertEqual(list(product(*(args*r))),
663 list(product(*args, **dict(repeat=r))))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000664 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
665 self.assertRaises(TypeError, product, range(6), None)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000666
667 def product1(*args, **kwds):
668 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
669 n = len(pools)
670 if n == 0:
671 yield ()
672 return
673 if any(len(pool) == 0 for pool in pools):
674 return
675 indices = [0] * n
676 yield tuple(pool[i] for pool, i in zip(pools, indices))
677 while 1:
678 for i in reversed(range(n)): # right to left
679 if indices[i] == len(pools[i]) - 1:
680 continue
681 indices[i] += 1
682 for j in range(i+1, n):
683 indices[j] = 0
684 yield tuple(pool[i] for pool, i in zip(pools, indices))
685 break
686 else:
687 return
688
689 def product2(*args, **kwds):
690 'Pure python version used in docs'
691 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
692 result = [[]]
693 for pool in pools:
694 result = [x+[y] for x in result for y in pool]
695 for prod in result:
696 yield tuple(prod)
697
Christian Heimesc3f30c42008-02-22 16:37:40 +0000698 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
699 set('abcdefg'), range(11), tuple(range(13))]
700 for i in range(100):
701 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Christian Heimes78644762008-03-04 23:39:23 +0000702 expected_len = prod(map(len, args))
703 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000704 self.assertEqual(list(product(*args)), list(product1(*args)))
705 self.assertEqual(list(product(*args)), list(product2(*args)))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000706 args = map(iter, args)
Christian Heimes78644762008-03-04 23:39:23 +0000707 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000708
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000709 # Test implementation detail: tuple re-use
710 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
711 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000712
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000713 def test_repeat(self):
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +0000714 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
Guido van Rossum805365e2007-05-07 22:24:25 +0000715 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000716 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000717 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000718 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000719 self.assertEqual(list(repeat('a', 0)), [])
720 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000721 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000722 self.assertRaises(TypeError, repeat, None, 3, 4)
723 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000724 r = repeat(1+0j)
725 self.assertEqual(repr(r), 'repeat((1+0j))')
726 r = repeat(1+0j, 5)
727 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
728 list(r)
729 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000730
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000731 def test_map(self):
732 self.assertEqual(list(map(operator.pow, range(3), range(1,7))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000733 [0**1, 1**2, 2**3])
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000734 def tupleize(*args):
735 return args
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000736 self.assertEqual(list(map(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000737 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000738 self.assertEqual(list(map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000739 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000740 self.assertEqual(take(2,map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000741 [('a',0),('b',1)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000742 self.assertEqual(list(map(operator.pow, [])), [])
743 self.assertRaises(TypeError, map)
744 self.assertRaises(TypeError, list, map(None, range(3), range(3)))
745 self.assertRaises(TypeError, map, operator.neg)
746 self.assertRaises(TypeError, next, map(10, range(5)))
747 self.assertRaises(ValueError, next, map(errfunc, [4], [5]))
748 self.assertRaises(TypeError, next, map(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000749
750 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000751 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
752 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000753 self.assertEqual(take(3, starmap(operator.pow, zip(count(), count(1)))),
Raymond Hettinger02420702003-06-29 20:36:23 +0000754 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000755 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +0000756 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
757 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000758 self.assertRaises(TypeError, starmap)
759 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000760 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
761 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
762 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000763
764 def test_islice(self):
765 for args in [ # islice(args) should agree with range(args)
766 (10, 20, 3),
767 (10, 3, 20),
768 (10, 20),
769 (10, 3),
770 (20,)
771 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000772 self.assertEqual(list(islice(range(100), *args)),
773 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000774
775 for args, tgtargs in [ # Stop when seqn is exhausted
776 ((10, 110, 3), ((10, 100, 3))),
777 ((10, 110), ((10, 100))),
778 ((110,), (100,))
779 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000780 self.assertEqual(list(islice(range(100), *args)),
781 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000782
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000783 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +0000784 self.assertEqual(list(islice(range(10), None)), list(range(10)))
785 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
786 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
787 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
788 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000789
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000790 # Test number of items consumed SF #1171417
791 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +0000792 self.assertEqual(list(islice(it, 3)), list(range(3)))
793 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000794
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000795 # Test invalid arguments
Guido van Rossum805365e2007-05-07 22:24:25 +0000796 self.assertRaises(TypeError, islice, range(10))
797 self.assertRaises(TypeError, islice, range(10), 1, 2, 3, 4)
798 self.assertRaises(ValueError, islice, range(10), -5, 10, 1)
799 self.assertRaises(ValueError, islice, range(10), 1, -5, -1)
800 self.assertRaises(ValueError, islice, range(10), 1, 10, -1)
801 self.assertRaises(ValueError, islice, range(10), 1, 10, 0)
802 self.assertRaises(ValueError, islice, range(10), 'a')
803 self.assertRaises(ValueError, islice, range(10), 'a', 1)
804 self.assertRaises(ValueError, islice, range(10), 1, 'a')
805 self.assertRaises(ValueError, islice, range(10), 'a', 1, 1)
806 self.assertRaises(ValueError, islice, range(10), 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +0000807 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000808
Raymond Hettinger69b34bf2010-11-30 02:49:29 +0000809 # Issue #10323: Less islice in a predictable state
810 c = count()
811 self.assertEqual(list(islice(c, 1, 3, 50)), [1])
812 self.assertEqual(next(c), 3)
813
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000814 def test_takewhile(self):
815 data = [1, 3, 5, 20, 2, 4, 6, 8]
816 underten = lambda x: x<10
817 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000818 self.assertEqual(list(takewhile(underten, [])), [])
819 self.assertRaises(TypeError, takewhile)
820 self.assertRaises(TypeError, takewhile, operator.pow)
821 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000822 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
823 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000824 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
825 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +0000826 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000827
828 def test_dropwhile(self):
829 data = [1, 3, 5, 20, 2, 4, 6, 8]
830 underten = lambda x: x<10
831 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000832 self.assertEqual(list(dropwhile(underten, [])), [])
833 self.assertRaises(TypeError, dropwhile)
834 self.assertRaises(TypeError, dropwhile, operator.pow)
835 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000836 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
837 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000838
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000839 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000840 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000841 def irange(n):
Guido van Rossum805365e2007-05-07 22:24:25 +0000842 for i in range(n):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000843 yield i
844
845 a, b = tee([]) # test empty iterator
846 self.assertEqual(list(a), [])
847 self.assertEqual(list(b), [])
848
849 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +0000850 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000851
852 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +0000853 self.assertEqual(list(a), list(range(n)))
854 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000855
856 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000857 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000858 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000859 del a
Guido van Rossum805365e2007-05-07 22:24:25 +0000860 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000861
862 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000863 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000864 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000865 del b
Guido van Rossum805365e2007-05-07 22:24:25 +0000866 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000867
Guido van Rossum805365e2007-05-07 22:24:25 +0000868 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000869 order = [0]*n + [1]*n
870 random.shuffle(order)
871 lists = ([], [])
872 its = tee(irange(n))
873 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +0000874 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000875 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +0000876 self.assertEqual(lists[0], list(range(n)))
877 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000878
Raymond Hettingerad983e72003-11-12 14:32:26 +0000879 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000880 self.assertRaises(TypeError, tee)
881 self.assertRaises(TypeError, tee, 3)
882 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000883 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000884
Raymond Hettingerad983e72003-11-12 14:32:26 +0000885 # tee object should be instantiable
886 a, b = tee('abc')
887 c = type(a)('def')
888 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000889
Raymond Hettingerad983e72003-11-12 14:32:26 +0000890 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +0000891 a, b, c = tee(range(2000), 3)
892 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000893 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +0000894 self.assertEqual(list(b), list(range(2000)))
895 self.assertEqual([next(c), next(c)], list(range(2)))
896 self.assertEqual(list(a), list(range(100,2000)))
897 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000898
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000899 # test values of n
900 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000901 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +0000902 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000903 result = tee('abc', n)
904 self.assertEqual(type(result), tuple)
905 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000906 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000907
Raymond Hettingerad983e72003-11-12 14:32:26 +0000908 # tee pass-through to copyable iterator
909 a, b = tee('abc')
910 c, d = tee(a)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000911 self.assertTrue(a is c)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000912
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000913 # test tee_new
914 t1, t2 = tee('abc')
915 tnew = type(t1)
916 self.assertRaises(TypeError, tnew)
917 self.assertRaises(TypeError, tnew, 10)
918 t3 = tnew(t1)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000919 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000920
Raymond Hettingera9f60922004-10-17 16:40:14 +0000921 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +0000922 a, b = tee(range(10))
Raymond Hettingera9f60922004-10-17 16:40:14 +0000923 p = proxy(a)
924 self.assertEqual(getattr(p, '__class__'), type(b))
925 del a
926 self.assertRaises(ReferenceError, getattr, p, '__class__')
927
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000928 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000929 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000930
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000931 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +0000932 self.assertRaises(StopIteration, next, f([]))
933 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000934
Georg Brandla18af4e2007-04-21 15:47:16 +0000935 self.assertRaises(StopIteration, next, islice([], None))
936 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000937
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000938 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +0000939 self.assertRaises(StopIteration, next, p)
940 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000941 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +0000942 self.assertRaises(StopIteration, next, p)
943 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000944
Georg Brandla18af4e2007-04-21 15:47:16 +0000945 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000946
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000947 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +0000948 self.assertRaises(StopIteration, next, f(lambda x:x, []))
949 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000950
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000951class TestExamples(unittest.TestCase):
952
Raymond Hettinger482ba772010-12-01 22:48:00 +0000953 def test_accumlate(self):
954 self.assertEqual(list(accumulate([1,2,3,4,5])), [1, 3, 6, 10, 15])
955
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000956 def test_chain(self):
957 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
958
959 def test_chain_from_iterable(self):
960 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
961
962 def test_combinations(self):
963 self.assertEqual(list(combinations('ABCD', 2)),
964 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
965 self.assertEqual(list(combinations(range(4), 3)),
966 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
967
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000968 def test_combinations_with_replacement(self):
969 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
970 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
971
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000972 def test_compress(self):
973 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
974
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000975 def test_count(self):
976 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
977
978 def test_cycle(self):
979 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
980
981 def test_dropwhile(self):
982 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
983
984 def test_groupby(self):
985 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
986 list('ABCDAB'))
987 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
988 [list('AAAA'), list('BBB'), list('CC'), list('D')])
989
990 def test_filter(self):
991 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
992
993 def test_filterfalse(self):
994 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
995
996 def test_map(self):
997 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
998
999 def test_islice(self):
1000 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
1001 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
1002 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
1003 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1004
1005 def test_zip(self):
1006 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
1007
1008 def test_zip_longest(self):
1009 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
1010 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
1011
1012 def test_permutations(self):
1013 self.assertEqual(list(permutations('ABCD', 2)),
1014 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
1015 self.assertEqual(list(permutations(range(3))),
1016 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
1017
1018 def test_product(self):
1019 self.assertEqual(list(product('ABCD', 'xy')),
1020 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
1021 self.assertEqual(list(product(range(2), repeat=3)),
1022 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
1023 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
1024
1025 def test_repeat(self):
1026 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
1027
1028 def test_stapmap(self):
1029 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
1030 [32, 9, 1000])
1031
1032 def test_takewhile(self):
1033 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
1034
1035
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001036class TestGC(unittest.TestCase):
1037
1038 def makecycle(self, iterator, container):
1039 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +00001040 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001041 del container, iterator
1042
Raymond Hettinger482ba772010-12-01 22:48:00 +00001043 def test_accumulate(self):
1044 a = []
1045 self.makecycle(accumulate([1,2,a,3]), a)
1046
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001047 def test_chain(self):
1048 a = []
1049 self.makecycle(chain(a), a)
1050
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001051 def test_chain_from_iterable(self):
1052 a = []
1053 self.makecycle(chain.from_iterable([a]), a)
1054
1055 def test_combinations(self):
1056 a = []
1057 self.makecycle(combinations([1,2,a,3], 3), a)
1058
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001059 def test_combinations_with_replacement(self):
1060 a = []
1061 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1062
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001063 def test_compress(self):
1064 a = []
1065 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1066
Raymond Hettingerd6280f42009-02-16 20:50:56 +00001067 def test_count(self):
1068 a = []
1069 Int = type('Int', (int,), dict(x=a))
1070 self.makecycle(count(Int(0), Int(1)), a)
1071
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001072 def test_cycle(self):
1073 a = []
1074 self.makecycle(cycle([a]*2), a)
1075
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001076 def test_dropwhile(self):
1077 a = []
1078 self.makecycle(dropwhile(bool, [0, a, a]), a)
1079
1080 def test_groupby(self):
1081 a = []
1082 self.makecycle(groupby([a]*2, lambda x:x), a)
1083
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001084 def test_issue2246(self):
1085 # Issue 2246 -- the _grouper iterator was not included in GC
1086 n = 10
1087 keyfunc = lambda x: x
1088 for i, j in groupby(range(n), key=keyfunc):
1089 keyfunc.__dict__.setdefault('x',[]).append(j)
1090
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001091 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001092 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001093 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001094
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001095 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001096 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001097 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001098
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001099 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001100 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001101 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001102
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001103 def test_zip_longest(self):
1104 a = []
1105 self.makecycle(zip_longest([a]*2, [a]*3), a)
1106 b = [a, None]
1107 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1108
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001109 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001110 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001111 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001112
1113 def test_islice(self):
1114 a = []
1115 self.makecycle(islice([a]*2, None), a)
1116
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001117 def test_permutations(self):
1118 a = []
1119 self.makecycle(permutations([1,2,a,3], 3), a)
1120
1121 def test_product(self):
1122 a = []
1123 self.makecycle(product([1,2,a,3], repeat=3), a)
1124
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001125 def test_repeat(self):
1126 a = []
1127 self.makecycle(repeat(a), a)
1128
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001129 def test_starmap(self):
1130 a = []
1131 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1132
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001133 def test_takewhile(self):
1134 a = []
1135 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1136
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001137def R(seqn):
1138 'Regular generator'
1139 for i in seqn:
1140 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001141
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001142class G:
1143 'Sequence using __getitem__'
1144 def __init__(self, seqn):
1145 self.seqn = seqn
1146 def __getitem__(self, i):
1147 return self.seqn[i]
1148
1149class I:
1150 'Sequence using iterator protocol'
1151 def __init__(self, seqn):
1152 self.seqn = seqn
1153 self.i = 0
1154 def __iter__(self):
1155 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001156 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001157 if self.i >= len(self.seqn): raise StopIteration
1158 v = self.seqn[self.i]
1159 self.i += 1
1160 return v
1161
1162class Ig:
1163 'Sequence using iterator protocol defined with a generator'
1164 def __init__(self, seqn):
1165 self.seqn = seqn
1166 self.i = 0
1167 def __iter__(self):
1168 for val in self.seqn:
1169 yield val
1170
1171class X:
1172 'Missing __getitem__ and __iter__'
1173 def __init__(self, seqn):
1174 self.seqn = seqn
1175 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001176 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001177 if self.i >= len(self.seqn): raise StopIteration
1178 v = self.seqn[self.i]
1179 self.i += 1
1180 return v
1181
1182class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001183 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001184 def __init__(self, seqn):
1185 self.seqn = seqn
1186 self.i = 0
1187 def __iter__(self):
1188 return self
1189
1190class E:
1191 'Test propagation of exceptions'
1192 def __init__(self, seqn):
1193 self.seqn = seqn
1194 self.i = 0
1195 def __iter__(self):
1196 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001197 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001198 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001199
1200class S:
1201 'Test immediate stop'
1202 def __init__(self, seqn):
1203 pass
1204 def __iter__(self):
1205 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001206 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001207 raise StopIteration
1208
1209def L(seqn):
1210 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001211 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001212
1213
1214class TestVariousIteratorArgs(unittest.TestCase):
1215
Raymond Hettinger482ba772010-12-01 22:48:00 +00001216 def test_accumulate(self):
1217 s = [1,2,3,4,5]
1218 r = [1,3,6,10,15]
1219 n = len(s)
1220 for g in (G, I, Ig, L, R):
1221 self.assertEqual(list(accumulate(g(s))), r)
1222 self.assertEqual(list(accumulate(S(s))), [])
1223 self.assertRaises(TypeError, accumulate, X(s))
1224 self.assertRaises(TypeError, accumulate, N(s))
1225 self.assertRaises(ZeroDivisionError, list, accumulate(E(s)))
1226
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001227 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001228 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001229 for g in (G, I, Ig, S, L, R):
1230 self.assertEqual(list(chain(g(s))), list(g(s)))
1231 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001232 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001233 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001234 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1235
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001236 def test_compress(self):
1237 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1238 n = len(s)
1239 for g in (G, I, Ig, S, L, R):
1240 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1241 self.assertRaises(TypeError, compress, X(s), repeat(1))
1242 self.assertRaises(TypeError, compress, N(s), repeat(1))
1243 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1244
Christian Heimesc3f30c42008-02-22 16:37:40 +00001245 def test_product(self):
1246 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1247 self.assertRaises(TypeError, product, X(s))
1248 self.assertRaises(TypeError, product, N(s))
1249 self.assertRaises(ZeroDivisionError, product, E(s))
1250
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001251 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001252 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001253 for g in (G, I, Ig, S, L, R):
1254 tgtlen = len(s) * 3
1255 expected = list(g(s))*3
1256 actual = list(islice(cycle(g(s)), tgtlen))
1257 self.assertEqual(actual, expected)
1258 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001259 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001260 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1261
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001262 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001263 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001264 for g in (G, I, Ig, S, L, R):
1265 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1266 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001267 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001268 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1269
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001270 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001271 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001272 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001273 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001274 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001275 self.assertRaises(TypeError, filter, isEven, X(s))
1276 self.assertRaises(TypeError, filter, isEven, N(s))
1277 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001278
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001279 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001280 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001281 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001282 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001283 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001284 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1285 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1286 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001287
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001288 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001289 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001290 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001291 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1292 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1293 self.assertRaises(TypeError, zip, X(s))
1294 self.assertRaises(TypeError, zip, N(s))
1295 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001296
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001297 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001298 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001299 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001300 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1301 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1302 self.assertRaises(TypeError, zip_longest, X(s))
1303 self.assertRaises(TypeError, zip_longest, N(s))
1304 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001305
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001306 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001307 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001308 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001309 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001310 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001311 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001312 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001313 self.assertRaises(TypeError, map, onearg, X(s))
1314 self.assertRaises(TypeError, map, onearg, N(s))
1315 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001316
1317 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001318 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001319 for g in (G, I, Ig, S, L, R):
1320 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1321 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001322 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001323 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1324
1325 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001326 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001327 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001328 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001329 self.assertEqual(list(starmap(operator.pow, g(ss))),
1330 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001331 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001332 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001333 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1334
1335 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001336 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001337 for g in (G, I, Ig, S, L, R):
1338 tgt = []
1339 for elem in g(s):
1340 if not isEven(elem): break
1341 tgt.append(elem)
1342 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1343 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001344 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001345 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1346
1347 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001348 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001349 for g in (G, I, Ig, S, L, R):
1350 tgt = []
1351 for elem in g(s):
1352 if not tgt and isOdd(elem): continue
1353 tgt.append(elem)
1354 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1355 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001356 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001357 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1358
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001359 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001360 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001361 for g in (G, I, Ig, S, L, R):
1362 it1, it2 = tee(g(s))
1363 self.assertEqual(list(it1), list(g(s)))
1364 self.assertEqual(list(it2), list(g(s)))
1365 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001366 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001367 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1368
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001369class LengthTransparency(unittest.TestCase):
1370
1371 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001372 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001373 self.assertEqual(len(repeat(None, 50)), 50)
1374 self.assertRaises(TypeError, len, repeat(None))
1375
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001376class RegressionTests(unittest.TestCase):
1377
1378 def test_sf_793826(self):
1379 # Fix Armin Rigo's successful efforts to wreak havoc
1380
1381 def mutatingtuple(tuple1, f, tuple2):
1382 # this builds a tuple t which is a copy of tuple1,
1383 # then calls f(t), then mutates t to be equal to tuple2
1384 # (needs len(tuple1) == len(tuple2)).
1385 def g(value, first=[1]):
1386 if first:
1387 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00001388 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001389 return value
1390 items = list(tuple2)
1391 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001392 gen = map(g, items)
1393 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00001394 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001395
1396 def f(t):
1397 global T
1398 T = t
1399 first[:] = list(T)
1400
1401 first = []
1402 mutatingtuple((1,2,3), f, (4,5,6))
1403 second = list(T)
1404 self.assertEqual(first, second)
1405
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001406
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001407 def test_sf_950057(self):
1408 # Make sure that chain() and cycle() catch exceptions immediately
1409 # rather than when shifting between input sources
1410
1411 def gen1():
1412 hist.append(0)
1413 yield 1
1414 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001415 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001416 hist.append(2)
1417
1418 def gen2(x):
1419 hist.append(3)
1420 yield 2
1421 hist.append(4)
1422 if x:
1423 raise StopIteration
1424
1425 hist = []
1426 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1427 self.assertEqual(hist, [0,1])
1428
1429 hist = []
1430 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1431 self.assertEqual(hist, [0,1])
1432
1433 hist = []
1434 self.assertRaises(AssertionError, list, cycle(gen1()))
1435 self.assertEqual(hist, [0,1])
1436
Thomas Woutersb2137042007-02-01 18:02:27 +00001437class SubclassWithKwargsTest(unittest.TestCase):
1438 def test_keywords_in_subclass(self):
1439 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001440 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001441 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00001442 class Subclass(cls):
1443 def __init__(self, newarg=None, *args):
1444 cls.__init__(self, *args)
1445 try:
1446 Subclass(newarg=1)
1447 except TypeError as err:
1448 # we expect type errors because of wrong argument count
Ezio Melottib58e0bd2010-01-23 15:40:09 +00001449 self.assertNotIn("does not take keyword arguments", err.args[0])
Thomas Woutersb2137042007-02-01 18:02:27 +00001450
1451
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001452libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001453
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001454
1455>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001456>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00001457... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001458...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001459Check 1200 is for $120.15
1460Check 1201 is for $764.05
1461Check 1202 is for $823.14
1462
1463>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001464>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00001465... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00001466...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000014671
14688
146927
1470
1471>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001472>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00001473... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00001474...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001475Alex
1476Laura
1477Martin
1478Walter
1479Samuele
1480
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001481>>> from operator import itemgetter
1482>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00001483>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001484>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001485... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001486...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000014871 ['a', 'c', 'e']
14882 ['b', 'd', 'f']
14893 ['g']
1490
Raymond Hettinger734fb572004-01-20 20:04:40 +00001491# Find runs of consecutive numbers using groupby. The key to the solution
1492# is differencing with a range so that consecutive numbers all appear in
1493# same group.
1494>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00001495>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001496... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001497...
Raymond Hettinger734fb572004-01-20 20:04:40 +00001498[1]
1499[4, 5, 6]
1500[10]
1501[15, 16, 17, 18]
1502[22]
1503[25, 26, 27, 28]
1504
Georg Brandl3dbca812008-07-23 16:10:53 +00001505>>> def take(n, iterable):
1506... "Return first n items of the iterable as a list"
1507... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001508
Georg Brandl3dbca812008-07-23 16:10:53 +00001509>>> def enumerate(iterable, start=0):
1510... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001511
Georg Brandl3dbca812008-07-23 16:10:53 +00001512>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001513... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00001514... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001515
Raymond Hettingercdf8ba32009-02-19 04:45:07 +00001516>>> def nth(iterable, n, default=None):
1517... "Returns the nth item or a default value"
1518... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001519
Georg Brandl3dbca812008-07-23 16:10:53 +00001520>>> def quantify(iterable, pred=bool):
1521... "Count how many times the predicate is true"
1522... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001523
Georg Brandl3dbca812008-07-23 16:10:53 +00001524>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001525... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00001526... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001527
Georg Brandl3dbca812008-07-23 16:10:53 +00001528>>> def ncycles(iterable, n):
1529... "Returns the seqeuence elements n times"
1530... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001531
1532>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001533... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001534
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001535>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001536... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001537
1538>>> def repeatfunc(func, times=None, *args):
1539... "Repeat calls to func with specified arguments."
1540... " Example: repeatfunc(random.random)"
1541... if times is None:
1542... return starmap(func, repeat(args))
1543... else:
1544... return starmap(func, repeat(args, times))
1545
Raymond Hettingerd591f662003-10-26 15:34:50 +00001546>>> def pairwise(iterable):
1547... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1548... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001549... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00001550... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001551... except StopIteration:
1552... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001553... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001554
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001555>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001556... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001557... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00001558... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001559
1560>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001561... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001562... # Recipe credited to George Sakkis
1563... pending = len(iterables)
1564... nexts = cycle(iter(it).__next__ for it in iterables)
1565... while pending:
1566... try:
1567... for next in nexts:
1568... yield next()
1569... except StopIteration:
1570... pending -= 1
1571... nexts = cycle(islice(nexts, pending))
1572
1573>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00001574... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1575... s = list(iterable)
1576... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001577
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001578>>> def unique_everseen(iterable, key=None):
1579... "List unique elements, preserving order. Remember all elements ever seen."
1580... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1581... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1582... seen = set()
1583... seen_add = seen.add
1584... if key is None:
1585... for element in iterable:
1586... if element not in seen:
1587... seen_add(element)
1588... yield element
1589... else:
1590... for element in iterable:
1591... k = key(element)
1592... if k not in seen:
1593... seen_add(k)
1594... yield element
1595
1596>>> def unique_justseen(iterable, key=None):
1597... "List unique elements, preserving order. Remember only the element just seen."
1598... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1599... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1600... return map(next, map(itemgetter(1), groupby(iterable, key)))
1601
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001602This is not part of the examples but it tests to make sure the definitions
1603perform as purported.
1604
Raymond Hettingera098b332003-09-08 23:58:40 +00001605>>> take(10, count())
1606[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1607
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001608>>> list(enumerate('abc'))
1609[(0, 'a'), (1, 'b'), (2, 'c')]
1610
1611>>> list(islice(tabulate(lambda x: 2*x), 4))
1612[0, 2, 4, 6]
1613
1614>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00001615'd'
1616
1617>>> nth('abcde', 9) is None
1618True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001619
Guido van Rossum805365e2007-05-07 22:24:25 +00001620>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000162150
1622
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001623>>> a = [[1, 2, 3], [4, 5, 6]]
1624>>> flatten(a)
1625[1, 2, 3, 4, 5, 6]
1626
1627>>> list(repeatfunc(pow, 5, 2, 3))
1628[8, 8, 8, 8, 8]
1629
1630>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001631>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001632[0, 0, 0, 0, 0]
1633
Raymond Hettingerd591f662003-10-26 15:34:50 +00001634>>> list(pairwise('abcd'))
1635[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001636
Raymond Hettingerd591f662003-10-26 15:34:50 +00001637>>> list(pairwise([]))
1638[]
1639
1640>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001641[]
1642
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001643>>> list(islice(padnone('abc'), 0, 6))
1644['a', 'b', 'c', None, None, None]
1645
1646>>> list(ncycles('abc', 3))
1647['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1648
1649>>> dotproduct([1,2,3], [4,5,6])
165032
1651
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001652>>> list(grouper(3, 'abcdefg', 'x'))
1653[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1654
1655>>> list(roundrobin('abc', 'd', 'ef'))
1656['a', 'd', 'e', 'b', 'f', 'c']
1657
Raymond Hettingerace67332009-01-26 02:23:50 +00001658>>> list(powerset([1,2,3]))
1659[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001660
Raymond Hettinger191e8502009-01-27 13:29:43 +00001661>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1662True
1663
1664>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1665True
1666
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001667>>> list(unique_everseen('AAAABBBCCDAABBB'))
1668['A', 'B', 'C', 'D']
1669
1670>>> list(unique_everseen('ABBCcAD', str.lower))
1671['A', 'B', 'C', 'D']
1672
1673>>> list(unique_justseen('AAAABBBCCDAABBB'))
1674['A', 'B', 'C', 'D', 'A', 'B']
1675
1676>>> list(unique_justseen('ABBCcAD', str.lower))
1677['A', 'B', 'C', 'A', 'D']
1678
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001679"""
1680
1681__test__ = {'libreftest' : libreftest}
1682
1683def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001684 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00001685 RegressionTests, LengthTransparency,
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001686 SubclassWithKwargsTest, TestExamples)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001687 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001688
1689 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001690 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001691 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001692 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00001693 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001694 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001695 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001696 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001697 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001698
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001699 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001700 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001701
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001702if __name__ == "__main__":
1703 test_main(verbose=True)