blob: 53926a9df5cd0c13e2a0eb14d3f38b8aaf7b1f94 [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
Serhiy Storchakaa3e91282013-01-25 13:19:31 +02001270 # Issue 13454: Crash when deleting backward iterator from tee()
1271 def test_tee_del_backward(self):
Serhiy Storchaka5bb893c2013-01-26 11:52:06 +02001272 forward, backward = tee(repeat(None, 20000000))
1273 any(forward) # exhaust the iterator
Serhiy Storchakaa3e91282013-01-25 13:19:31 +02001274 del backward
1275
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001276 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001277 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +00001278
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001279 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +00001280 self.assertRaises(StopIteration, next, f([]))
1281 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001282
Georg Brandla18af4e2007-04-21 15:47:16 +00001283 self.assertRaises(StopIteration, next, islice([], None))
1284 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001285
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001286 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +00001287 self.assertRaises(StopIteration, next, p)
1288 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001289 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +00001290 self.assertRaises(StopIteration, next, p)
1291 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001292
Georg Brandla18af4e2007-04-21 15:47:16 +00001293 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001294
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001295 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +00001296 self.assertRaises(StopIteration, next, f(lambda x:x, []))
1297 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001298
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001299class TestExamples(unittest.TestCase):
1300
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001301 def test_accumulate(self):
Raymond Hettinger482ba772010-12-01 22:48:00 +00001302 self.assertEqual(list(accumulate([1,2,3,4,5])), [1, 3, 6, 10, 15])
1303
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001304 def test_accumulate_reducible(self):
1305 # check copy, deepcopy, pickle
1306 data = [1, 2, 3, 4, 5]
1307 accumulated = [1, 3, 6, 10, 15]
1308 it = accumulate(data)
1309
1310 self.assertEqual(list(pickle.loads(pickle.dumps(it))), accumulated[:])
1311 self.assertEqual(next(it), 1)
1312 self.assertEqual(list(pickle.loads(pickle.dumps(it))), accumulated[1:])
1313 self.assertEqual(list(copy.deepcopy(it)), accumulated[1:])
1314 self.assertEqual(list(copy.copy(it)), accumulated[1:])
1315
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001316 def test_chain(self):
1317 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
1318
1319 def test_chain_from_iterable(self):
1320 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
1321
1322 def test_combinations(self):
1323 self.assertEqual(list(combinations('ABCD', 2)),
1324 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
1325 self.assertEqual(list(combinations(range(4), 3)),
1326 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
1327
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001328 def test_combinations_with_replacement(self):
1329 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
1330 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
1331
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001332 def test_compress(self):
1333 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
1334
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001335 def test_count(self):
1336 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
1337
1338 def test_cycle(self):
1339 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
1340
1341 def test_dropwhile(self):
1342 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
1343
1344 def test_groupby(self):
1345 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
1346 list('ABCDAB'))
1347 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
1348 [list('AAAA'), list('BBB'), list('CC'), list('D')])
1349
1350 def test_filter(self):
1351 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
1352
1353 def test_filterfalse(self):
1354 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
1355
1356 def test_map(self):
1357 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
1358
1359 def test_islice(self):
1360 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
1361 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
1362 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
1363 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1364
1365 def test_zip(self):
1366 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
1367
1368 def test_zip_longest(self):
1369 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
1370 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
1371
1372 def test_permutations(self):
1373 self.assertEqual(list(permutations('ABCD', 2)),
1374 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
1375 self.assertEqual(list(permutations(range(3))),
1376 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
1377
1378 def test_product(self):
1379 self.assertEqual(list(product('ABCD', 'xy')),
1380 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
1381 self.assertEqual(list(product(range(2), repeat=3)),
1382 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
1383 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
1384
1385 def test_repeat(self):
1386 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
1387
1388 def test_stapmap(self):
1389 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
1390 [32, 9, 1000])
1391
1392 def test_takewhile(self):
1393 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
1394
1395
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001396class TestGC(unittest.TestCase):
1397
1398 def makecycle(self, iterator, container):
1399 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +00001400 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001401 del container, iterator
1402
Raymond Hettinger482ba772010-12-01 22:48:00 +00001403 def test_accumulate(self):
1404 a = []
1405 self.makecycle(accumulate([1,2,a,3]), a)
1406
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001407 def test_chain(self):
1408 a = []
1409 self.makecycle(chain(a), a)
1410
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001411 def test_chain_from_iterable(self):
1412 a = []
1413 self.makecycle(chain.from_iterable([a]), a)
1414
1415 def test_combinations(self):
1416 a = []
1417 self.makecycle(combinations([1,2,a,3], 3), a)
1418
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001419 def test_combinations_with_replacement(self):
1420 a = []
1421 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1422
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001423 def test_compress(self):
1424 a = []
1425 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1426
Raymond Hettingerd6280f42009-02-16 20:50:56 +00001427 def test_count(self):
1428 a = []
1429 Int = type('Int', (int,), dict(x=a))
1430 self.makecycle(count(Int(0), Int(1)), a)
1431
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001432 def test_cycle(self):
1433 a = []
1434 self.makecycle(cycle([a]*2), a)
1435
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001436 def test_dropwhile(self):
1437 a = []
1438 self.makecycle(dropwhile(bool, [0, a, a]), a)
1439
1440 def test_groupby(self):
1441 a = []
1442 self.makecycle(groupby([a]*2, lambda x:x), a)
1443
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001444 def test_issue2246(self):
1445 # Issue 2246 -- the _grouper iterator was not included in GC
1446 n = 10
1447 keyfunc = lambda x: x
1448 for i, j in groupby(range(n), key=keyfunc):
1449 keyfunc.__dict__.setdefault('x',[]).append(j)
1450
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001451 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001452 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001453 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001454
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001455 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001456 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001457 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001458
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001459 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001460 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001461 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001462
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001463 def test_zip_longest(self):
1464 a = []
1465 self.makecycle(zip_longest([a]*2, [a]*3), a)
1466 b = [a, None]
1467 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1468
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001469 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001470 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001471 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001472
1473 def test_islice(self):
1474 a = []
1475 self.makecycle(islice([a]*2, None), a)
1476
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001477 def test_permutations(self):
1478 a = []
1479 self.makecycle(permutations([1,2,a,3], 3), a)
1480
1481 def test_product(self):
1482 a = []
1483 self.makecycle(product([1,2,a,3], repeat=3), a)
1484
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001485 def test_repeat(self):
1486 a = []
1487 self.makecycle(repeat(a), a)
1488
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001489 def test_starmap(self):
1490 a = []
1491 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1492
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001493 def test_takewhile(self):
1494 a = []
1495 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1496
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001497def R(seqn):
1498 'Regular generator'
1499 for i in seqn:
1500 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001501
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001502class G:
1503 'Sequence using __getitem__'
1504 def __init__(self, seqn):
1505 self.seqn = seqn
1506 def __getitem__(self, i):
1507 return self.seqn[i]
1508
1509class I:
1510 'Sequence using iterator protocol'
1511 def __init__(self, seqn):
1512 self.seqn = seqn
1513 self.i = 0
1514 def __iter__(self):
1515 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001516 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001517 if self.i >= len(self.seqn): raise StopIteration
1518 v = self.seqn[self.i]
1519 self.i += 1
1520 return v
1521
1522class Ig:
1523 'Sequence using iterator protocol defined with a generator'
1524 def __init__(self, seqn):
1525 self.seqn = seqn
1526 self.i = 0
1527 def __iter__(self):
1528 for val in self.seqn:
1529 yield val
1530
1531class X:
1532 'Missing __getitem__ and __iter__'
1533 def __init__(self, seqn):
1534 self.seqn = seqn
1535 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001536 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001537 if self.i >= len(self.seqn): raise StopIteration
1538 v = self.seqn[self.i]
1539 self.i += 1
1540 return v
1541
1542class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001543 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001544 def __init__(self, seqn):
1545 self.seqn = seqn
1546 self.i = 0
1547 def __iter__(self):
1548 return self
1549
1550class E:
1551 'Test propagation of exceptions'
1552 def __init__(self, seqn):
1553 self.seqn = seqn
1554 self.i = 0
1555 def __iter__(self):
1556 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001557 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001558 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001559
1560class S:
1561 'Test immediate stop'
1562 def __init__(self, seqn):
1563 pass
1564 def __iter__(self):
1565 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001566 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001567 raise StopIteration
1568
1569def L(seqn):
1570 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001571 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001572
1573
1574class TestVariousIteratorArgs(unittest.TestCase):
1575
Raymond Hettinger482ba772010-12-01 22:48:00 +00001576 def test_accumulate(self):
1577 s = [1,2,3,4,5]
1578 r = [1,3,6,10,15]
1579 n = len(s)
1580 for g in (G, I, Ig, L, R):
1581 self.assertEqual(list(accumulate(g(s))), r)
1582 self.assertEqual(list(accumulate(S(s))), [])
1583 self.assertRaises(TypeError, accumulate, X(s))
1584 self.assertRaises(TypeError, accumulate, N(s))
1585 self.assertRaises(ZeroDivisionError, list, accumulate(E(s)))
1586
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001587 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001588 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001589 for g in (G, I, Ig, S, L, R):
1590 self.assertEqual(list(chain(g(s))), list(g(s)))
1591 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001592 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001593 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001594 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1595
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001596 def test_compress(self):
1597 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1598 n = len(s)
1599 for g in (G, I, Ig, S, L, R):
1600 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1601 self.assertRaises(TypeError, compress, X(s), repeat(1))
1602 self.assertRaises(TypeError, compress, N(s), repeat(1))
1603 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1604
Christian Heimesc3f30c42008-02-22 16:37:40 +00001605 def test_product(self):
1606 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1607 self.assertRaises(TypeError, product, X(s))
1608 self.assertRaises(TypeError, product, N(s))
1609 self.assertRaises(ZeroDivisionError, product, E(s))
1610
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001611 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001612 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001613 for g in (G, I, Ig, S, L, R):
1614 tgtlen = len(s) * 3
1615 expected = list(g(s))*3
1616 actual = list(islice(cycle(g(s)), tgtlen))
1617 self.assertEqual(actual, expected)
1618 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001619 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001620 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1621
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001622 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001623 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001624 for g in (G, I, Ig, S, L, R):
1625 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1626 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001627 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001628 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1629
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001630 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001631 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001632 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001633 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001634 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001635 self.assertRaises(TypeError, filter, isEven, X(s))
1636 self.assertRaises(TypeError, filter, isEven, N(s))
1637 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001638
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001639 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001640 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001641 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001642 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001643 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001644 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1645 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1646 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001647
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001648 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001649 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001650 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001651 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1652 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1653 self.assertRaises(TypeError, zip, X(s))
1654 self.assertRaises(TypeError, zip, N(s))
1655 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001656
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001657 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001658 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001659 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001660 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1661 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1662 self.assertRaises(TypeError, zip_longest, X(s))
1663 self.assertRaises(TypeError, zip_longest, N(s))
1664 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001665
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001666 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001667 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001668 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001669 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001670 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001671 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001672 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001673 self.assertRaises(TypeError, map, onearg, X(s))
1674 self.assertRaises(TypeError, map, onearg, N(s))
1675 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001676
1677 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001678 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001679 for g in (G, I, Ig, S, L, R):
1680 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1681 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001682 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001683 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1684
1685 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001686 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001687 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001688 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001689 self.assertEqual(list(starmap(operator.pow, g(ss))),
1690 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001691 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001692 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001693 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1694
1695 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001696 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001697 for g in (G, I, Ig, S, L, R):
1698 tgt = []
1699 for elem in g(s):
1700 if not isEven(elem): break
1701 tgt.append(elem)
1702 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1703 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001704 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001705 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1706
1707 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001708 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001709 for g in (G, I, Ig, S, L, R):
1710 tgt = []
1711 for elem in g(s):
1712 if not tgt and isOdd(elem): continue
1713 tgt.append(elem)
1714 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1715 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001716 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001717 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1718
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001719 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001720 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001721 for g in (G, I, Ig, S, L, R):
1722 it1, it2 = tee(g(s))
1723 self.assertEqual(list(it1), list(g(s)))
1724 self.assertEqual(list(it2), list(g(s)))
1725 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001726 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001727 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1728
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001729class LengthTransparency(unittest.TestCase):
1730
1731 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001732 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001733 self.assertEqual(len(repeat(None, 50)), 50)
1734 self.assertRaises(TypeError, len, repeat(None))
1735
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001736class RegressionTests(unittest.TestCase):
1737
1738 def test_sf_793826(self):
1739 # Fix Armin Rigo's successful efforts to wreak havoc
1740
1741 def mutatingtuple(tuple1, f, tuple2):
1742 # this builds a tuple t which is a copy of tuple1,
1743 # then calls f(t), then mutates t to be equal to tuple2
1744 # (needs len(tuple1) == len(tuple2)).
1745 def g(value, first=[1]):
1746 if first:
1747 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00001748 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001749 return value
1750 items = list(tuple2)
1751 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001752 gen = map(g, items)
1753 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00001754 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001755
1756 def f(t):
1757 global T
1758 T = t
1759 first[:] = list(T)
1760
1761 first = []
1762 mutatingtuple((1,2,3), f, (4,5,6))
1763 second = list(T)
1764 self.assertEqual(first, second)
1765
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001766
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001767 def test_sf_950057(self):
1768 # Make sure that chain() and cycle() catch exceptions immediately
1769 # rather than when shifting between input sources
1770
1771 def gen1():
1772 hist.append(0)
1773 yield 1
1774 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001775 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001776 hist.append(2)
1777
1778 def gen2(x):
1779 hist.append(3)
1780 yield 2
1781 hist.append(4)
1782 if x:
1783 raise StopIteration
1784
1785 hist = []
1786 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1787 self.assertEqual(hist, [0,1])
1788
1789 hist = []
1790 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1791 self.assertEqual(hist, [0,1])
1792
1793 hist = []
1794 self.assertRaises(AssertionError, list, cycle(gen1()))
1795 self.assertEqual(hist, [0,1])
1796
Thomas Woutersb2137042007-02-01 18:02:27 +00001797class SubclassWithKwargsTest(unittest.TestCase):
1798 def test_keywords_in_subclass(self):
1799 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001800 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001801 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00001802 class Subclass(cls):
1803 def __init__(self, newarg=None, *args):
1804 cls.__init__(self, *args)
1805 try:
1806 Subclass(newarg=1)
1807 except TypeError as err:
1808 # we expect type errors because of wrong argument count
Ezio Melottib58e0bd2010-01-23 15:40:09 +00001809 self.assertNotIn("does not take keyword arguments", err.args[0])
Thomas Woutersb2137042007-02-01 18:02:27 +00001810
1811
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001812libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001813
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001814
1815>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001816>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00001817... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001818...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001819Check 1200 is for $120.15
1820Check 1201 is for $764.05
1821Check 1202 is for $823.14
1822
1823>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001824>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00001825... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00001826...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000018271
18288
182927
1830
1831>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001832>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00001833... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00001834...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001835Alex
1836Laura
1837Martin
1838Walter
1839Samuele
1840
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001841>>> from operator import itemgetter
1842>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00001843>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001844>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001845... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001846...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000018471 ['a', 'c', 'e']
18482 ['b', 'd', 'f']
18493 ['g']
1850
Raymond Hettinger734fb572004-01-20 20:04:40 +00001851# Find runs of consecutive numbers using groupby. The key to the solution
1852# is differencing with a range so that consecutive numbers all appear in
1853# same group.
1854>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00001855>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001856... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001857...
Raymond Hettinger734fb572004-01-20 20:04:40 +00001858[1]
1859[4, 5, 6]
1860[10]
1861[15, 16, 17, 18]
1862[22]
1863[25, 26, 27, 28]
1864
Georg Brandl3dbca812008-07-23 16:10:53 +00001865>>> def take(n, iterable):
1866... "Return first n items of the iterable as a list"
1867... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001868
Georg Brandl3dbca812008-07-23 16:10:53 +00001869>>> def enumerate(iterable, start=0):
1870... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001871
Georg Brandl3dbca812008-07-23 16:10:53 +00001872>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001873... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00001874... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001875
Raymond Hettingercdf8ba32009-02-19 04:45:07 +00001876>>> def nth(iterable, n, default=None):
1877... "Returns the nth item or a default value"
1878... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001879
Georg Brandl3dbca812008-07-23 16:10:53 +00001880>>> def quantify(iterable, pred=bool):
1881... "Count how many times the predicate is true"
1882... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001883
Georg Brandl3dbca812008-07-23 16:10:53 +00001884>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001885... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00001886... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001887
Georg Brandl3dbca812008-07-23 16:10:53 +00001888>>> def ncycles(iterable, n):
Ezio Melotti13925002011-03-16 11:05:33 +02001889... "Returns the sequence elements n times"
Georg Brandl3dbca812008-07-23 16:10:53 +00001890... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001891
1892>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001893... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001894
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001895>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001896... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001897
1898>>> def repeatfunc(func, times=None, *args):
1899... "Repeat calls to func with specified arguments."
1900... " Example: repeatfunc(random.random)"
1901... if times is None:
1902... return starmap(func, repeat(args))
1903... else:
1904... return starmap(func, repeat(args, times))
1905
Raymond Hettingerd591f662003-10-26 15:34:50 +00001906>>> def pairwise(iterable):
1907... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1908... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001909... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00001910... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001911... except StopIteration:
1912... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001913... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001914
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001915>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001916... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001917... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00001918... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001919
1920>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001921... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001922... # Recipe credited to George Sakkis
1923... pending = len(iterables)
1924... nexts = cycle(iter(it).__next__ for it in iterables)
1925... while pending:
1926... try:
1927... for next in nexts:
1928... yield next()
1929... except StopIteration:
1930... pending -= 1
1931... nexts = cycle(islice(nexts, pending))
1932
1933>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00001934... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1935... s = list(iterable)
1936... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001937
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001938>>> def unique_everseen(iterable, key=None):
1939... "List unique elements, preserving order. Remember all elements ever seen."
1940... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1941... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1942... seen = set()
1943... seen_add = seen.add
1944... if key is None:
1945... for element in iterable:
1946... if element not in seen:
1947... seen_add(element)
1948... yield element
1949... else:
1950... for element in iterable:
1951... k = key(element)
1952... if k not in seen:
1953... seen_add(k)
1954... yield element
1955
1956>>> def unique_justseen(iterable, key=None):
1957... "List unique elements, preserving order. Remember only the element just seen."
1958... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1959... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1960... return map(next, map(itemgetter(1), groupby(iterable, key)))
1961
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001962This is not part of the examples but it tests to make sure the definitions
1963perform as purported.
1964
Raymond Hettingera098b332003-09-08 23:58:40 +00001965>>> take(10, count())
1966[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1967
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001968>>> list(enumerate('abc'))
1969[(0, 'a'), (1, 'b'), (2, 'c')]
1970
1971>>> list(islice(tabulate(lambda x: 2*x), 4))
1972[0, 2, 4, 6]
1973
1974>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00001975'd'
1976
1977>>> nth('abcde', 9) is None
1978True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001979
Guido van Rossum805365e2007-05-07 22:24:25 +00001980>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000198150
1982
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001983>>> a = [[1, 2, 3], [4, 5, 6]]
1984>>> flatten(a)
1985[1, 2, 3, 4, 5, 6]
1986
1987>>> list(repeatfunc(pow, 5, 2, 3))
1988[8, 8, 8, 8, 8]
1989
1990>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001991>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001992[0, 0, 0, 0, 0]
1993
Raymond Hettingerd591f662003-10-26 15:34:50 +00001994>>> list(pairwise('abcd'))
1995[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001996
Raymond Hettingerd591f662003-10-26 15:34:50 +00001997>>> list(pairwise([]))
1998[]
1999
2000>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00002001[]
2002
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002003>>> list(islice(padnone('abc'), 0, 6))
2004['a', 'b', 'c', None, None, None]
2005
2006>>> list(ncycles('abc', 3))
2007['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
2008
2009>>> dotproduct([1,2,3], [4,5,6])
201032
2011
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002012>>> list(grouper(3, 'abcdefg', 'x'))
2013[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
2014
2015>>> list(roundrobin('abc', 'd', 'ef'))
2016['a', 'd', 'e', 'b', 'f', 'c']
2017
Raymond Hettingerace67332009-01-26 02:23:50 +00002018>>> list(powerset([1,2,3]))
2019[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002020
Raymond Hettinger191e8502009-01-27 13:29:43 +00002021>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
2022True
2023
2024>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
2025True
2026
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00002027>>> list(unique_everseen('AAAABBBCCDAABBB'))
2028['A', 'B', 'C', 'D']
2029
2030>>> list(unique_everseen('ABBCcAD', str.lower))
2031['A', 'B', 'C', 'D']
2032
2033>>> list(unique_justseen('AAAABBBCCDAABBB'))
2034['A', 'B', 'C', 'D', 'A', 'B']
2035
2036>>> list(unique_justseen('ABBCcAD', str.lower))
2037['A', 'B', 'C', 'A', 'D']
2038
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002039"""
2040
2041__test__ = {'libreftest' : libreftest}
2042
2043def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002044 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00002045 RegressionTests, LengthTransparency,
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03002046 SubclassWithKwargsTest, TestExamples)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002047 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002048
2049 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002050 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00002051 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00002052 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00002053 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002054 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00002055 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00002056 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00002057 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002058
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002059 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002060 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002061
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002062if __name__ == "__main__":
2063 test_main(verbose=True)