blob: 90e85a9b25d2dc17fc0d58be52faa8f692fecfc3 [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
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000040def tupleize(*args):
41 return args
42
43def irange(n):
44 for i in range(n):
45 yield i
46
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000047class StopNow:
48 'Class emulating an empty iterable.'
49 def __iter__(self):
50 return self
Georg Brandla18af4e2007-04-21 15:47:16 +000051 def __next__(self):
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000052 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000053
Raymond Hettinger02420702003-06-29 20:36:23 +000054def take(n, seq):
55 'Convenience function for partially consuming a long of infinite iterable'
56 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000057
Christian Heimes78644762008-03-04 23:39:23 +000058def prod(iterable):
59 return reduce(operator.mul, iterable, 1)
60
Christian Heimes380f7f22008-02-28 11:19:05 +000061def fact(n):
62 'Factorial'
Christian Heimes78644762008-03-04 23:39:23 +000063 return prod(range(1, n+1))
64
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000065# root level methods for pickling ability
66def testR(r):
67 return r[0]
68
69def testR2(r):
70 return r[2]
71
72def underten(x):
73 return x<10
74
Raymond Hettinger96ef8112003-02-01 00:10:11 +000075class TestBasicOps(unittest.TestCase):
Raymond Hettinger482ba772010-12-01 22:48:00 +000076
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000077 def pickletest(self, it, stop=4, take=1, compare=None):
78 """Test that an iterator is the same after pickling, also when part-consumed"""
79 def expand(it, i=0):
80 # Recursively expand iterables, within sensible bounds
81 if i > 10:
82 raise RuntimeError("infinite recursion encountered")
83 if isinstance(it, str):
84 return it
85 try:
86 l = list(islice(it, stop))
87 except TypeError:
88 return it # can't expand it
89 return [expand(e, i+1) for e in l]
90
91 # Test the initial copy against the original
92 dump = pickle.dumps(it)
93 i2 = pickle.loads(dump)
94 self.assertEqual(type(it), type(i2))
95 a, b = expand(it), expand(i2)
96 self.assertEqual(a, b)
97 if compare:
98 c = expand(compare)
99 self.assertEqual(a, c)
100
101 # Take from the copy, and create another copy and compare them.
102 i3 = pickle.loads(dump)
103 took = 0
104 try:
105 for i in range(take):
106 next(i3)
107 took += 1
108 except StopIteration:
109 pass #in case there is less data than 'take'
110 dump = pickle.dumps(i3)
111 i4 = pickle.loads(dump)
112 a, b = expand(i3), expand(i4)
113 self.assertEqual(a, b)
114 if compare:
115 c = expand(compare[took:])
116 self.assertEqual(a, c);
117
Raymond Hettinger482ba772010-12-01 22:48:00 +0000118 def test_accumulate(self):
119 self.assertEqual(list(accumulate(range(10))), # one positional arg
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000120 [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])
121 self.assertEqual(list(accumulate(iterable=range(10))), # kw arg
122 [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])
Raymond Hettinger482ba772010-12-01 22:48:00 +0000123 for typ in int, complex, Decimal, Fraction: # multiple types
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000124 self.assertEqual(
125 list(accumulate(map(typ, range(10)))),
Raymond Hettinger482ba772010-12-01 22:48:00 +0000126 list(map(typ, [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])))
Raymond Hettinger2d93e6e2010-12-03 02:33:53 +0000127 self.assertEqual(list(accumulate('abc')), ['a', 'ab', 'abc']) # works with non-numeric
Raymond Hettinger482ba772010-12-01 22:48:00 +0000128 self.assertEqual(list(accumulate([])), []) # empty iterable
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000129 self.assertEqual(list(accumulate([7])), [7]) # iterable of length one
Raymond Hettinger5d446132011-03-27 18:52:10 -0700130 self.assertRaises(TypeError, accumulate, range(10), 5, 6) # too many args
Raymond Hettinger482ba772010-12-01 22:48:00 +0000131 self.assertRaises(TypeError, accumulate) # too few args
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000132 self.assertRaises(TypeError, accumulate, x=range(10)) # unexpected kwd arg
Raymond Hettinger482ba772010-12-01 22:48:00 +0000133 self.assertRaises(TypeError, list, accumulate([1, []])) # args that don't add
134
Raymond Hettinger5d446132011-03-27 18:52:10 -0700135 s = [2, 8, 9, 5, 7, 0, 3, 4, 1, 6]
136 self.assertEqual(list(accumulate(s, min)),
137 [2, 2, 2, 2, 2, 0, 0, 0, 0, 0])
138 self.assertEqual(list(accumulate(s, max)),
139 [2, 8, 9, 9, 9, 9, 9, 9, 9, 9])
140 self.assertEqual(list(accumulate(s, operator.mul)),
141 [2, 16, 144, 720, 5040, 0, 0, 0, 0, 0])
142 with self.assertRaises(TypeError):
143 list(accumulate(s, chr)) # unary-operation
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000144 self.pickletest(accumulate(range(10))) # test pickling
Raymond Hettinger5d446132011-03-27 18:52:10 -0700145
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000146 def test_chain(self):
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000147
148 def chain2(*iterables):
149 'Pure python version in the docs'
150 for it in iterables:
151 for element in it:
152 yield element
153
154 for c in (chain, chain2):
155 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
156 self.assertEqual(list(c('abc')), list('abc'))
157 self.assertEqual(list(c('')), [])
158 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
159 self.assertRaises(TypeError, list,c(2, 3))
Christian Heimesf16baeb2008-02-29 14:57:44 +0000160
161 def test_chain_from_iterable(self):
162 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
163 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
164 self.assertEqual(list(chain.from_iterable([''])), [])
165 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
166 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000167
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000168 def test_chain_reducible(self):
169 operators = [copy.deepcopy,
170 lambda s: pickle.loads(pickle.dumps(s))]
171 for oper in operators:
172 it = chain('abc', 'def')
173 self.assertEqual(list(oper(it)), list('abcdef'))
174 self.assertEqual(next(it), 'a')
175 self.assertEqual(list(oper(it)), list('bcdef'))
176
177 self.assertEqual(list(oper(chain(''))), [])
178 self.assertEqual(take(4, oper(chain('abc', 'def'))), list('abcd'))
179 self.assertRaises(TypeError, list, oper(chain(2, 3)))
180 self.pickletest(chain('abc', 'def'), compare=list('abcdef'))
181
Christian Heimes380f7f22008-02-28 11:19:05 +0000182 def test_combinations(self):
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000183 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Christian Heimes380f7f22008-02-28 11:19:05 +0000184 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Christian Heimes78644762008-03-04 23:39:23 +0000185 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Christian Heimes380f7f22008-02-28 11:19:05 +0000186 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000187
188 for op in (lambda a:a, lambda a:pickle.loads(pickle.dumps(a))):
189 self.assertEqual(list(op(combinations('abc', 32))), []) # r > n
190
191 self.assertEqual(list(op(combinations('ABCD', 2))),
192 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
193 testIntermediate = combinations('ABCD', 2)
194 next(testIntermediate)
195 self.assertEqual(list(op(testIntermediate)),
196 [('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
197
198 self.assertEqual(list(op(combinations(range(4), 3))),
199 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
200 testIntermediate = combinations(range(4), 3)
201 next(testIntermediate)
202 self.assertEqual(list(op(testIntermediate)),
203 [(0,1,3), (0,2,3), (1,2,3)])
204
Christian Heimes78644762008-03-04 23:39:23 +0000205
206 def combinations1(iterable, r):
207 'Pure python version shown in the docs'
208 pool = tuple(iterable)
209 n = len(pool)
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000210 if r > n:
211 return
Christian Heimes78644762008-03-04 23:39:23 +0000212 indices = list(range(r))
213 yield tuple(pool[i] for i in indices)
214 while 1:
215 for i in reversed(range(r)):
216 if indices[i] != i + n - r:
217 break
218 else:
219 return
220 indices[i] += 1
221 for j in range(i+1, r):
222 indices[j] = indices[j-1] + 1
223 yield tuple(pool[i] for i in indices)
224
225 def combinations2(iterable, r):
226 'Pure python version shown in the docs'
227 pool = tuple(iterable)
228 n = len(pool)
229 for indices in permutations(range(n), r):
230 if sorted(indices) == list(indices):
231 yield tuple(pool[i] for i in indices)
232
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000233 def combinations3(iterable, r):
234 'Pure python version from cwr()'
235 pool = tuple(iterable)
236 n = len(pool)
237 for indices in combinations_with_replacement(range(n), r):
238 if len(set(indices)) == r:
239 yield tuple(pool[i] for i in indices)
240
Christian Heimes78644762008-03-04 23:39:23 +0000241 for n in range(7):
Christian Heimes380f7f22008-02-28 11:19:05 +0000242 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000243 for r in range(n+2):
Christian Heimes380f7f22008-02-28 11:19:05 +0000244 result = list(combinations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000245 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 +0000246 self.assertEqual(len(result), len(set(result))) # no repeats
247 self.assertEqual(result, sorted(result)) # lexicographic order
248 for c in result:
249 self.assertEqual(len(c), r) # r-length combinations
250 self.assertEqual(len(set(c)), r) # no duplicate elements
251 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000252 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000253 self.assertEqual(list(c),
254 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000255 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000256 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000257 self.assertEqual(result, list(combinations3(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000258
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000259 self.pickletest(combinations(values, r)) # test pickling
260
261 # Test implementation detail: tuple re-use
Alex Gaynore151d212011-07-17 16:21:30 -0700262 @support.impl_detail("tuple reuse is specific to CPython")
263 def test_combinations_tuple_reuse(self):
Christian Heimes78644762008-03-04 23:39:23 +0000264 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
265 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
266
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000267 def test_combinations_with_replacement(self):
268 cwr = combinations_with_replacement
269 self.assertRaises(TypeError, cwr, 'abc') # missing r argument
270 self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
271 self.assertRaises(TypeError, cwr, None) # pool is not iterable
272 self.assertRaises(ValueError, cwr, 'abc', -2) # r is negative
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000273
274 for op in (lambda a:a, lambda a:pickle.loads(pickle.dumps(a))):
275 self.assertEqual(list(op(cwr('ABC', 2))),
276 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
277 testIntermediate = cwr('ABC', 2)
278 next(testIntermediate)
279 self.assertEqual(list(op(testIntermediate)),
280 [('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
281
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000282
283 def cwr1(iterable, r):
284 'Pure python version shown in the docs'
285 # number items returned: (n+r-1)! / r! / (n-1)! when n>0
286 pool = tuple(iterable)
287 n = len(pool)
288 if not n and r:
289 return
290 indices = [0] * r
291 yield tuple(pool[i] for i in indices)
292 while 1:
293 for i in reversed(range(r)):
294 if indices[i] != n - 1:
295 break
296 else:
297 return
298 indices[i:] = [indices[i] + 1] * (r - i)
299 yield tuple(pool[i] for i in indices)
300
301 def cwr2(iterable, r):
302 'Pure python version shown in the docs'
303 pool = tuple(iterable)
304 n = len(pool)
305 for indices in product(range(n), repeat=r):
306 if sorted(indices) == list(indices):
307 yield tuple(pool[i] for i in indices)
308
309 def numcombs(n, r):
310 if not n:
311 return 0 if r else 1
312 return fact(n+r-1) / fact(r)/ fact(n-1)
313
314 for n in range(7):
315 values = [5*x-12 for x in range(n)]
316 for r in range(n+2):
317 result = list(cwr(values, r))
318
319 self.assertEqual(len(result), numcombs(n, r)) # right number of combs
320 self.assertEqual(len(result), len(set(result))) # no repeats
321 self.assertEqual(result, sorted(result)) # lexicographic order
322
323 regular_combs = list(combinations(values, r)) # compare to combs without replacement
324 if n == 0 or r <= 1:
Ezio Melottib3aedd42010-11-20 19:04:17 +0000325 self.assertEqual(result, regular_combs) # cases that should be identical
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000326 else:
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000327 self.assertTrue(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000328
329 for c in result:
330 self.assertEqual(len(c), r) # r-length combinations
331 noruns = [k for k,v in groupby(c)] # combo without consecutive repeats
332 self.assertEqual(len(noruns), len(set(noruns))) # no repeats other than consecutive
333 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000334 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000335 self.assertEqual(noruns,
336 [e for e in values if e in c]) # comb is a subsequence of the input iterable
337 self.assertEqual(result, list(cwr1(values, r))) # matches first pure python version
338 self.assertEqual(result, list(cwr2(values, r))) # matches second pure python version
339
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000340 self.pickletest(cwr(values,r)) # test pickling
341
342 # Test implementation detail: tuple re-use
343
Alex Gaynore151d212011-07-17 16:21:30 -0700344 @support.impl_detail("tuple reuse is specific to CPython")
345 def test_combinations_with_replacement_tuple_reuse(self):
346 cwr = combinations_with_replacement
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000347 self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
348 self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
349
Christian Heimes78644762008-03-04 23:39:23 +0000350 def test_permutations(self):
351 self.assertRaises(TypeError, permutations) # too few arguments
352 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000353 self.assertRaises(TypeError, permutations, None) # pool is not iterable
354 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000355 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000356 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Christian Heimes78644762008-03-04 23:39:23 +0000357 self.assertEqual(list(permutations(range(3), 2)),
358 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
359
360 def permutations1(iterable, r=None):
361 'Pure python version shown in the docs'
362 pool = tuple(iterable)
363 n = len(pool)
364 r = n if r is None else r
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000365 if r > n:
366 return
Christian Heimes78644762008-03-04 23:39:23 +0000367 indices = list(range(n))
368 cycles = list(range(n-r+1, n+1))[::-1]
369 yield tuple(pool[i] for i in indices[:r])
370 while n:
371 for i in reversed(range(r)):
372 cycles[i] -= 1
373 if cycles[i] == 0:
374 indices[i:] = indices[i+1:] + indices[i:i+1]
375 cycles[i] = n - i
376 else:
377 j = cycles[i]
378 indices[i], indices[-j] = indices[-j], indices[i]
379 yield tuple(pool[i] for i in indices[:r])
380 break
381 else:
382 return
383
384 def permutations2(iterable, r=None):
385 'Pure python version shown in the docs'
386 pool = tuple(iterable)
387 n = len(pool)
388 r = n if r is None else r
389 for indices in product(range(n), repeat=r):
390 if len(set(indices)) == r:
391 yield tuple(pool[i] for i in indices)
392
393 for n in range(7):
394 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000395 for r in range(n+2):
Christian Heimes78644762008-03-04 23:39:23 +0000396 result = list(permutations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000397 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 +0000398 self.assertEqual(len(result), len(set(result))) # no repeats
399 self.assertEqual(result, sorted(result)) # lexicographic order
400 for p in result:
401 self.assertEqual(len(p), r) # r-length permutations
402 self.assertEqual(len(set(p)), r) # no duplicate elements
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000403 self.assertTrue(all(e in values for e in p)) # elements taken from input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000404 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000405 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000406 if r == n:
407 self.assertEqual(result, list(permutations(values, None))) # test r as None
408 self.assertEqual(result, list(permutations(values))) # test default r
409
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000410 self.pickletest(permutations(values, r)) # test pickling
411
Alex Gaynore151d212011-07-17 16:21:30 -0700412 @support.impl_detail("tuple resuse is CPython specific")
413 def test_permutations_tuple_reuse(self):
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000414 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Christian Heimes78644762008-03-04 23:39:23 +0000415 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Christian Heimes380f7f22008-02-28 11:19:05 +0000416
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000417 def test_combinatorics(self):
418 # Test relationships between product(), permutations(),
419 # combinations() and combinations_with_replacement().
420
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000421 for n in range(6):
422 s = 'ABCDEFG'[:n]
423 for r in range(8):
424 prod = list(product(s, repeat=r))
425 cwr = list(combinations_with_replacement(s, r))
426 perm = list(permutations(s, r))
427 comb = list(combinations(s, r))
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000428
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000429 # Check size
Ezio Melottib3aedd42010-11-20 19:04:17 +0000430 self.assertEqual(len(prod), n**r)
431 self.assertEqual(len(cwr), (fact(n+r-1) / fact(r)/ fact(n-1)) if n else (not r))
432 self.assertEqual(len(perm), 0 if r>n else fact(n) / fact(n-r))
433 self.assertEqual(len(comb), 0 if r>n else fact(n) / fact(r) / fact(n-r))
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000434
435 # Check lexicographic order without repeated tuples
Ezio Melottib3aedd42010-11-20 19:04:17 +0000436 self.assertEqual(prod, sorted(set(prod)))
437 self.assertEqual(cwr, sorted(set(cwr)))
438 self.assertEqual(perm, sorted(set(perm)))
439 self.assertEqual(comb, sorted(set(comb)))
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000440
441 # Check interrelationships
Ezio Melottib3aedd42010-11-20 19:04:17 +0000442 self.assertEqual(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
443 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 +0000444 self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
445 self.assertEqual(comb, [t for t in cwr if len(set(t))==r]) # comb: cwrs without dups
446 self.assertEqual(comb, list(filter(set(cwr).__contains__, perm))) # comb: perm that is a cwr
447 self.assertEqual(comb, list(filter(set(perm).__contains__, cwr))) # comb: cwr that is a perm
448 self.assertEqual(comb, sorted(set(cwr) & set(perm))) # comb: both a cwr and a perm
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000449
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000450 def test_compress(self):
Raymond Hettinger15a49502009-02-19 02:17:09 +0000451 self.assertEqual(list(compress(data='ABCDEF', selectors=[1,0,1,0,1,1])), list('ACEF'))
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000452 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
453 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
454 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
455 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
456 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
457 n = 10000
458 data = chain.from_iterable(repeat(range(6), n))
459 selectors = chain.from_iterable(repeat((0, 1)))
460 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
461 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
462 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
463 self.assertRaises(TypeError, compress, range(6)) # too few args
464 self.assertRaises(TypeError, compress, range(6), None) # too many args
465
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000466 # check copy, deepcopy, pickle
467 for op in (lambda a:copy.copy(a), lambda a:copy.deepcopy(a), lambda a:pickle.loads(pickle.dumps(a))):
468 for data, selectors, result1, result2 in [
469 ('ABCDEF', [1,0,1,0,1,1], 'ACEF', 'CEF'),
470 ('ABCDEF', [0,0,0,0,0,0], '', ''),
471 ('ABCDEF', [1,1,1,1,1,1], 'ABCDEF', 'BCDEF'),
472 ('ABCDEF', [1,0,1], 'AC', 'C'),
473 ('ABC', [0,1,1,1,1,1], 'BC', 'C'),
474 ]:
475
476 self.assertEqual(list(op(compress(data=data, selectors=selectors))), list(result1))
477 self.assertEqual(list(op(compress(data, selectors))), list(result1))
478 testIntermediate = compress(data, selectors)
479 if result1:
480 next(testIntermediate)
481 self.assertEqual(list(op(testIntermediate)), list(result2))
482
483
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000484 def test_count(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000485 self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
486 self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
487 self.assertEqual(take(2, lzip('abc',count(3))), [('a', 3), ('b', 4)])
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000488 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
489 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000490 self.assertRaises(TypeError, count, 2, 3, 4)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000491 self.assertRaises(TypeError, count, 'a')
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000492 self.assertEqual(list(islice(count(maxsize-5), 10)),
493 list(range(maxsize-5, maxsize+5)))
494 self.assertEqual(list(islice(count(-maxsize-5), 10)),
495 list(range(-maxsize-5, -maxsize+5)))
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +0000496 self.assertEqual(list(islice(count(10, maxsize+5), 3)),
497 list(range(10, 10+3*(maxsize+5), maxsize+5)))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000498 c = count(3)
499 self.assertEqual(repr(c), 'count(3)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000500 next(c)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000501 self.assertEqual(repr(c), 'count(4)')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000502 c = count(-9)
503 self.assertEqual(repr(c), 'count(-9)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000504 next(c)
Raymond Hettinger30729212009-02-12 06:28:27 +0000505 self.assertEqual(repr(count(10.25)), 'count(10.25)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000506 self.assertEqual(next(c), -8)
Christian Heimesa37d4c62007-12-04 23:02:19 +0000507 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 +0000508 # Test repr (ignoring the L in longs)
509 r1 = repr(count(i)).replace('L', '')
510 r2 = 'count(%r)'.__mod__(i).replace('L', '')
511 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000512
Raymond Hettingera3e1ad22009-11-30 22:02:31 +0000513 # check copy, deepcopy, pickle
514 for value in -3, 3, maxsize-5, maxsize+5:
515 c = count(value)
516 self.assertEqual(next(copy.copy(c)), value)
517 self.assertEqual(next(copy.deepcopy(c)), value)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000518 self.pickletest(count(value))
Raymond Hettingera3e1ad22009-11-30 22:02:31 +0000519
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +0000520 #check proper internal error handling for large "step' sizes
521 count(1, maxsize+5); sys.exc_info()
522
Raymond Hettinger30729212009-02-12 06:28:27 +0000523 def test_count_with_stride(self):
524 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingereb13fdd2009-02-21 22:30:12 +0000525 self.assertEqual(lzip('abc',count(start=2,step=3)),
526 [('a', 2), ('b', 5), ('c', 8)])
527 self.assertEqual(lzip('abc',count(step=-1)),
528 [('a', 0), ('b', -1), ('c', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000529 self.assertEqual(lzip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
530 self.assertEqual(lzip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000531 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000532 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
533 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
534 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettinger8a882a72009-02-12 12:08:18 +0000535 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
536 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerde09acf2009-02-12 12:53:53 +0000537 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
538 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000539 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
540 c = count(3, 5)
541 self.assertEqual(repr(c), 'count(3, 5)')
542 next(c)
543 self.assertEqual(repr(c), 'count(8, 5)')
544 c = count(-9, 0)
545 self.assertEqual(repr(c), 'count(-9, 0)')
546 next(c)
547 self.assertEqual(repr(c), 'count(-9, 0)')
548 c = count(-9, -3)
549 self.assertEqual(repr(c), 'count(-9, -3)')
550 next(c)
551 self.assertEqual(repr(c), 'count(-12, -3)')
552 self.assertEqual(repr(c), 'count(-12, -3)')
553 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
554 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
555 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
556 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
557 for j in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 1, 10, sys.maxsize-5, sys.maxsize+5):
558 # Test repr (ignoring the L in longs)
559 r1 = repr(count(i, j)).replace('L', '')
560 if j == 1:
561 r2 = ('count(%r)' % i).replace('L', '')
562 else:
563 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
564 self.assertEqual(r1, r2)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000565 self.pickletest(count(i, j))
Raymond Hettinger30729212009-02-12 06:28:27 +0000566
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000567 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000568 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000569 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000570 self.assertRaises(TypeError, cycle)
571 self.assertRaises(TypeError, cycle, 5)
572 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000573
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000574 # check copy, deepcopy, pickle
575 c = cycle('abc')
576 self.assertEqual(next(c), 'a')
577 #simple copy currently not supported, because __reduce__ returns
578 #an internal iterator
579 #self.assertEqual(take(10, copy.copy(c)), list('bcabcabcab'))
580 self.assertEqual(take(10, copy.deepcopy(c)), list('bcabcabcab'))
581 self.assertEqual(take(10, pickle.loads(pickle.dumps(c))), list('bcabcabcab'))
582 next(c)
583 self.assertEqual(take(10, pickle.loads(pickle.dumps(c))), list('cabcabcabc'))
584 self.pickletest(cycle('abc'))
585
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000586 def test_groupby(self):
587 # Check whether it accepts arguments correctly
588 self.assertEqual([], list(groupby([])))
589 self.assertEqual([], list(groupby([], key=id)))
590 self.assertRaises(TypeError, list, groupby('abc', []))
591 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000592 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000593
594 # Check normal input
595 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
596 (2,15,22), (3,16,23), (3,17,23)]
597 dup = []
598 for k, g in groupby(s, lambda r:r[0]):
599 for elem in g:
600 self.assertEqual(k, elem[0])
601 dup.append(elem)
602 self.assertEqual(s, dup)
603
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000604 # Check normal pickled
605 dup = []
606 for k, g in pickle.loads(pickle.dumps(groupby(s, testR))):
607 for elem in g:
608 self.assertEqual(k, elem[0])
609 dup.append(elem)
610 self.assertEqual(s, dup)
611
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000612 # Check nested case
613 dup = []
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000614 for k, g in groupby(s, testR):
615 for ik, ig in groupby(g, testR2):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000616 for elem in ig:
617 self.assertEqual(k, elem[0])
618 self.assertEqual(ik, elem[2])
619 dup.append(elem)
620 self.assertEqual(s, dup)
621
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000622 # Check nested and pickled
623 dup = []
624 for k, g in pickle.loads(pickle.dumps(groupby(s, testR))):
625 for ik, ig in pickle.loads(pickle.dumps(groupby(g, testR2))):
626 for elem in ig:
627 self.assertEqual(k, elem[0])
628 self.assertEqual(ik, elem[2])
629 dup.append(elem)
630 self.assertEqual(s, dup)
631
632
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000633 # Check case where inner iterator is not used
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000634 keys = [k for k, g in groupby(s, testR)]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000635 expectedkeys = set([r[0] for r in s])
636 self.assertEqual(set(keys), expectedkeys)
637 self.assertEqual(len(keys), len(expectedkeys))
638
639 # Exercise pipes and filters style
640 s = 'abracadabra'
641 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000642 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000643 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
644 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000645 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000646 self.assertEqual(r, ['a', 'b', 'r'])
647 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000648 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000649 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
650 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000651 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000652 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
653
Georg Brandla18af4e2007-04-21 15:47:16 +0000654 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000655 class ExpectedError(Exception):
656 pass
657 def delayed_raise(n=0):
658 for i in range(n):
659 yield 'yo'
660 raise ExpectedError
661 def gulp(iterable, keyp=None, func=list):
662 return [func(g) for k, g in groupby(iterable, keyp)]
663
Georg Brandla18af4e2007-04-21 15:47:16 +0000664 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000665 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000666 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000667 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
668
669 # __cmp__ failure
670 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000671 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000672 raise ExpectedError
673 s = [DummyCmp(), DummyCmp(), None]
674
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000675 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000676 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000677 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000678 self.assertRaises(ExpectedError, gulp, s)
679
680 # keyfunc failure
681 def keyfunc(obj):
682 if keyfunc.skip > 0:
683 keyfunc.skip -= 1
684 return obj
685 else:
686 raise ExpectedError
687
688 # keyfunc failure on outer object
689 keyfunc.skip = 0
690 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
691 keyfunc.skip = 1
692 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
693
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000694 def test_filter(self):
695 self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
696 self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])
697 self.assertEqual(list(filter(bool, [0,1,0,2,0])), [1,2])
698 self.assertEqual(take(4, filter(isEven, count())), [0,2,4,6])
699 self.assertRaises(TypeError, filter)
700 self.assertRaises(TypeError, filter, lambda x:x)
701 self.assertRaises(TypeError, filter, lambda x:x, range(6), 7)
702 self.assertRaises(TypeError, filter, isEven, 3)
703 self.assertRaises(TypeError, next, filter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000704
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000705 # check copy, deepcopy, pickle
706 ans = [0,2,4]
707
708 c = filter(isEven, range(6))
709 self.assertEqual(list(copy.copy(c)), ans)
710 c = filter(isEven, range(6))
711 self.assertEqual(list(copy.deepcopy(c)), ans)
712 c = filter(isEven, range(6))
713 self.assertEqual(list(pickle.loads(pickle.dumps(c))), ans)
714 next(c)
715 self.assertEqual(list(pickle.loads(pickle.dumps(c))), ans[1:])
716 c = filter(isEven, range(6))
717 self.pickletest(c)
718
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000719 def test_filterfalse(self):
720 self.assertEqual(list(filterfalse(isEven, range(6))), [1,3,5])
721 self.assertEqual(list(filterfalse(None, [0,1,0,2,0])), [0,0,0])
722 self.assertEqual(list(filterfalse(bool, [0,1,0,2,0])), [0,0,0])
723 self.assertEqual(take(4, filterfalse(isEven, count())), [1,3,5,7])
724 self.assertRaises(TypeError, filterfalse)
725 self.assertRaises(TypeError, filterfalse, lambda x:x)
726 self.assertRaises(TypeError, filterfalse, lambda x:x, range(6), 7)
727 self.assertRaises(TypeError, filterfalse, isEven, 3)
728 self.assertRaises(TypeError, next, filterfalse(range(6), range(6)))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000729 self.pickletest(filterfalse(isEven, range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000730
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000731 def test_zip(self):
732 # XXX This is rather silly now that builtin zip() calls zip()...
733 ans = [(x,y) for x, y in zip('abc',count())]
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000734 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000735 self.assertEqual(list(zip('abc', range(6))), lzip('abc', range(6)))
736 self.assertEqual(list(zip('abcdef', range(3))), lzip('abcdef', range(3)))
737 self.assertEqual(take(3,zip('abcdef', count())), lzip('abcdef', range(3)))
738 self.assertEqual(list(zip('abcdef')), lzip('abcdef'))
739 self.assertEqual(list(zip()), lzip())
740 self.assertRaises(TypeError, zip, 3)
741 self.assertRaises(TypeError, zip, range(3), 3)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000742 self.assertEqual([tuple(list(pair)) for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000743 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000744 self.assertEqual([pair for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000745 lzip('abc', 'def'))
Alex Gaynore151d212011-07-17 16:21:30 -0700746
747 @support.impl_detail("tuple reuse is specific to CPython")
748 def test_zip_tuple_reuse(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000749 ids = list(map(id, zip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000750 self.assertEqual(min(ids), max(ids))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000751 ids = list(map(id, list(zip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000752 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000753
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000754 # check copy, deepcopy, pickle
755 ans = [(x,y) for x, y in copy.copy(zip('abc',count()))]
756 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
757
758 ans = [(x,y) for x, y in copy.deepcopy(zip('abc',count()))]
759 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
760
761 ans = [(x,y) for x, y in pickle.loads(pickle.dumps(zip('abc',count())))]
762 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
763
764 testIntermediate = zip('abc',count())
765 next(testIntermediate)
766 ans = [(x,y) for x, y in pickle.loads(pickle.dumps(testIntermediate))]
767 self.assertEqual(ans, [('b', 1), ('c', 2)])
768
769 self.pickletest(zip('abc', count()))
770
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000771 def test_ziplongest(self):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000772 for args in [
773 ['abc', range(6)],
774 [range(6), 'abc'],
775 [range(1000), range(2000,2100), range(3000,3050)],
776 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
777 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
778 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000779 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
780 for i in range(max(map(len, args)))]
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000781 self.assertEqual(list(zip_longest(*args)), target)
782 self.assertEqual(list(zip_longest(*args, **{})), target)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000783 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 +0000784 self.assertEqual(list(zip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000785
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000786 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 +0000787
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000788 self.assertEqual(list(zip_longest()), list(zip()))
789 self.assertEqual(list(zip_longest([])), list(zip([])))
790 self.assertEqual(list(zip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000791
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000792 self.assertEqual(list(zip_longest('abc', 'defg', **{})),
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000793 list(zip(list('abc')+[None], 'defg'))) # empty keyword dict
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000794 self.assertRaises(TypeError, zip_longest, 3)
795 self.assertRaises(TypeError, zip_longest, range(3), 3)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000796
797 for stmt in [
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000798 "zip_longest('abc', fv=1)",
799 "zip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000800 ]:
801 try:
802 eval(stmt, globals(), locals())
803 except TypeError:
804 pass
805 else:
806 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000807
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000808 self.assertEqual([tuple(list(pair)) for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000809 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000810 self.assertEqual([pair for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000811 list(zip('abc', 'def')))
Alex Gaynore151d212011-07-17 16:21:30 -0700812
813 @support.impl_detail("tuple reuse is specific to CPython")
814 def test_zip_longest_tuple_reuse(self):
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000815 ids = list(map(id, zip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000816 self.assertEqual(min(ids), max(ids))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000817 ids = list(map(id, list(zip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000818 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
819
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000820 def test_zip_longest_pickling(self):
821 self.pickletest(zip_longest("abc", "def"))
822 self.pickletest(zip_longest("abc", "defgh"))
823 self.pickletest(zip_longest("abc", "defgh", fillvalue=1))
824 self.pickletest(zip_longest("", "defgh"))
825
Raymond Hettingerfc438512009-11-01 20:55:33 +0000826 def test_bug_7244(self):
827
828 class Repeater:
829 # this class is similar to itertools.repeat
830 def __init__(self, o, t, e):
831 self.o = o
832 self.t = int(t)
833 self.e = e
834 def __iter__(self): # its iterator is itself
835 return self
836 def __next__(self):
837 if self.t > 0:
838 self.t -= 1
839 return self.o
840 else:
841 raise self.e
842
843 # Formerly this code in would fail in debug mode
844 # with Undetected Error and Stop Iteration
845 r1 = Repeater(1, 3, StopIteration)
846 r2 = Repeater(2, 4, StopIteration)
847 def run(r1, r2):
848 result = []
849 for i, j in zip_longest(r1, r2, fillvalue=0):
850 with support.captured_output('stdout'):
851 print((i, j))
852 result.append((i, j))
853 return result
854 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
855
856 # Formerly, the RuntimeError would be lost
857 # and StopIteration would stop as expected
858 r1 = Repeater(1, 3, RuntimeError)
859 r2 = Repeater(2, 4, StopIteration)
860 it = zip_longest(r1, r2, fillvalue=0)
861 self.assertEqual(next(it), (1, 2))
862 self.assertEqual(next(it), (1, 2))
863 self.assertEqual(next(it), (1, 2))
864 self.assertRaises(RuntimeError, next, it)
865
Christian Heimesc3f30c42008-02-22 16:37:40 +0000866 def test_product(self):
867 for args, result in [
Christian Heimes78644762008-03-04 23:39:23 +0000868 ([], [()]), # zero iterables
Christian Heimesc3f30c42008-02-22 16:37:40 +0000869 (['ab'], [('a',), ('b',)]), # one iterable
870 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
871 ([range(0), range(2), range(3)], []), # first iterable with zero length
872 ([range(2), range(0), range(3)], []), # middle iterable with zero length
873 ([range(2), range(3), range(0)], []), # last iterable with zero length
874 ]:
875 self.assertEqual(list(product(*args)), result)
Christian Heimesf16baeb2008-02-29 14:57:44 +0000876 for r in range(4):
877 self.assertEqual(list(product(*(args*r))),
878 list(product(*args, **dict(repeat=r))))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000879 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
880 self.assertRaises(TypeError, product, range(6), None)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000881
882 def product1(*args, **kwds):
883 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
884 n = len(pools)
885 if n == 0:
886 yield ()
887 return
888 if any(len(pool) == 0 for pool in pools):
889 return
890 indices = [0] * n
891 yield tuple(pool[i] for pool, i in zip(pools, indices))
892 while 1:
893 for i in reversed(range(n)): # right to left
894 if indices[i] == len(pools[i]) - 1:
895 continue
896 indices[i] += 1
897 for j in range(i+1, n):
898 indices[j] = 0
899 yield tuple(pool[i] for pool, i in zip(pools, indices))
900 break
901 else:
902 return
903
904 def product2(*args, **kwds):
905 'Pure python version used in docs'
906 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
907 result = [[]]
908 for pool in pools:
909 result = [x+[y] for x in result for y in pool]
910 for prod in result:
911 yield tuple(prod)
912
Christian Heimesc3f30c42008-02-22 16:37:40 +0000913 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
914 set('abcdefg'), range(11), tuple(range(13))]
915 for i in range(100):
916 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Christian Heimes78644762008-03-04 23:39:23 +0000917 expected_len = prod(map(len, args))
918 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000919 self.assertEqual(list(product(*args)), list(product1(*args)))
920 self.assertEqual(list(product(*args)), list(product2(*args)))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000921 args = map(iter, args)
Christian Heimes78644762008-03-04 23:39:23 +0000922 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000923
Alex Gaynore151d212011-07-17 16:21:30 -0700924 @support.impl_detail("tuple reuse is specific to CPython")
925 def test_product_tuple_reuse(self):
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000926 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
927 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000928
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000929 def test_product_pickling(self):
930 # check copy, deepcopy, pickle
931 for args, result in [
932 ([], [()]), # zero iterables
933 (['ab'], [('a',), ('b',)]), # one iterable
934 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
935 ([range(0), range(2), range(3)], []), # first iterable with zero length
936 ([range(2), range(0), range(3)], []), # middle iterable with zero length
937 ([range(2), range(3), range(0)], []), # last iterable with zero length
938 ]:
939 self.assertEqual(list(copy.copy(product(*args))), result)
940 self.assertEqual(list(copy.deepcopy(product(*args))), result)
941 self.pickletest(product(*args))
942
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000943 def test_repeat(self):
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +0000944 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
Guido van Rossum805365e2007-05-07 22:24:25 +0000945 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000946 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000947 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000948 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000949 self.assertEqual(list(repeat('a', 0)), [])
950 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000951 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000952 self.assertRaises(TypeError, repeat, None, 3, 4)
953 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000954 r = repeat(1+0j)
955 self.assertEqual(repr(r), 'repeat((1+0j))')
956 r = repeat(1+0j, 5)
957 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
958 list(r)
959 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000960
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000961 # check copy, deepcopy, pickle
962 c = repeat(object='a', times=10)
963 self.assertEqual(next(c), 'a')
964 self.assertEqual(take(2, copy.copy(c)), list('a' * 2))
965 self.assertEqual(take(2, copy.deepcopy(c)), list('a' * 2))
966 self.pickletest(repeat(object='a', times=10))
967
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000968 def test_map(self):
969 self.assertEqual(list(map(operator.pow, range(3), range(1,7))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000970 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000971 self.assertEqual(list(map(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000972 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000973 self.assertEqual(list(map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000974 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000975 self.assertEqual(take(2,map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000976 [('a',0),('b',1)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000977 self.assertEqual(list(map(operator.pow, [])), [])
978 self.assertRaises(TypeError, map)
979 self.assertRaises(TypeError, list, map(None, range(3), range(3)))
980 self.assertRaises(TypeError, map, operator.neg)
981 self.assertRaises(TypeError, next, map(10, range(5)))
982 self.assertRaises(ValueError, next, map(errfunc, [4], [5]))
983 self.assertRaises(TypeError, next, map(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000984
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000985 # check copy, deepcopy, pickle
986 ans = [('a',0),('b',1),('c',2)]
987
988 c = map(tupleize, 'abc', count())
989 self.assertEqual(list(copy.copy(c)), ans)
990
991 c = map(tupleize, 'abc', count())
992 self.assertEqual(list(copy.deepcopy(c)), ans)
993
994 c = map(tupleize, 'abc', count())
995 self.pickletest(c)
996
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000997 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000998 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
999 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001000 self.assertEqual(take(3, starmap(operator.pow, zip(count(), count(1)))),
Raymond Hettinger02420702003-06-29 20:36:23 +00001001 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001002 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +00001003 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
1004 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001005 self.assertRaises(TypeError, starmap)
1006 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001007 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
1008 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
1009 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001010
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001011 # check copy, deepcopy, pickle
1012 ans = [0**1, 1**2, 2**3]
1013
1014 c = starmap(operator.pow, zip(range(3), range(1,7)))
1015 self.assertEqual(list(copy.copy(c)), ans)
1016
1017 c = starmap(operator.pow, zip(range(3), range(1,7)))
1018 self.assertEqual(list(copy.deepcopy(c)), ans)
1019
1020 c = starmap(operator.pow, zip(range(3), range(1,7)))
1021 self.pickletest(c)
1022
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001023 def test_islice(self):
1024 for args in [ # islice(args) should agree with range(args)
1025 (10, 20, 3),
1026 (10, 3, 20),
1027 (10, 20),
1028 (10, 3),
1029 (20,)
1030 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +00001031 self.assertEqual(list(islice(range(100), *args)),
1032 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001033
1034 for args, tgtargs in [ # Stop when seqn is exhausted
1035 ((10, 110, 3), ((10, 100, 3))),
1036 ((10, 110), ((10, 100))),
1037 ((110,), (100,))
1038 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +00001039 self.assertEqual(list(islice(range(100), *args)),
1040 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001041
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001042 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +00001043 self.assertEqual(list(islice(range(10), None)), list(range(10)))
1044 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
1045 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
1046 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
1047 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001048
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +00001049 # Test number of items consumed SF #1171417
1050 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +00001051 self.assertEqual(list(islice(it, 3)), list(range(3)))
1052 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +00001053
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001054 # Test invalid arguments
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001055 ra = range(10)
1056 self.assertRaises(TypeError, islice, ra)
1057 self.assertRaises(TypeError, islice, ra, 1, 2, 3, 4)
1058 self.assertRaises(ValueError, islice, ra, -5, 10, 1)
1059 self.assertRaises(ValueError, islice, ra, 1, -5, -1)
1060 self.assertRaises(ValueError, islice, ra, 1, 10, -1)
1061 self.assertRaises(ValueError, islice, ra, 1, 10, 0)
1062 self.assertRaises(ValueError, islice, ra, 'a')
1063 self.assertRaises(ValueError, islice, ra, 'a', 1)
1064 self.assertRaises(ValueError, islice, ra, 1, 'a')
1065 self.assertRaises(ValueError, islice, ra, 'a', 1, 1)
1066 self.assertRaises(ValueError, islice, ra, 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +00001067 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001068
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001069 # Issue #10323: Less islice in a predictable state
1070 c = count()
1071 self.assertEqual(list(islice(c, 1, 3, 50)), [1])
1072 self.assertEqual(next(c), 3)
1073
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001074 # check copy, deepcopy, pickle
1075 for args in [ # islice(args) should agree with range(args)
1076 (10, 20, 3),
1077 (10, 3, 20),
1078 (10, 20),
1079 (10, 3),
1080 (20,)
1081 ]:
1082 self.assertEqual(list(copy.copy(islice(range(100), *args))),
1083 list(range(*args)))
1084 self.assertEqual(list(copy.deepcopy(islice(range(100), *args))),
1085 list(range(*args)))
1086 self.pickletest(islice(range(100), *args))
1087
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001088 def test_takewhile(self):
1089 data = [1, 3, 5, 20, 2, 4, 6, 8]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001090 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001091 self.assertEqual(list(takewhile(underten, [])), [])
1092 self.assertRaises(TypeError, takewhile)
1093 self.assertRaises(TypeError, takewhile, operator.pow)
1094 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001095 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
1096 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001097 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
1098 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +00001099 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001100
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001101 # check copy, deepcopy, pickle
1102 self.assertEqual(list(copy.copy(takewhile(underten, data))), [1, 3, 5])
1103 self.assertEqual(list(copy.deepcopy(takewhile(underten, data))),
1104 [1, 3, 5])
1105 self.pickletest(takewhile(underten, data))
1106
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001107 def test_dropwhile(self):
1108 data = [1, 3, 5, 20, 2, 4, 6, 8]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001109 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001110 self.assertEqual(list(dropwhile(underten, [])), [])
1111 self.assertRaises(TypeError, dropwhile)
1112 self.assertRaises(TypeError, dropwhile, operator.pow)
1113 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001114 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
1115 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001116
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001117 # check copy, deepcopy, pickle
1118 self.assertEqual(list(copy.copy(dropwhile(underten, data))), [20, 2, 4, 6, 8])
1119 self.assertEqual(list(copy.deepcopy(dropwhile(underten, data))),
1120 [20, 2, 4, 6, 8])
1121 self.pickletest(dropwhile(underten, data))
1122
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001123 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +00001124 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001125
1126 a, b = tee([]) # test empty iterator
1127 self.assertEqual(list(a), [])
1128 self.assertEqual(list(b), [])
1129
1130 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +00001131 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001132
1133 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +00001134 self.assertEqual(list(a), list(range(n)))
1135 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001136
1137 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +00001138 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001139 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001140 del a
Guido van Rossum805365e2007-05-07 22:24:25 +00001141 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001142
1143 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +00001144 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001145 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001146 del b
Guido van Rossum805365e2007-05-07 22:24:25 +00001147 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001148
Guido van Rossum805365e2007-05-07 22:24:25 +00001149 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001150 order = [0]*n + [1]*n
1151 random.shuffle(order)
1152 lists = ([], [])
1153 its = tee(irange(n))
1154 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +00001155 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001156 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +00001157 self.assertEqual(lists[0], list(range(n)))
1158 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001159
Raymond Hettingerad983e72003-11-12 14:32:26 +00001160 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001161 self.assertRaises(TypeError, tee)
1162 self.assertRaises(TypeError, tee, 3)
1163 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +00001164 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001165
Raymond Hettingerad983e72003-11-12 14:32:26 +00001166 # tee object should be instantiable
1167 a, b = tee('abc')
1168 c = type(a)('def')
1169 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +00001170
Raymond Hettingerad983e72003-11-12 14:32:26 +00001171 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +00001172 a, b, c = tee(range(2000), 3)
1173 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001174 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +00001175 self.assertEqual(list(b), list(range(2000)))
1176 self.assertEqual([next(c), next(c)], list(range(2)))
1177 self.assertEqual(list(a), list(range(100,2000)))
1178 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +00001179
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001180 # test values of n
1181 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +00001182 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +00001183 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001184 result = tee('abc', n)
1185 self.assertEqual(type(result), tuple)
1186 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001187 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001188
Raymond Hettingerad983e72003-11-12 14:32:26 +00001189 # tee pass-through to copyable iterator
1190 a, b = tee('abc')
1191 c, d = tee(a)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +00001192 self.assertTrue(a is c)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001193
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001194 # test tee_new
1195 t1, t2 = tee('abc')
1196 tnew = type(t1)
1197 self.assertRaises(TypeError, tnew)
1198 self.assertRaises(TypeError, tnew, 10)
1199 t3 = tnew(t1)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +00001200 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +00001201
Raymond Hettingera9f60922004-10-17 16:40:14 +00001202 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +00001203 a, b = tee(range(10))
Raymond Hettingera9f60922004-10-17 16:40:14 +00001204 p = proxy(a)
1205 self.assertEqual(getattr(p, '__class__'), type(b))
1206 del a
1207 self.assertRaises(ReferenceError, getattr, p, '__class__')
1208
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001209 ans = list('abc')
1210 long_ans = list(range(10000))
1211
1212 # check copy
1213 a, b = tee('abc')
1214 self.assertEqual(list(copy.copy(a)), ans)
1215 self.assertEqual(list(copy.copy(b)), ans)
1216 a, b = tee(list(range(10000)))
1217 self.assertEqual(list(copy.copy(a)), long_ans)
1218 self.assertEqual(list(copy.copy(b)), long_ans)
1219
1220 # check partially consumed copy
1221 a, b = tee('abc')
1222 take(2, a)
1223 take(1, b)
1224 self.assertEqual(list(copy.copy(a)), ans[2:])
1225 self.assertEqual(list(copy.copy(b)), ans[1:])
1226 self.assertEqual(list(a), ans[2:])
1227 self.assertEqual(list(b), ans[1:])
1228 a, b = tee(range(10000))
1229 take(100, a)
1230 take(60, b)
1231 self.assertEqual(list(copy.copy(a)), long_ans[100:])
1232 self.assertEqual(list(copy.copy(b)), long_ans[60:])
1233 self.assertEqual(list(a), long_ans[100:])
1234 self.assertEqual(list(b), long_ans[60:])
1235
1236 # check deepcopy
1237 a, b = tee('abc')
1238 self.assertEqual(list(copy.deepcopy(a)), ans)
1239 self.assertEqual(list(copy.deepcopy(b)), ans)
1240 self.assertEqual(list(a), ans)
1241 self.assertEqual(list(b), ans)
1242 a, b = tee(range(10000))
1243 self.assertEqual(list(copy.deepcopy(a)), long_ans)
1244 self.assertEqual(list(copy.deepcopy(b)), long_ans)
1245 self.assertEqual(list(a), long_ans)
1246 self.assertEqual(list(b), long_ans)
1247
1248 # check partially consumed deepcopy
1249 a, b = tee('abc')
1250 take(2, a)
1251 take(1, b)
1252 self.assertEqual(list(copy.deepcopy(a)), ans[2:])
1253 self.assertEqual(list(copy.deepcopy(b)), ans[1:])
1254 self.assertEqual(list(a), ans[2:])
1255 self.assertEqual(list(b), ans[1:])
1256 a, b = tee(range(10000))
1257 take(100, a)
1258 take(60, b)
1259 self.assertEqual(list(copy.deepcopy(a)), long_ans[100:])
1260 self.assertEqual(list(copy.deepcopy(b)), long_ans[60:])
1261 self.assertEqual(list(a), long_ans[100:])
1262 self.assertEqual(list(b), long_ans[60:])
1263
1264 # check pickle
1265 self.pickletest(iter(tee('abc')))
1266 a, b = tee('abc')
1267 self.pickletest(a, compare=ans)
1268 self.pickletest(b, compare=ans)
1269
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001270 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001271 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +00001272
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001273 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +00001274 self.assertRaises(StopIteration, next, f([]))
1275 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001276
Georg Brandla18af4e2007-04-21 15:47:16 +00001277 self.assertRaises(StopIteration, next, islice([], None))
1278 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001279
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001280 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +00001281 self.assertRaises(StopIteration, next, p)
1282 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001283 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +00001284 self.assertRaises(StopIteration, next, p)
1285 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001286
Georg Brandla18af4e2007-04-21 15:47:16 +00001287 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001288
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001289 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +00001290 self.assertRaises(StopIteration, next, f(lambda x:x, []))
1291 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001292
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001293class TestExamples(unittest.TestCase):
1294
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001295 def test_accumulate(self):
Raymond Hettinger482ba772010-12-01 22:48:00 +00001296 self.assertEqual(list(accumulate([1,2,3,4,5])), [1, 3, 6, 10, 15])
1297
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001298 def test_accumulate_reducible(self):
1299 # check copy, deepcopy, pickle
1300 data = [1, 2, 3, 4, 5]
1301 accumulated = [1, 3, 6, 10, 15]
1302 it = accumulate(data)
1303
1304 self.assertEqual(list(pickle.loads(pickle.dumps(it))), accumulated[:])
1305 self.assertEqual(next(it), 1)
1306 self.assertEqual(list(pickle.loads(pickle.dumps(it))), accumulated[1:])
1307 self.assertEqual(list(copy.deepcopy(it)), accumulated[1:])
1308 self.assertEqual(list(copy.copy(it)), accumulated[1:])
1309
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001310 def test_chain(self):
1311 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
1312
1313 def test_chain_from_iterable(self):
1314 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
1315
1316 def test_combinations(self):
1317 self.assertEqual(list(combinations('ABCD', 2)),
1318 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
1319 self.assertEqual(list(combinations(range(4), 3)),
1320 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
1321
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001322 def test_combinations_with_replacement(self):
1323 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
1324 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
1325
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001326 def test_compress(self):
1327 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
1328
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001329 def test_count(self):
1330 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
1331
1332 def test_cycle(self):
1333 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
1334
1335 def test_dropwhile(self):
1336 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
1337
1338 def test_groupby(self):
1339 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
1340 list('ABCDAB'))
1341 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
1342 [list('AAAA'), list('BBB'), list('CC'), list('D')])
1343
1344 def test_filter(self):
1345 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
1346
1347 def test_filterfalse(self):
1348 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
1349
1350 def test_map(self):
1351 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
1352
1353 def test_islice(self):
1354 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
1355 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
1356 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
1357 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1358
1359 def test_zip(self):
1360 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
1361
1362 def test_zip_longest(self):
1363 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
1364 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
1365
1366 def test_permutations(self):
1367 self.assertEqual(list(permutations('ABCD', 2)),
1368 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
1369 self.assertEqual(list(permutations(range(3))),
1370 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
1371
1372 def test_product(self):
1373 self.assertEqual(list(product('ABCD', 'xy')),
1374 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
1375 self.assertEqual(list(product(range(2), repeat=3)),
1376 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
1377 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
1378
1379 def test_repeat(self):
1380 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
1381
1382 def test_stapmap(self):
1383 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
1384 [32, 9, 1000])
1385
1386 def test_takewhile(self):
1387 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
1388
1389
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001390class TestGC(unittest.TestCase):
1391
1392 def makecycle(self, iterator, container):
1393 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +00001394 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001395 del container, iterator
1396
Raymond Hettinger482ba772010-12-01 22:48:00 +00001397 def test_accumulate(self):
1398 a = []
1399 self.makecycle(accumulate([1,2,a,3]), a)
1400
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001401 def test_chain(self):
1402 a = []
1403 self.makecycle(chain(a), a)
1404
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001405 def test_chain_from_iterable(self):
1406 a = []
1407 self.makecycle(chain.from_iterable([a]), a)
1408
1409 def test_combinations(self):
1410 a = []
1411 self.makecycle(combinations([1,2,a,3], 3), a)
1412
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001413 def test_combinations_with_replacement(self):
1414 a = []
1415 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1416
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001417 def test_compress(self):
1418 a = []
1419 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1420
Raymond Hettingerd6280f42009-02-16 20:50:56 +00001421 def test_count(self):
1422 a = []
1423 Int = type('Int', (int,), dict(x=a))
1424 self.makecycle(count(Int(0), Int(1)), a)
1425
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001426 def test_cycle(self):
1427 a = []
1428 self.makecycle(cycle([a]*2), a)
1429
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001430 def test_dropwhile(self):
1431 a = []
1432 self.makecycle(dropwhile(bool, [0, a, a]), a)
1433
1434 def test_groupby(self):
1435 a = []
1436 self.makecycle(groupby([a]*2, lambda x:x), a)
1437
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001438 def test_issue2246(self):
1439 # Issue 2246 -- the _grouper iterator was not included in GC
1440 n = 10
1441 keyfunc = lambda x: x
1442 for i, j in groupby(range(n), key=keyfunc):
1443 keyfunc.__dict__.setdefault('x',[]).append(j)
1444
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001445 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001446 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001447 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001448
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001449 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001450 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001451 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001452
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001453 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001454 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001455 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001456
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001457 def test_zip_longest(self):
1458 a = []
1459 self.makecycle(zip_longest([a]*2, [a]*3), a)
1460 b = [a, None]
1461 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1462
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001463 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001464 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001465 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001466
1467 def test_islice(self):
1468 a = []
1469 self.makecycle(islice([a]*2, None), a)
1470
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001471 def test_permutations(self):
1472 a = []
1473 self.makecycle(permutations([1,2,a,3], 3), a)
1474
1475 def test_product(self):
1476 a = []
1477 self.makecycle(product([1,2,a,3], repeat=3), a)
1478
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001479 def test_repeat(self):
1480 a = []
1481 self.makecycle(repeat(a), a)
1482
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001483 def test_starmap(self):
1484 a = []
1485 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1486
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001487 def test_takewhile(self):
1488 a = []
1489 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1490
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001491def R(seqn):
1492 'Regular generator'
1493 for i in seqn:
1494 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001495
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001496class G:
1497 'Sequence using __getitem__'
1498 def __init__(self, seqn):
1499 self.seqn = seqn
1500 def __getitem__(self, i):
1501 return self.seqn[i]
1502
1503class I:
1504 'Sequence using iterator protocol'
1505 def __init__(self, seqn):
1506 self.seqn = seqn
1507 self.i = 0
1508 def __iter__(self):
1509 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001510 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001511 if self.i >= len(self.seqn): raise StopIteration
1512 v = self.seqn[self.i]
1513 self.i += 1
1514 return v
1515
1516class Ig:
1517 'Sequence using iterator protocol defined with a generator'
1518 def __init__(self, seqn):
1519 self.seqn = seqn
1520 self.i = 0
1521 def __iter__(self):
1522 for val in self.seqn:
1523 yield val
1524
1525class X:
1526 'Missing __getitem__ and __iter__'
1527 def __init__(self, seqn):
1528 self.seqn = seqn
1529 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001530 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001531 if self.i >= len(self.seqn): raise StopIteration
1532 v = self.seqn[self.i]
1533 self.i += 1
1534 return v
1535
1536class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001537 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001538 def __init__(self, seqn):
1539 self.seqn = seqn
1540 self.i = 0
1541 def __iter__(self):
1542 return self
1543
1544class E:
1545 'Test propagation of exceptions'
1546 def __init__(self, seqn):
1547 self.seqn = seqn
1548 self.i = 0
1549 def __iter__(self):
1550 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001551 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001552 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001553
1554class S:
1555 'Test immediate stop'
1556 def __init__(self, seqn):
1557 pass
1558 def __iter__(self):
1559 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001560 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001561 raise StopIteration
1562
1563def L(seqn):
1564 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001565 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001566
1567
1568class TestVariousIteratorArgs(unittest.TestCase):
1569
Raymond Hettinger482ba772010-12-01 22:48:00 +00001570 def test_accumulate(self):
1571 s = [1,2,3,4,5]
1572 r = [1,3,6,10,15]
1573 n = len(s)
1574 for g in (G, I, Ig, L, R):
1575 self.assertEqual(list(accumulate(g(s))), r)
1576 self.assertEqual(list(accumulate(S(s))), [])
1577 self.assertRaises(TypeError, accumulate, X(s))
1578 self.assertRaises(TypeError, accumulate, N(s))
1579 self.assertRaises(ZeroDivisionError, list, accumulate(E(s)))
1580
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001581 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001582 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001583 for g in (G, I, Ig, S, L, R):
1584 self.assertEqual(list(chain(g(s))), list(g(s)))
1585 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001586 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001587 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001588 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1589
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001590 def test_compress(self):
1591 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1592 n = len(s)
1593 for g in (G, I, Ig, S, L, R):
1594 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1595 self.assertRaises(TypeError, compress, X(s), repeat(1))
1596 self.assertRaises(TypeError, compress, N(s), repeat(1))
1597 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1598
Christian Heimesc3f30c42008-02-22 16:37:40 +00001599 def test_product(self):
1600 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1601 self.assertRaises(TypeError, product, X(s))
1602 self.assertRaises(TypeError, product, N(s))
1603 self.assertRaises(ZeroDivisionError, product, E(s))
1604
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001605 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001606 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001607 for g in (G, I, Ig, S, L, R):
1608 tgtlen = len(s) * 3
1609 expected = list(g(s))*3
1610 actual = list(islice(cycle(g(s)), tgtlen))
1611 self.assertEqual(actual, expected)
1612 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001613 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001614 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1615
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001616 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001617 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001618 for g in (G, I, Ig, S, L, R):
1619 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1620 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001621 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001622 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1623
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001624 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001625 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001626 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001627 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001628 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001629 self.assertRaises(TypeError, filter, isEven, X(s))
1630 self.assertRaises(TypeError, filter, isEven, N(s))
1631 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001632
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001633 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001634 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001635 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001636 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001637 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001638 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1639 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1640 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001641
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001642 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001643 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001644 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001645 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1646 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1647 self.assertRaises(TypeError, zip, X(s))
1648 self.assertRaises(TypeError, zip, N(s))
1649 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001650
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001651 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001652 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001653 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001654 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1655 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1656 self.assertRaises(TypeError, zip_longest, X(s))
1657 self.assertRaises(TypeError, zip_longest, N(s))
1658 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001659
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001660 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001661 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001662 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001663 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001664 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001665 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001666 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001667 self.assertRaises(TypeError, map, onearg, X(s))
1668 self.assertRaises(TypeError, map, onearg, N(s))
1669 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001670
1671 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001672 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001673 for g in (G, I, Ig, S, L, R):
1674 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1675 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001676 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001677 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1678
1679 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001680 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001681 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001682 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001683 self.assertEqual(list(starmap(operator.pow, g(ss))),
1684 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001685 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001686 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001687 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1688
1689 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001690 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001691 for g in (G, I, Ig, S, L, R):
1692 tgt = []
1693 for elem in g(s):
1694 if not isEven(elem): break
1695 tgt.append(elem)
1696 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1697 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001698 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001699 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1700
1701 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001702 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001703 for g in (G, I, Ig, S, L, R):
1704 tgt = []
1705 for elem in g(s):
1706 if not tgt and isOdd(elem): continue
1707 tgt.append(elem)
1708 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1709 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001710 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001711 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1712
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001713 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001714 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001715 for g in (G, I, Ig, S, L, R):
1716 it1, it2 = tee(g(s))
1717 self.assertEqual(list(it1), list(g(s)))
1718 self.assertEqual(list(it2), list(g(s)))
1719 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001720 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001721 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1722
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001723class LengthTransparency(unittest.TestCase):
1724
1725 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001726 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001727 self.assertEqual(len(repeat(None, 50)), 50)
1728 self.assertRaises(TypeError, len, repeat(None))
1729
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001730class RegressionTests(unittest.TestCase):
1731
1732 def test_sf_793826(self):
1733 # Fix Armin Rigo's successful efforts to wreak havoc
1734
1735 def mutatingtuple(tuple1, f, tuple2):
1736 # this builds a tuple t which is a copy of tuple1,
1737 # then calls f(t), then mutates t to be equal to tuple2
1738 # (needs len(tuple1) == len(tuple2)).
1739 def g(value, first=[1]):
1740 if first:
1741 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00001742 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001743 return value
1744 items = list(tuple2)
1745 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001746 gen = map(g, items)
1747 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00001748 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001749
1750 def f(t):
1751 global T
1752 T = t
1753 first[:] = list(T)
1754
1755 first = []
1756 mutatingtuple((1,2,3), f, (4,5,6))
1757 second = list(T)
1758 self.assertEqual(first, second)
1759
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001760
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001761 def test_sf_950057(self):
1762 # Make sure that chain() and cycle() catch exceptions immediately
1763 # rather than when shifting between input sources
1764
1765 def gen1():
1766 hist.append(0)
1767 yield 1
1768 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001769 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001770 hist.append(2)
1771
1772 def gen2(x):
1773 hist.append(3)
1774 yield 2
1775 hist.append(4)
1776 if x:
1777 raise StopIteration
1778
1779 hist = []
1780 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1781 self.assertEqual(hist, [0,1])
1782
1783 hist = []
1784 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1785 self.assertEqual(hist, [0,1])
1786
1787 hist = []
1788 self.assertRaises(AssertionError, list, cycle(gen1()))
1789 self.assertEqual(hist, [0,1])
1790
Thomas Woutersb2137042007-02-01 18:02:27 +00001791class SubclassWithKwargsTest(unittest.TestCase):
1792 def test_keywords_in_subclass(self):
1793 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001794 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001795 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00001796 class Subclass(cls):
1797 def __init__(self, newarg=None, *args):
1798 cls.__init__(self, *args)
1799 try:
1800 Subclass(newarg=1)
1801 except TypeError as err:
1802 # we expect type errors because of wrong argument count
Ezio Melottib58e0bd2010-01-23 15:40:09 +00001803 self.assertNotIn("does not take keyword arguments", err.args[0])
Thomas Woutersb2137042007-02-01 18:02:27 +00001804
1805
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001806libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001807
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001808
1809>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001810>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00001811... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001812...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001813Check 1200 is for $120.15
1814Check 1201 is for $764.05
1815Check 1202 is for $823.14
1816
1817>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001818>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00001819... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00001820...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000018211
18228
182327
1824
1825>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001826>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00001827... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00001828...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001829Alex
1830Laura
1831Martin
1832Walter
1833Samuele
1834
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001835>>> from operator import itemgetter
1836>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00001837>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001838>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001839... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001840...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000018411 ['a', 'c', 'e']
18422 ['b', 'd', 'f']
18433 ['g']
1844
Raymond Hettinger734fb572004-01-20 20:04:40 +00001845# Find runs of consecutive numbers using groupby. The key to the solution
1846# is differencing with a range so that consecutive numbers all appear in
1847# same group.
1848>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00001849>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001850... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001851...
Raymond Hettinger734fb572004-01-20 20:04:40 +00001852[1]
1853[4, 5, 6]
1854[10]
1855[15, 16, 17, 18]
1856[22]
1857[25, 26, 27, 28]
1858
Georg Brandl3dbca812008-07-23 16:10:53 +00001859>>> def take(n, iterable):
1860... "Return first n items of the iterable as a list"
1861... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001862
Georg Brandl3dbca812008-07-23 16:10:53 +00001863>>> def enumerate(iterable, start=0):
1864... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001865
Georg Brandl3dbca812008-07-23 16:10:53 +00001866>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001867... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00001868... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001869
Raymond Hettingercdf8ba32009-02-19 04:45:07 +00001870>>> def nth(iterable, n, default=None):
1871... "Returns the nth item or a default value"
1872... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001873
Georg Brandl3dbca812008-07-23 16:10:53 +00001874>>> def quantify(iterable, pred=bool):
1875... "Count how many times the predicate is true"
1876... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001877
Georg Brandl3dbca812008-07-23 16:10:53 +00001878>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001879... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00001880... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001881
Georg Brandl3dbca812008-07-23 16:10:53 +00001882>>> def ncycles(iterable, n):
Ezio Melotti13925002011-03-16 11:05:33 +02001883... "Returns the sequence elements n times"
Georg Brandl3dbca812008-07-23 16:10:53 +00001884... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001885
1886>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001887... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001888
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001889>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001890... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001891
1892>>> def repeatfunc(func, times=None, *args):
1893... "Repeat calls to func with specified arguments."
1894... " Example: repeatfunc(random.random)"
1895... if times is None:
1896... return starmap(func, repeat(args))
1897... else:
1898... return starmap(func, repeat(args, times))
1899
Raymond Hettingerd591f662003-10-26 15:34:50 +00001900>>> def pairwise(iterable):
1901... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1902... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001903... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00001904... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001905... except StopIteration:
1906... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001907... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001908
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001909>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001910... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001911... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00001912... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001913
1914>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001915... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001916... # Recipe credited to George Sakkis
1917... pending = len(iterables)
1918... nexts = cycle(iter(it).__next__ for it in iterables)
1919... while pending:
1920... try:
1921... for next in nexts:
1922... yield next()
1923... except StopIteration:
1924... pending -= 1
1925... nexts = cycle(islice(nexts, pending))
1926
1927>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00001928... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1929... s = list(iterable)
1930... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001931
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001932>>> def unique_everseen(iterable, key=None):
1933... "List unique elements, preserving order. Remember all elements ever seen."
1934... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1935... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1936... seen = set()
1937... seen_add = seen.add
1938... if key is None:
1939... for element in iterable:
1940... if element not in seen:
1941... seen_add(element)
1942... yield element
1943... else:
1944... for element in iterable:
1945... k = key(element)
1946... if k not in seen:
1947... seen_add(k)
1948... yield element
1949
1950>>> def unique_justseen(iterable, key=None):
1951... "List unique elements, preserving order. Remember only the element just seen."
1952... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1953... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1954... return map(next, map(itemgetter(1), groupby(iterable, key)))
1955
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001956This is not part of the examples but it tests to make sure the definitions
1957perform as purported.
1958
Raymond Hettingera098b332003-09-08 23:58:40 +00001959>>> take(10, count())
1960[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1961
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001962>>> list(enumerate('abc'))
1963[(0, 'a'), (1, 'b'), (2, 'c')]
1964
1965>>> list(islice(tabulate(lambda x: 2*x), 4))
1966[0, 2, 4, 6]
1967
1968>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00001969'd'
1970
1971>>> nth('abcde', 9) is None
1972True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001973
Guido van Rossum805365e2007-05-07 22:24:25 +00001974>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000197550
1976
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001977>>> a = [[1, 2, 3], [4, 5, 6]]
1978>>> flatten(a)
1979[1, 2, 3, 4, 5, 6]
1980
1981>>> list(repeatfunc(pow, 5, 2, 3))
1982[8, 8, 8, 8, 8]
1983
1984>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001985>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001986[0, 0, 0, 0, 0]
1987
Raymond Hettingerd591f662003-10-26 15:34:50 +00001988>>> list(pairwise('abcd'))
1989[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001990
Raymond Hettingerd591f662003-10-26 15:34:50 +00001991>>> list(pairwise([]))
1992[]
1993
1994>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001995[]
1996
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001997>>> list(islice(padnone('abc'), 0, 6))
1998['a', 'b', 'c', None, None, None]
1999
2000>>> list(ncycles('abc', 3))
2001['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
2002
2003>>> dotproduct([1,2,3], [4,5,6])
200432
2005
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002006>>> list(grouper(3, 'abcdefg', 'x'))
2007[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
2008
2009>>> list(roundrobin('abc', 'd', 'ef'))
2010['a', 'd', 'e', 'b', 'f', 'c']
2011
Raymond Hettingerace67332009-01-26 02:23:50 +00002012>>> list(powerset([1,2,3]))
2013[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002014
Raymond Hettinger191e8502009-01-27 13:29:43 +00002015>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
2016True
2017
2018>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
2019True
2020
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00002021>>> list(unique_everseen('AAAABBBCCDAABBB'))
2022['A', 'B', 'C', 'D']
2023
2024>>> list(unique_everseen('ABBCcAD', str.lower))
2025['A', 'B', 'C', 'D']
2026
2027>>> list(unique_justseen('AAAABBBCCDAABBB'))
2028['A', 'B', 'C', 'D', 'A', 'B']
2029
2030>>> list(unique_justseen('ABBCcAD', str.lower))
2031['A', 'B', 'C', 'A', 'D']
2032
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002033"""
2034
2035__test__ = {'libreftest' : libreftest}
2036
2037def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002038 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00002039 RegressionTests, LengthTransparency,
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002040 SubclassWithKwargsTest, TestExamples)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002041 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002042
2043 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002044 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00002045 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00002046 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00002047 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002048 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00002049 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00002050 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00002051 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002052
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002053 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002054 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002055
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002056if __name__ == "__main__":
2057 test_main(verbose=True)