blob: ea1f57caade12b91bf7c51f7eed93e33d9d76267 [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 *
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02004import weakref
Raymond Hettinger8a882a72009-02-12 12:08:18 +00005from decimal import Decimal
Raymond Hettingerde09acf2009-02-12 12:53:53 +00006from fractions import Fraction
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00007import operator
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00008import random
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00009import copy
10import pickle
Christian Heimesc3f30c42008-02-22 16:37:40 +000011from functools import reduce
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +020012import sys
13import struct
Benjamin Petersonee8712c2008-05-20 21:35:26 +000014maxsize = support.MAX_Py_ssize_t
Guido van Rossum360e4b82007-05-14 22:51:27 +000015minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000016
Guido van Rossum801f0d72006-08-24 19:48:10 +000017def lzip(*args):
18 return list(zip(*args))
19
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000020def onearg(x):
21 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000022 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000023
24def errfunc(*args):
25 'Test function that raises an error'
26 raise ValueError
27
28def gen3():
29 'Non-restartable source sequence'
30 for i in (0, 1, 2):
31 yield i
32
33def isEven(x):
34 'Test predicate'
35 return x%2==0
36
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000037def isOdd(x):
38 'Test predicate'
39 return x%2==1
40
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000041def tupleize(*args):
42 return args
43
44def irange(n):
45 for i in range(n):
46 yield i
47
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000048class StopNow:
49 'Class emulating an empty iterable.'
50 def __iter__(self):
51 return self
Georg Brandla18af4e2007-04-21 15:47:16 +000052 def __next__(self):
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000053 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000054
Raymond Hettinger02420702003-06-29 20:36:23 +000055def take(n, seq):
56 'Convenience function for partially consuming a long of infinite iterable'
57 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000058
Christian Heimes78644762008-03-04 23:39:23 +000059def prod(iterable):
60 return reduce(operator.mul, iterable, 1)
61
Christian Heimes380f7f22008-02-28 11:19:05 +000062def fact(n):
63 'Factorial'
Christian Heimes78644762008-03-04 23:39:23 +000064 return prod(range(1, n+1))
65
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000066# root level methods for pickling ability
67def testR(r):
68 return r[0]
69
70def testR2(r):
71 return r[2]
72
73def underten(x):
74 return x<10
75
Serhiy Storchakabad12572014-12-15 14:03:42 +020076picklecopiers = [lambda s, proto=proto: pickle.loads(pickle.dumps(s, proto))
77 for proto in range(pickle.HIGHEST_PROTOCOL + 1)]
78
Raymond Hettinger96ef8112003-02-01 00:10:11 +000079class TestBasicOps(unittest.TestCase):
Raymond Hettinger482ba772010-12-01 22:48:00 +000080
Serhiy Storchakabad12572014-12-15 14:03:42 +020081 def pickletest(self, protocol, it, stop=4, take=1, compare=None):
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000082 """Test that an iterator is the same after pickling, also when part-consumed"""
83 def expand(it, i=0):
84 # Recursively expand iterables, within sensible bounds
85 if i > 10:
86 raise RuntimeError("infinite recursion encountered")
87 if isinstance(it, str):
88 return it
89 try:
90 l = list(islice(it, stop))
91 except TypeError:
92 return it # can't expand it
93 return [expand(e, i+1) for e in l]
94
95 # Test the initial copy against the original
Serhiy Storchakabad12572014-12-15 14:03:42 +020096 dump = pickle.dumps(it, protocol)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000097 i2 = pickle.loads(dump)
98 self.assertEqual(type(it), type(i2))
99 a, b = expand(it), expand(i2)
100 self.assertEqual(a, b)
101 if compare:
102 c = expand(compare)
103 self.assertEqual(a, c)
104
105 # Take from the copy, and create another copy and compare them.
106 i3 = pickle.loads(dump)
107 took = 0
108 try:
109 for i in range(take):
110 next(i3)
111 took += 1
112 except StopIteration:
113 pass #in case there is less data than 'take'
Serhiy Storchakabad12572014-12-15 14:03:42 +0200114 dump = pickle.dumps(i3, protocol)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000115 i4 = pickle.loads(dump)
116 a, b = expand(i3), expand(i4)
117 self.assertEqual(a, b)
118 if compare:
119 c = expand(compare[took:])
120 self.assertEqual(a, c);
121
Raymond Hettinger482ba772010-12-01 22:48:00 +0000122 def test_accumulate(self):
123 self.assertEqual(list(accumulate(range(10))), # one positional arg
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000124 [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])
125 self.assertEqual(list(accumulate(iterable=range(10))), # kw arg
126 [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])
Raymond Hettinger482ba772010-12-01 22:48:00 +0000127 for typ in int, complex, Decimal, Fraction: # multiple types
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000128 self.assertEqual(
129 list(accumulate(map(typ, range(10)))),
Raymond Hettinger482ba772010-12-01 22:48:00 +0000130 list(map(typ, [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])))
Raymond Hettinger2d93e6e2010-12-03 02:33:53 +0000131 self.assertEqual(list(accumulate('abc')), ['a', 'ab', 'abc']) # works with non-numeric
Raymond Hettinger482ba772010-12-01 22:48:00 +0000132 self.assertEqual(list(accumulate([])), []) # empty iterable
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000133 self.assertEqual(list(accumulate([7])), [7]) # iterable of length one
Raymond Hettinger5d446132011-03-27 18:52:10 -0700134 self.assertRaises(TypeError, accumulate, range(10), 5, 6) # too many args
Raymond Hettinger482ba772010-12-01 22:48:00 +0000135 self.assertRaises(TypeError, accumulate) # too few args
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000136 self.assertRaises(TypeError, accumulate, x=range(10)) # unexpected kwd arg
Raymond Hettinger482ba772010-12-01 22:48:00 +0000137 self.assertRaises(TypeError, list, accumulate([1, []])) # args that don't add
138
Raymond Hettinger5d446132011-03-27 18:52:10 -0700139 s = [2, 8, 9, 5, 7, 0, 3, 4, 1, 6]
140 self.assertEqual(list(accumulate(s, min)),
141 [2, 2, 2, 2, 2, 0, 0, 0, 0, 0])
142 self.assertEqual(list(accumulate(s, max)),
143 [2, 8, 9, 9, 9, 9, 9, 9, 9, 9])
144 self.assertEqual(list(accumulate(s, operator.mul)),
145 [2, 16, 144, 720, 5040, 0, 0, 0, 0, 0])
146 with self.assertRaises(TypeError):
147 list(accumulate(s, chr)) # unary-operation
Serhiy Storchakabad12572014-12-15 14:03:42 +0200148 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
149 self.pickletest(proto, accumulate(range(10))) # test pickling
Raymond Hettinger5d446132011-03-27 18:52:10 -0700150
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000151 def test_chain(self):
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000152
153 def chain2(*iterables):
154 'Pure python version in the docs'
155 for it in iterables:
156 for element in it:
157 yield element
158
159 for c in (chain, chain2):
160 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
161 self.assertEqual(list(c('abc')), list('abc'))
162 self.assertEqual(list(c('')), [])
163 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
164 self.assertRaises(TypeError, list,c(2, 3))
Christian Heimesf16baeb2008-02-29 14:57:44 +0000165
166 def test_chain_from_iterable(self):
167 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
168 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
169 self.assertEqual(list(chain.from_iterable([''])), [])
170 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
171 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000172
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000173 def test_chain_reducible(self):
Serhiy Storchakabad12572014-12-15 14:03:42 +0200174 for oper in [copy.deepcopy] + picklecopiers:
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000175 it = chain('abc', 'def')
176 self.assertEqual(list(oper(it)), list('abcdef'))
177 self.assertEqual(next(it), 'a')
178 self.assertEqual(list(oper(it)), list('bcdef'))
179
180 self.assertEqual(list(oper(chain(''))), [])
181 self.assertEqual(take(4, oper(chain('abc', 'def'))), list('abcd'))
182 self.assertRaises(TypeError, list, oper(chain(2, 3)))
Serhiy Storchakabad12572014-12-15 14:03:42 +0200183 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
184 self.pickletest(proto, chain('abc', 'def'), compare=list('abcdef'))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000185
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300186 def test_chain_setstate(self):
187 self.assertRaises(TypeError, chain().__setstate__, ())
188 self.assertRaises(TypeError, chain().__setstate__, [])
189 self.assertRaises(TypeError, chain().__setstate__, 0)
190 self.assertRaises(TypeError, chain().__setstate__, ([],))
191 self.assertRaises(TypeError, chain().__setstate__, (iter([]), []))
192 it = chain()
193 it.__setstate__((iter(['abc', 'def']),))
194 self.assertEqual(list(it), ['a', 'b', 'c', 'd', 'e', 'f'])
195 it = chain()
196 it.__setstate__((iter(['abc', 'def']), iter(['ghi'])))
197 self.assertEqual(list(it), ['ghi', 'a', 'b', 'c', 'd', 'e', 'f'])
198
Christian Heimes380f7f22008-02-28 11:19:05 +0000199 def test_combinations(self):
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000200 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Christian Heimes380f7f22008-02-28 11:19:05 +0000201 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Christian Heimes78644762008-03-04 23:39:23 +0000202 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Christian Heimes380f7f22008-02-28 11:19:05 +0000203 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000204
Serhiy Storchakabad12572014-12-15 14:03:42 +0200205 for op in [lambda a:a] + picklecopiers:
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000206 self.assertEqual(list(op(combinations('abc', 32))), []) # r > n
207
208 self.assertEqual(list(op(combinations('ABCD', 2))),
209 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
210 testIntermediate = combinations('ABCD', 2)
211 next(testIntermediate)
212 self.assertEqual(list(op(testIntermediate)),
213 [('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
214
215 self.assertEqual(list(op(combinations(range(4), 3))),
216 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
217 testIntermediate = combinations(range(4), 3)
218 next(testIntermediate)
219 self.assertEqual(list(op(testIntermediate)),
220 [(0,1,3), (0,2,3), (1,2,3)])
221
Christian Heimes78644762008-03-04 23:39:23 +0000222
223 def combinations1(iterable, r):
224 'Pure python version shown in the docs'
225 pool = tuple(iterable)
226 n = len(pool)
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000227 if r > n:
228 return
Christian Heimes78644762008-03-04 23:39:23 +0000229 indices = list(range(r))
230 yield tuple(pool[i] for i in indices)
231 while 1:
232 for i in reversed(range(r)):
233 if indices[i] != i + n - r:
234 break
235 else:
236 return
237 indices[i] += 1
238 for j in range(i+1, r):
239 indices[j] = indices[j-1] + 1
240 yield tuple(pool[i] for i in indices)
241
242 def combinations2(iterable, r):
243 'Pure python version shown in the docs'
244 pool = tuple(iterable)
245 n = len(pool)
246 for indices in permutations(range(n), r):
247 if sorted(indices) == list(indices):
248 yield tuple(pool[i] for i in indices)
249
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000250 def combinations3(iterable, r):
251 'Pure python version from cwr()'
252 pool = tuple(iterable)
253 n = len(pool)
254 for indices in combinations_with_replacement(range(n), r):
255 if len(set(indices)) == r:
256 yield tuple(pool[i] for i in indices)
257
Christian Heimes78644762008-03-04 23:39:23 +0000258 for n in range(7):
Christian Heimes380f7f22008-02-28 11:19:05 +0000259 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000260 for r in range(n+2):
Christian Heimes380f7f22008-02-28 11:19:05 +0000261 result = list(combinations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000262 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 +0000263 self.assertEqual(len(result), len(set(result))) # no repeats
264 self.assertEqual(result, sorted(result)) # lexicographic order
265 for c in result:
266 self.assertEqual(len(c), r) # r-length combinations
267 self.assertEqual(len(set(c)), r) # no duplicate elements
268 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000269 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000270 self.assertEqual(list(c),
271 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000272 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000273 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000274 self.assertEqual(result, list(combinations3(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000275
Serhiy Storchakabad12572014-12-15 14:03:42 +0200276 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
277 self.pickletest(proto, combinations(values, r)) # test pickling
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000278
Benjamin Peterson4b40eeb2015-02-01 20:59:00 -0500279 @support.bigaddrspacetest
280 def test_combinations_overflow(self):
Serhiy Storchakadee948b2015-02-03 01:34:09 +0200281 with self.assertRaises((OverflowError, MemoryError)):
Benjamin Peterson4b40eeb2015-02-01 20:59:00 -0500282 combinations("AA", 2**29)
283
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000284 # Test implementation detail: tuple re-use
Alex Gaynore151d212011-07-17 16:21:30 -0700285 @support.impl_detail("tuple reuse is specific to CPython")
286 def test_combinations_tuple_reuse(self):
Christian Heimes78644762008-03-04 23:39:23 +0000287 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
288 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
289
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000290 def test_combinations_with_replacement(self):
291 cwr = combinations_with_replacement
292 self.assertRaises(TypeError, cwr, 'abc') # missing r argument
293 self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
294 self.assertRaises(TypeError, cwr, None) # pool is not iterable
295 self.assertRaises(ValueError, cwr, 'abc', -2) # r is negative
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000296
Serhiy Storchakabad12572014-12-15 14:03:42 +0200297 for op in [lambda a:a] + picklecopiers:
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000298 self.assertEqual(list(op(cwr('ABC', 2))),
299 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
300 testIntermediate = cwr('ABC', 2)
301 next(testIntermediate)
302 self.assertEqual(list(op(testIntermediate)),
303 [('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
304
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000305
306 def cwr1(iterable, r):
307 'Pure python version shown in the docs'
308 # number items returned: (n+r-1)! / r! / (n-1)! when n>0
309 pool = tuple(iterable)
310 n = len(pool)
311 if not n and r:
312 return
313 indices = [0] * r
314 yield tuple(pool[i] for i in indices)
315 while 1:
316 for i in reversed(range(r)):
317 if indices[i] != n - 1:
318 break
319 else:
320 return
321 indices[i:] = [indices[i] + 1] * (r - i)
322 yield tuple(pool[i] for i in indices)
323
324 def cwr2(iterable, r):
325 'Pure python version shown in the docs'
326 pool = tuple(iterable)
327 n = len(pool)
328 for indices in product(range(n), repeat=r):
329 if sorted(indices) == list(indices):
330 yield tuple(pool[i] for i in indices)
331
332 def numcombs(n, r):
333 if not n:
334 return 0 if r else 1
335 return fact(n+r-1) / fact(r)/ fact(n-1)
336
337 for n in range(7):
338 values = [5*x-12 for x in range(n)]
339 for r in range(n+2):
340 result = list(cwr(values, r))
341
342 self.assertEqual(len(result), numcombs(n, r)) # right number of combs
343 self.assertEqual(len(result), len(set(result))) # no repeats
344 self.assertEqual(result, sorted(result)) # lexicographic order
345
346 regular_combs = list(combinations(values, r)) # compare to combs without replacement
347 if n == 0 or r <= 1:
Ezio Melottib3aedd42010-11-20 19:04:17 +0000348 self.assertEqual(result, regular_combs) # cases that should be identical
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000349 else:
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000350 self.assertTrue(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000351
352 for c in result:
353 self.assertEqual(len(c), r) # r-length combinations
354 noruns = [k for k,v in groupby(c)] # combo without consecutive repeats
355 self.assertEqual(len(noruns), len(set(noruns))) # no repeats other than consecutive
356 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000357 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000358 self.assertEqual(noruns,
359 [e for e in values if e in c]) # comb is a subsequence of the input iterable
360 self.assertEqual(result, list(cwr1(values, r))) # matches first pure python version
361 self.assertEqual(result, list(cwr2(values, r))) # matches second pure python version
362
Serhiy Storchakabad12572014-12-15 14:03:42 +0200363 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
364 self.pickletest(proto, cwr(values,r)) # test pickling
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000365
Benjamin Peterson6f082292015-02-01 21:10:47 -0500366 @support.bigaddrspacetest
367 def test_combinations_with_replacement_overflow(self):
Serhiy Storchakadee948b2015-02-03 01:34:09 +0200368 with self.assertRaises((OverflowError, MemoryError)):
Benjamin Peterson6f082292015-02-01 21:10:47 -0500369 combinations_with_replacement("AA", 2**30)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000370
Benjamin Peterson6f082292015-02-01 21:10:47 -0500371 # Test implementation detail: tuple re-use
Alex Gaynore151d212011-07-17 16:21:30 -0700372 @support.impl_detail("tuple reuse is specific to CPython")
373 def test_combinations_with_replacement_tuple_reuse(self):
374 cwr = combinations_with_replacement
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000375 self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
376 self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
377
Christian Heimes78644762008-03-04 23:39:23 +0000378 def test_permutations(self):
379 self.assertRaises(TypeError, permutations) # too few arguments
380 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000381 self.assertRaises(TypeError, permutations, None) # pool is not iterable
382 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000383 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000384 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Christian Heimes78644762008-03-04 23:39:23 +0000385 self.assertEqual(list(permutations(range(3), 2)),
386 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
387
388 def permutations1(iterable, r=None):
389 'Pure python version shown in the docs'
390 pool = tuple(iterable)
391 n = len(pool)
392 r = n if r is None else r
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000393 if r > n:
394 return
Christian Heimes78644762008-03-04 23:39:23 +0000395 indices = list(range(n))
396 cycles = list(range(n-r+1, n+1))[::-1]
397 yield tuple(pool[i] for i in indices[:r])
398 while n:
399 for i in reversed(range(r)):
400 cycles[i] -= 1
401 if cycles[i] == 0:
402 indices[i:] = indices[i+1:] + indices[i:i+1]
403 cycles[i] = n - i
404 else:
405 j = cycles[i]
406 indices[i], indices[-j] = indices[-j], indices[i]
407 yield tuple(pool[i] for i in indices[:r])
408 break
409 else:
410 return
411
412 def permutations2(iterable, r=None):
413 'Pure python version shown in the docs'
414 pool = tuple(iterable)
415 n = len(pool)
416 r = n if r is None else r
417 for indices in product(range(n), repeat=r):
418 if len(set(indices)) == r:
419 yield tuple(pool[i] for i in indices)
420
421 for n in range(7):
422 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000423 for r in range(n+2):
Christian Heimes78644762008-03-04 23:39:23 +0000424 result = list(permutations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000425 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 +0000426 self.assertEqual(len(result), len(set(result))) # no repeats
427 self.assertEqual(result, sorted(result)) # lexicographic order
428 for p in result:
429 self.assertEqual(len(p), r) # r-length permutations
430 self.assertEqual(len(set(p)), r) # no duplicate elements
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000431 self.assertTrue(all(e in values for e in p)) # elements taken from input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000432 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000433 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000434 if r == n:
435 self.assertEqual(result, list(permutations(values, None))) # test r as None
436 self.assertEqual(result, list(permutations(values))) # test default r
437
Serhiy Storchakabad12572014-12-15 14:03:42 +0200438 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
439 self.pickletest(proto, permutations(values, r)) # test pickling
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000440
Benjamin Peterson0eaabf12015-02-01 21:34:07 -0500441 @support.bigaddrspacetest
442 def test_permutations_overflow(self):
Serhiy Storchakadee948b2015-02-03 01:34:09 +0200443 with self.assertRaises((OverflowError, MemoryError)):
Benjamin Peterson0eaabf12015-02-01 21:34:07 -0500444 permutations("A", 2**30)
Benjamin Peterson0eaabf12015-02-01 21:34:07 -0500445
Zachary Waredca807b2014-04-24 13:22:05 -0500446 @support.impl_detail("tuple reuse is specific to CPython")
Alex Gaynore151d212011-07-17 16:21:30 -0700447 def test_permutations_tuple_reuse(self):
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000448 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Christian Heimes78644762008-03-04 23:39:23 +0000449 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Christian Heimes380f7f22008-02-28 11:19:05 +0000450
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000451 def test_combinatorics(self):
452 # Test relationships between product(), permutations(),
453 # combinations() and combinations_with_replacement().
454
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000455 for n in range(6):
456 s = 'ABCDEFG'[:n]
457 for r in range(8):
458 prod = list(product(s, repeat=r))
459 cwr = list(combinations_with_replacement(s, r))
460 perm = list(permutations(s, r))
461 comb = list(combinations(s, r))
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000462
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000463 # Check size
Ezio Melottib3aedd42010-11-20 19:04:17 +0000464 self.assertEqual(len(prod), n**r)
465 self.assertEqual(len(cwr), (fact(n+r-1) / fact(r)/ fact(n-1)) if n else (not r))
466 self.assertEqual(len(perm), 0 if r>n else fact(n) / fact(n-r))
467 self.assertEqual(len(comb), 0 if r>n else fact(n) / fact(r) / fact(n-r))
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000468
469 # Check lexicographic order without repeated tuples
Ezio Melottib3aedd42010-11-20 19:04:17 +0000470 self.assertEqual(prod, sorted(set(prod)))
471 self.assertEqual(cwr, sorted(set(cwr)))
472 self.assertEqual(perm, sorted(set(perm)))
473 self.assertEqual(comb, sorted(set(comb)))
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000474
475 # Check interrelationships
Ezio Melottib3aedd42010-11-20 19:04:17 +0000476 self.assertEqual(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
477 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 +0000478 self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
479 self.assertEqual(comb, [t for t in cwr if len(set(t))==r]) # comb: cwrs without dups
480 self.assertEqual(comb, list(filter(set(cwr).__contains__, perm))) # comb: perm that is a cwr
481 self.assertEqual(comb, list(filter(set(perm).__contains__, cwr))) # comb: cwr that is a perm
482 self.assertEqual(comb, sorted(set(cwr) & set(perm))) # comb: both a cwr and a perm
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000483
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000484 def test_compress(self):
Raymond Hettinger15a49502009-02-19 02:17:09 +0000485 self.assertEqual(list(compress(data='ABCDEF', selectors=[1,0,1,0,1,1])), list('ACEF'))
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000486 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
487 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
488 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
489 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
490 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
491 n = 10000
492 data = chain.from_iterable(repeat(range(6), n))
493 selectors = chain.from_iterable(repeat((0, 1)))
494 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
495 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
496 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
497 self.assertRaises(TypeError, compress, range(6)) # too few args
498 self.assertRaises(TypeError, compress, range(6), None) # too many args
499
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000500 # check copy, deepcopy, pickle
Serhiy Storchakabad12572014-12-15 14:03:42 +0200501 for op in [lambda a:copy.copy(a), lambda a:copy.deepcopy(a)] + picklecopiers:
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000502 for data, selectors, result1, result2 in [
503 ('ABCDEF', [1,0,1,0,1,1], 'ACEF', 'CEF'),
504 ('ABCDEF', [0,0,0,0,0,0], '', ''),
505 ('ABCDEF', [1,1,1,1,1,1], 'ABCDEF', 'BCDEF'),
506 ('ABCDEF', [1,0,1], 'AC', 'C'),
507 ('ABC', [0,1,1,1,1,1], 'BC', 'C'),
508 ]:
509
510 self.assertEqual(list(op(compress(data=data, selectors=selectors))), list(result1))
511 self.assertEqual(list(op(compress(data, selectors))), list(result1))
512 testIntermediate = compress(data, selectors)
513 if result1:
514 next(testIntermediate)
515 self.assertEqual(list(op(testIntermediate)), list(result2))
516
517
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000518 def test_count(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000519 self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
520 self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
521 self.assertEqual(take(2, lzip('abc',count(3))), [('a', 3), ('b', 4)])
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000522 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
523 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000524 self.assertRaises(TypeError, count, 2, 3, 4)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000525 self.assertRaises(TypeError, count, 'a')
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300526 self.assertEqual(take(10, count(maxsize-5)),
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000527 list(range(maxsize-5, maxsize+5)))
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300528 self.assertEqual(take(10, count(-maxsize-5)),
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000529 list(range(-maxsize-5, -maxsize+5)))
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300530 self.assertEqual(take(3, count(3.25)), [3.25, 4.25, 5.25])
531 self.assertEqual(take(3, count(3.25-4j)), [3.25-4j, 4.25-4j, 5.25-4j])
532 self.assertEqual(take(3, count(Decimal('1.1'))),
533 [Decimal('1.1'), Decimal('2.1'), Decimal('3.1')])
534 self.assertEqual(take(3, count(Fraction(2, 3))),
535 [Fraction(2, 3), Fraction(5, 3), Fraction(8, 3)])
536 BIGINT = 1<<1000
537 self.assertEqual(take(3, count(BIGINT)), [BIGINT, BIGINT+1, BIGINT+2])
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000538 c = count(3)
539 self.assertEqual(repr(c), 'count(3)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000540 next(c)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000541 self.assertEqual(repr(c), 'count(4)')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000542 c = count(-9)
543 self.assertEqual(repr(c), 'count(-9)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000544 next(c)
545 self.assertEqual(next(c), -8)
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300546 self.assertEqual(repr(count(10.25)), 'count(10.25)')
547 self.assertEqual(repr(count(10.0)), 'count(10.0)')
548 self.assertEqual(type(next(count(10.0))), float)
Christian Heimesa37d4c62007-12-04 23:02:19 +0000549 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
Serhiy Storchaka95949422013-08-27 19:40:23 +0300550 # Test repr
551 r1 = repr(count(i))
552 r2 = 'count(%r)'.__mod__(i)
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000553 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000554
Raymond Hettingera3e1ad22009-11-30 22:02:31 +0000555 # check copy, deepcopy, pickle
556 for value in -3, 3, maxsize-5, maxsize+5:
557 c = count(value)
558 self.assertEqual(next(copy.copy(c)), value)
559 self.assertEqual(next(copy.deepcopy(c)), value)
Serhiy Storchakabad12572014-12-15 14:03:42 +0200560 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
561 self.pickletest(proto, count(value))
Raymond Hettingera3e1ad22009-11-30 22:02:31 +0000562
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +0000563 #check proper internal error handling for large "step' sizes
564 count(1, maxsize+5); sys.exc_info()
565
Raymond Hettinger30729212009-02-12 06:28:27 +0000566 def test_count_with_stride(self):
567 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingereb13fdd2009-02-21 22:30:12 +0000568 self.assertEqual(lzip('abc',count(start=2,step=3)),
569 [('a', 2), ('b', 5), ('c', 8)])
570 self.assertEqual(lzip('abc',count(step=-1)),
571 [('a', 0), ('b', -1), ('c', -2)])
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300572 self.assertRaises(TypeError, count, 'a', 'b')
Raymond Hettinger30729212009-02-12 06:28:27 +0000573 self.assertEqual(lzip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
574 self.assertEqual(lzip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000575 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000576 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
577 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300578 self.assertEqual(take(3, count(10, maxsize+5)),
579 list(range(10, 10+3*(maxsize+5), maxsize+5)))
580 self.assertEqual(take(3, count(2, 1.25)), [2, 3.25, 4.5])
Raymond Hettinger30729212009-02-12 06:28:27 +0000581 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettinger8a882a72009-02-12 12:08:18 +0000582 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
583 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerde09acf2009-02-12 12:53:53 +0000584 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
585 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300586 BIGINT = 1<<1000
587 self.assertEqual(take(3, count(step=BIGINT)), [0, BIGINT, 2*BIGINT])
Raymond Hettinger30729212009-02-12 06:28:27 +0000588 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
589 c = count(3, 5)
590 self.assertEqual(repr(c), 'count(3, 5)')
591 next(c)
592 self.assertEqual(repr(c), 'count(8, 5)')
593 c = count(-9, 0)
594 self.assertEqual(repr(c), 'count(-9, 0)')
595 next(c)
596 self.assertEqual(repr(c), 'count(-9, 0)')
597 c = count(-9, -3)
598 self.assertEqual(repr(c), 'count(-9, -3)')
599 next(c)
600 self.assertEqual(repr(c), 'count(-12, -3)')
601 self.assertEqual(repr(c), 'count(-12, -3)')
602 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
603 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
604 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300605 self.assertEqual(repr(count(10, 1.00)), 'count(10, 1.0)')
606 c = count(10, 1.0)
607 self.assertEqual(type(next(c)), int)
608 self.assertEqual(type(next(c)), float)
Raymond Hettinger30729212009-02-12 06:28:27 +0000609 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
610 for j in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 1, 10, sys.maxsize-5, sys.maxsize+5):
Serhiy Storchaka95949422013-08-27 19:40:23 +0300611 # Test repr
612 r1 = repr(count(i, j))
Raymond Hettinger30729212009-02-12 06:28:27 +0000613 if j == 1:
Serhiy Storchaka95949422013-08-27 19:40:23 +0300614 r2 = ('count(%r)' % i)
Raymond Hettinger30729212009-02-12 06:28:27 +0000615 else:
Serhiy Storchaka95949422013-08-27 19:40:23 +0300616 r2 = ('count(%r, %r)' % (i, j))
Raymond Hettinger30729212009-02-12 06:28:27 +0000617 self.assertEqual(r1, r2)
Serhiy Storchakabad12572014-12-15 14:03:42 +0200618 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
619 self.pickletest(proto, count(i, j))
Raymond Hettinger30729212009-02-12 06:28:27 +0000620
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000621 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000622 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000623 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000624 self.assertRaises(TypeError, cycle)
625 self.assertRaises(TypeError, cycle, 5)
626 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000627
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000628 # check copy, deepcopy, pickle
629 c = cycle('abc')
630 self.assertEqual(next(c), 'a')
631 #simple copy currently not supported, because __reduce__ returns
632 #an internal iterator
633 #self.assertEqual(take(10, copy.copy(c)), list('bcabcabcab'))
634 self.assertEqual(take(10, copy.deepcopy(c)), list('bcabcabcab'))
Serhiy Storchakabad12572014-12-15 14:03:42 +0200635 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
636 self.assertEqual(take(10, pickle.loads(pickle.dumps(c, proto))),
637 list('bcabcabcab'))
638 next(c)
639 self.assertEqual(take(10, pickle.loads(pickle.dumps(c, proto))),
640 list('cabcabcabc'))
641 next(c)
642 next(c)
643 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
644 self.pickletest(proto, cycle('abc'))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000645
Raymond Hettingera166ce52015-08-15 14:45:49 -0700646 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
647 # test with partial consumed input iterable
648 it = iter('abcde')
649 c = cycle(it)
Raymond Hettinger1cadf762015-08-15 14:47:27 -0700650 _ = [next(c) for i in range(2)] # consume 2 of 5 inputs
Raymond Hettingera166ce52015-08-15 14:45:49 -0700651 p = pickle.dumps(c, proto)
652 d = pickle.loads(p) # rebuild the cycle object
653 self.assertEqual(take(20, d), list('cdeabcdeabcdeabcdeab'))
654
655 # test with completely consumed input iterable
656 it = iter('abcde')
657 c = cycle(it)
Raymond Hettinger1cadf762015-08-15 14:47:27 -0700658 _ = [next(c) for i in range(7)] # consume 7 of 5 inputs
Raymond Hettingera166ce52015-08-15 14:45:49 -0700659 p = pickle.dumps(c, proto)
660 d = pickle.loads(p) # rebuild the cycle object
661 self.assertEqual(take(20, d), list('cdeabcdeabcdeabcdeab'))
662
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700663 def test_cycle_setstate(self):
664 # Verify both modes for restoring state
665
666 # Mode 0 is efficient. It uses an incompletely consumed input
667 # iterator to build a cycle object and then passes in state with
668 # a list of previously consumed values. There is no data
Martin Panterf1579822016-05-26 06:03:33 +0000669 # overlap between the two.
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700670 c = cycle('defg')
671 c.__setstate__((list('abc'), 0))
672 self.assertEqual(take(20, c), list('defgabcdefgabcdefgab'))
673
674 # Mode 1 is inefficient. It starts with a cycle object built
675 # from an iterator over the remaining elements in a partial
676 # cycle and then passes in state with all of the previously
677 # seen values (this overlaps values included in the iterator).
678 c = cycle('defg')
679 c.__setstate__((list('abcdefg'), 1))
680 self.assertEqual(take(20, c), list('defgabcdefgabcdefgab'))
681
682 # The first argument to setstate needs to be a tuple
Serhiy Storchaka8f0f2052016-10-02 09:13:14 +0300683 with self.assertRaises(TypeError):
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700684 cycle('defg').__setstate__([list('abcdefg'), 0])
685
686 # The first argument in the setstate tuple must be a list
687 with self.assertRaises(TypeError):
688 c = cycle('defg')
Serhiy Storchaka8f0f2052016-10-02 09:13:14 +0300689 c.__setstate__((tuple('defg'), 0))
690 take(20, c)
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700691
Serhiy Storchaka8f0f2052016-10-02 09:13:14 +0300692 # The second argument in the setstate tuple must be an int
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700693 with self.assertRaises(TypeError):
694 cycle('defg').__setstate__((list('abcdefg'), 'x'))
695
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300696 self.assertRaises(TypeError, cycle('').__setstate__, ())
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300697 self.assertRaises(TypeError, cycle('').__setstate__, ([],))
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300698
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000699 def test_groupby(self):
700 # Check whether it accepts arguments correctly
701 self.assertEqual([], list(groupby([])))
702 self.assertEqual([], list(groupby([], key=id)))
703 self.assertRaises(TypeError, list, groupby('abc', []))
704 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000705 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000706
707 # Check normal input
708 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
709 (2,15,22), (3,16,23), (3,17,23)]
710 dup = []
711 for k, g in groupby(s, lambda r:r[0]):
712 for elem in g:
713 self.assertEqual(k, elem[0])
714 dup.append(elem)
715 self.assertEqual(s, dup)
716
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000717 # Check normal pickled
Serhiy Storchakabad12572014-12-15 14:03:42 +0200718 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
719 dup = []
720 for k, g in pickle.loads(pickle.dumps(groupby(s, testR), proto)):
721 for elem in g:
722 self.assertEqual(k, elem[0])
723 dup.append(elem)
724 self.assertEqual(s, dup)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000725
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000726 # Check nested case
727 dup = []
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000728 for k, g in groupby(s, testR):
729 for ik, ig in groupby(g, testR2):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000730 for elem in ig:
731 self.assertEqual(k, elem[0])
732 self.assertEqual(ik, elem[2])
733 dup.append(elem)
734 self.assertEqual(s, dup)
735
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000736 # Check nested and pickled
Serhiy Storchakabad12572014-12-15 14:03:42 +0200737 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
738 dup = []
739 for k, g in pickle.loads(pickle.dumps(groupby(s, testR), proto)):
740 for ik, ig in pickle.loads(pickle.dumps(groupby(g, testR2), proto)):
741 for elem in ig:
742 self.assertEqual(k, elem[0])
743 self.assertEqual(ik, elem[2])
744 dup.append(elem)
745 self.assertEqual(s, dup)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000746
747
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000748 # Check case where inner iterator is not used
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000749 keys = [k for k, g in groupby(s, testR)]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000750 expectedkeys = set([r[0] for r in s])
751 self.assertEqual(set(keys), expectedkeys)
752 self.assertEqual(len(keys), len(expectedkeys))
753
754 # Exercise pipes and filters style
755 s = 'abracadabra'
756 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000757 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000758 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
759 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000760 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000761 self.assertEqual(r, ['a', 'b', 'r'])
762 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000763 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000764 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
765 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000766 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000767 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
768
Georg Brandla18af4e2007-04-21 15:47:16 +0000769 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000770 class ExpectedError(Exception):
771 pass
772 def delayed_raise(n=0):
773 for i in range(n):
774 yield 'yo'
775 raise ExpectedError
776 def gulp(iterable, keyp=None, func=list):
777 return [func(g) for k, g in groupby(iterable, keyp)]
778
Georg Brandla18af4e2007-04-21 15:47:16 +0000779 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000780 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000781 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000782 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
783
Serhiy Storchakaa60c2fe2015-03-12 21:56:08 +0200784 # __eq__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000785 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000786 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000787 raise ExpectedError
788 s = [DummyCmp(), DummyCmp(), None]
789
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000790 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000791 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000792 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000793 self.assertRaises(ExpectedError, gulp, s)
794
795 # keyfunc failure
796 def keyfunc(obj):
797 if keyfunc.skip > 0:
798 keyfunc.skip -= 1
799 return obj
800 else:
801 raise ExpectedError
802
803 # keyfunc failure on outer object
804 keyfunc.skip = 0
805 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
806 keyfunc.skip = 1
807 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
808
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000809 def test_filter(self):
810 self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
811 self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])
812 self.assertEqual(list(filter(bool, [0,1,0,2,0])), [1,2])
813 self.assertEqual(take(4, filter(isEven, count())), [0,2,4,6])
814 self.assertRaises(TypeError, filter)
815 self.assertRaises(TypeError, filter, lambda x:x)
816 self.assertRaises(TypeError, filter, lambda x:x, range(6), 7)
817 self.assertRaises(TypeError, filter, isEven, 3)
818 self.assertRaises(TypeError, next, filter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000819
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000820 # check copy, deepcopy, pickle
821 ans = [0,2,4]
822
823 c = filter(isEven, range(6))
824 self.assertEqual(list(copy.copy(c)), ans)
825 c = filter(isEven, range(6))
826 self.assertEqual(list(copy.deepcopy(c)), ans)
Serhiy Storchakabad12572014-12-15 14:03:42 +0200827 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
828 c = filter(isEven, range(6))
829 self.assertEqual(list(pickle.loads(pickle.dumps(c, proto))), ans)
830 next(c)
831 self.assertEqual(list(pickle.loads(pickle.dumps(c, proto))), ans[1:])
832 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
833 c = filter(isEven, range(6))
834 self.pickletest(proto, c)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000835
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000836 def test_filterfalse(self):
837 self.assertEqual(list(filterfalse(isEven, range(6))), [1,3,5])
838 self.assertEqual(list(filterfalse(None, [0,1,0,2,0])), [0,0,0])
839 self.assertEqual(list(filterfalse(bool, [0,1,0,2,0])), [0,0,0])
840 self.assertEqual(take(4, filterfalse(isEven, count())), [1,3,5,7])
841 self.assertRaises(TypeError, filterfalse)
842 self.assertRaises(TypeError, filterfalse, lambda x:x)
843 self.assertRaises(TypeError, filterfalse, lambda x:x, range(6), 7)
844 self.assertRaises(TypeError, filterfalse, isEven, 3)
845 self.assertRaises(TypeError, next, filterfalse(range(6), range(6)))
Serhiy Storchakabad12572014-12-15 14:03:42 +0200846 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
847 self.pickletest(proto, filterfalse(isEven, range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000848
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000849 def test_zip(self):
850 # XXX This is rather silly now that builtin zip() calls zip()...
851 ans = [(x,y) for x, y in zip('abc',count())]
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000852 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000853 self.assertEqual(list(zip('abc', range(6))), lzip('abc', range(6)))
854 self.assertEqual(list(zip('abcdef', range(3))), lzip('abcdef', range(3)))
855 self.assertEqual(take(3,zip('abcdef', count())), lzip('abcdef', range(3)))
856 self.assertEqual(list(zip('abcdef')), lzip('abcdef'))
857 self.assertEqual(list(zip()), lzip())
858 self.assertRaises(TypeError, zip, 3)
859 self.assertRaises(TypeError, zip, range(3), 3)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000860 self.assertEqual([tuple(list(pair)) for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000861 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000862 self.assertEqual([pair for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000863 lzip('abc', 'def'))
Alex Gaynore151d212011-07-17 16:21:30 -0700864
865 @support.impl_detail("tuple reuse is specific to CPython")
866 def test_zip_tuple_reuse(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000867 ids = list(map(id, zip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000868 self.assertEqual(min(ids), max(ids))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000869 ids = list(map(id, list(zip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000870 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000871
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000872 # check copy, deepcopy, pickle
873 ans = [(x,y) for x, y in copy.copy(zip('abc',count()))]
874 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
875
876 ans = [(x,y) for x, y in copy.deepcopy(zip('abc',count()))]
877 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
878
Serhiy Storchakabad12572014-12-15 14:03:42 +0200879 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
880 ans = [(x,y) for x, y in pickle.loads(pickle.dumps(zip('abc',count()), proto))]
881 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000882
Serhiy Storchakabad12572014-12-15 14:03:42 +0200883 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
884 testIntermediate = zip('abc',count())
885 next(testIntermediate)
886 ans = [(x,y) for x, y in pickle.loads(pickle.dumps(testIntermediate, proto))]
887 self.assertEqual(ans, [('b', 1), ('c', 2)])
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000888
Serhiy Storchakabad12572014-12-15 14:03:42 +0200889 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
890 self.pickletest(proto, zip('abc', count()))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000891
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000892 def test_ziplongest(self):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000893 for args in [
894 ['abc', range(6)],
895 [range(6), 'abc'],
896 [range(1000), range(2000,2100), range(3000,3050)],
897 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
898 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
899 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000900 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
901 for i in range(max(map(len, args)))]
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000902 self.assertEqual(list(zip_longest(*args)), target)
903 self.assertEqual(list(zip_longest(*args, **{})), target)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000904 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 +0000905 self.assertEqual(list(zip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000906
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000907 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 +0000908
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000909 self.assertEqual(list(zip_longest()), list(zip()))
910 self.assertEqual(list(zip_longest([])), list(zip([])))
911 self.assertEqual(list(zip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000912
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000913 self.assertEqual(list(zip_longest('abc', 'defg', **{})),
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000914 list(zip(list('abc')+[None], 'defg'))) # empty keyword dict
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000915 self.assertRaises(TypeError, zip_longest, 3)
916 self.assertRaises(TypeError, zip_longest, range(3), 3)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000917
918 for stmt in [
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000919 "zip_longest('abc', fv=1)",
920 "zip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000921 ]:
922 try:
923 eval(stmt, globals(), locals())
924 except TypeError:
925 pass
926 else:
927 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000928
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000929 self.assertEqual([tuple(list(pair)) for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000930 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000931 self.assertEqual([pair for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000932 list(zip('abc', 'def')))
Alex Gaynore151d212011-07-17 16:21:30 -0700933
934 @support.impl_detail("tuple reuse is specific to CPython")
935 def test_zip_longest_tuple_reuse(self):
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000936 ids = list(map(id, zip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000937 self.assertEqual(min(ids), max(ids))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000938 ids = list(map(id, list(zip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000939 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
940
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000941 def test_zip_longest_pickling(self):
Serhiy Storchakabad12572014-12-15 14:03:42 +0200942 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
943 self.pickletest(proto, zip_longest("abc", "def"))
944 self.pickletest(proto, zip_longest("abc", "defgh"))
945 self.pickletest(proto, zip_longest("abc", "defgh", fillvalue=1))
946 self.pickletest(proto, zip_longest("", "defgh"))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000947
Raymond Hettingerfc438512009-11-01 20:55:33 +0000948 def test_bug_7244(self):
949
950 class Repeater:
951 # this class is similar to itertools.repeat
952 def __init__(self, o, t, e):
953 self.o = o
954 self.t = int(t)
955 self.e = e
956 def __iter__(self): # its iterator is itself
957 return self
958 def __next__(self):
959 if self.t > 0:
960 self.t -= 1
961 return self.o
962 else:
963 raise self.e
964
965 # Formerly this code in would fail in debug mode
966 # with Undetected Error and Stop Iteration
967 r1 = Repeater(1, 3, StopIteration)
968 r2 = Repeater(2, 4, StopIteration)
969 def run(r1, r2):
970 result = []
971 for i, j in zip_longest(r1, r2, fillvalue=0):
972 with support.captured_output('stdout'):
973 print((i, j))
974 result.append((i, j))
975 return result
976 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
977
978 # Formerly, the RuntimeError would be lost
979 # and StopIteration would stop as expected
980 r1 = Repeater(1, 3, RuntimeError)
981 r2 = Repeater(2, 4, StopIteration)
982 it = zip_longest(r1, r2, fillvalue=0)
983 self.assertEqual(next(it), (1, 2))
984 self.assertEqual(next(it), (1, 2))
985 self.assertEqual(next(it), (1, 2))
986 self.assertRaises(RuntimeError, next, it)
987
Christian Heimesc3f30c42008-02-22 16:37:40 +0000988 def test_product(self):
989 for args, result in [
Christian Heimes78644762008-03-04 23:39:23 +0000990 ([], [()]), # zero iterables
Christian Heimesc3f30c42008-02-22 16:37:40 +0000991 (['ab'], [('a',), ('b',)]), # one iterable
992 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
993 ([range(0), range(2), range(3)], []), # first iterable with zero length
994 ([range(2), range(0), range(3)], []), # middle iterable with zero length
995 ([range(2), range(3), range(0)], []), # last iterable with zero length
996 ]:
997 self.assertEqual(list(product(*args)), result)
Christian Heimesf16baeb2008-02-29 14:57:44 +0000998 for r in range(4):
999 self.assertEqual(list(product(*(args*r))),
1000 list(product(*args, **dict(repeat=r))))
Christian Heimesc3f30c42008-02-22 16:37:40 +00001001 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
1002 self.assertRaises(TypeError, product, range(6), None)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001003
1004 def product1(*args, **kwds):
1005 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
1006 n = len(pools)
1007 if n == 0:
1008 yield ()
1009 return
1010 if any(len(pool) == 0 for pool in pools):
1011 return
1012 indices = [0] * n
1013 yield tuple(pool[i] for pool, i in zip(pools, indices))
1014 while 1:
1015 for i in reversed(range(n)): # right to left
1016 if indices[i] == len(pools[i]) - 1:
1017 continue
1018 indices[i] += 1
1019 for j in range(i+1, n):
1020 indices[j] = 0
1021 yield tuple(pool[i] for pool, i in zip(pools, indices))
1022 break
1023 else:
1024 return
1025
1026 def product2(*args, **kwds):
1027 'Pure python version used in docs'
1028 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
1029 result = [[]]
1030 for pool in pools:
1031 result = [x+[y] for x in result for y in pool]
1032 for prod in result:
1033 yield tuple(prod)
1034
Christian Heimesc3f30c42008-02-22 16:37:40 +00001035 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
1036 set('abcdefg'), range(11), tuple(range(13))]
1037 for i in range(100):
1038 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Christian Heimes78644762008-03-04 23:39:23 +00001039 expected_len = prod(map(len, args))
1040 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001041 self.assertEqual(list(product(*args)), list(product1(*args)))
1042 self.assertEqual(list(product(*args)), list(product2(*args)))
Christian Heimesc3f30c42008-02-22 16:37:40 +00001043 args = map(iter, args)
Christian Heimes78644762008-03-04 23:39:23 +00001044 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesc3f30c42008-02-22 16:37:40 +00001045
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05001046 @support.bigaddrspacetest
1047 def test_product_overflow(self):
Serhiy Storchakadee948b2015-02-03 01:34:09 +02001048 with self.assertRaises((OverflowError, MemoryError)):
1049 product(*(['ab']*2**5), repeat=2**25)
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05001050
Alex Gaynore151d212011-07-17 16:21:30 -07001051 @support.impl_detail("tuple reuse is specific to CPython")
1052 def test_product_tuple_reuse(self):
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001053 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
1054 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +00001055
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001056 def test_product_pickling(self):
1057 # check copy, deepcopy, pickle
1058 for args, result in [
1059 ([], [()]), # zero iterables
1060 (['ab'], [('a',), ('b',)]), # one iterable
1061 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
1062 ([range(0), range(2), range(3)], []), # first iterable with zero length
1063 ([range(2), range(0), range(3)], []), # middle iterable with zero length
1064 ([range(2), range(3), range(0)], []), # last iterable with zero length
1065 ]:
1066 self.assertEqual(list(copy.copy(product(*args))), result)
1067 self.assertEqual(list(copy.deepcopy(product(*args))), result)
Serhiy Storchakabad12572014-12-15 14:03:42 +02001068 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1069 self.pickletest(proto, product(*args))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001070
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00001071 def test_product_issue_25021(self):
1072 # test that indices are properly clamped to the length of the tuples
1073 p = product((1, 2),(3,))
1074 p.__setstate__((0, 0x1000)) # will access tuple element 1 if not clamped
1075 self.assertEqual(next(p), (2, 3))
1076 # test that empty tuple in the list will result in an immediate StopIteration
1077 p = product((1, 2), (), (3,))
1078 p.__setstate__((0, 0, 0x1000)) # will access tuple element 1 if not clamped
1079 self.assertRaises(StopIteration, next, p)
1080
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001081 def test_repeat(self):
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00001082 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
Guido van Rossum805365e2007-05-07 22:24:25 +00001083 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001084 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001085 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +00001086 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001087 self.assertEqual(list(repeat('a', 0)), [])
1088 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001089 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001090 self.assertRaises(TypeError, repeat, None, 3, 4)
1091 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001092 r = repeat(1+0j)
1093 self.assertEqual(repr(r), 'repeat((1+0j))')
1094 r = repeat(1+0j, 5)
1095 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
1096 list(r)
1097 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001098
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001099 # check copy, deepcopy, pickle
1100 c = repeat(object='a', times=10)
1101 self.assertEqual(next(c), 'a')
1102 self.assertEqual(take(2, copy.copy(c)), list('a' * 2))
1103 self.assertEqual(take(2, copy.deepcopy(c)), list('a' * 2))
Serhiy Storchakabad12572014-12-15 14:03:42 +02001104 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1105 self.pickletest(proto, repeat(object='a', times=10))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001106
Raymond Hettinger97d35552014-06-24 21:36:58 -07001107 def test_repeat_with_negative_times(self):
1108 self.assertEqual(repr(repeat('a', -1)), "repeat('a', 0)")
1109 self.assertEqual(repr(repeat('a', -2)), "repeat('a', 0)")
1110 self.assertEqual(repr(repeat('a', times=-1)), "repeat('a', 0)")
1111 self.assertEqual(repr(repeat('a', times=-2)), "repeat('a', 0)")
1112
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001113 def test_map(self):
1114 self.assertEqual(list(map(operator.pow, range(3), range(1,7))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001115 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001116 self.assertEqual(list(map(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001117 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001118 self.assertEqual(list(map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +00001119 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001120 self.assertEqual(take(2,map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +00001121 [('a',0),('b',1)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001122 self.assertEqual(list(map(operator.pow, [])), [])
1123 self.assertRaises(TypeError, map)
1124 self.assertRaises(TypeError, list, map(None, range(3), range(3)))
1125 self.assertRaises(TypeError, map, operator.neg)
1126 self.assertRaises(TypeError, next, map(10, range(5)))
1127 self.assertRaises(ValueError, next, map(errfunc, [4], [5]))
1128 self.assertRaises(TypeError, next, map(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001129
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001130 # check copy, deepcopy, pickle
1131 ans = [('a',0),('b',1),('c',2)]
1132
1133 c = map(tupleize, 'abc', count())
1134 self.assertEqual(list(copy.copy(c)), ans)
1135
1136 c = map(tupleize, 'abc', count())
1137 self.assertEqual(list(copy.deepcopy(c)), ans)
1138
Serhiy Storchakabad12572014-12-15 14:03:42 +02001139 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1140 c = map(tupleize, 'abc', count())
1141 self.pickletest(proto, c)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001142
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001143 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001144 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
1145 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001146 self.assertEqual(take(3, starmap(operator.pow, zip(count(), count(1)))),
Raymond Hettinger02420702003-06-29 20:36:23 +00001147 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001148 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +00001149 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
1150 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001151 self.assertRaises(TypeError, starmap)
1152 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001153 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
1154 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
1155 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001156
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001157 # check copy, deepcopy, pickle
1158 ans = [0**1, 1**2, 2**3]
1159
1160 c = starmap(operator.pow, zip(range(3), range(1,7)))
1161 self.assertEqual(list(copy.copy(c)), ans)
1162
1163 c = starmap(operator.pow, zip(range(3), range(1,7)))
1164 self.assertEqual(list(copy.deepcopy(c)), ans)
1165
Serhiy Storchakabad12572014-12-15 14:03:42 +02001166 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1167 c = starmap(operator.pow, zip(range(3), range(1,7)))
1168 self.pickletest(proto, c)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001169
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001170 def test_islice(self):
1171 for args in [ # islice(args) should agree with range(args)
1172 (10, 20, 3),
1173 (10, 3, 20),
1174 (10, 20),
1175 (10, 3),
1176 (20,)
1177 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +00001178 self.assertEqual(list(islice(range(100), *args)),
1179 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001180
1181 for args, tgtargs in [ # Stop when seqn is exhausted
1182 ((10, 110, 3), ((10, 100, 3))),
1183 ((10, 110), ((10, 100))),
1184 ((110,), (100,))
1185 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +00001186 self.assertEqual(list(islice(range(100), *args)),
1187 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001188
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001189 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +00001190 self.assertEqual(list(islice(range(10), None)), list(range(10)))
1191 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
1192 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
1193 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
1194 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001195
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +00001196 # Test number of items consumed SF #1171417
1197 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +00001198 self.assertEqual(list(islice(it, 3)), list(range(3)))
1199 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +00001200
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001201 # Test invalid arguments
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001202 ra = range(10)
1203 self.assertRaises(TypeError, islice, ra)
1204 self.assertRaises(TypeError, islice, ra, 1, 2, 3, 4)
1205 self.assertRaises(ValueError, islice, ra, -5, 10, 1)
1206 self.assertRaises(ValueError, islice, ra, 1, -5, -1)
1207 self.assertRaises(ValueError, islice, ra, 1, 10, -1)
1208 self.assertRaises(ValueError, islice, ra, 1, 10, 0)
1209 self.assertRaises(ValueError, islice, ra, 'a')
1210 self.assertRaises(ValueError, islice, ra, 'a', 1)
1211 self.assertRaises(ValueError, islice, ra, 1, 'a')
1212 self.assertRaises(ValueError, islice, ra, 'a', 1, 1)
1213 self.assertRaises(ValueError, islice, ra, 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +00001214 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001215
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001216 # Issue #10323: Less islice in a predictable state
1217 c = count()
1218 self.assertEqual(list(islice(c, 1, 3, 50)), [1])
1219 self.assertEqual(next(c), 3)
1220
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001221 # check copy, deepcopy, pickle
1222 for args in [ # islice(args) should agree with range(args)
1223 (10, 20, 3),
1224 (10, 3, 20),
1225 (10, 20),
1226 (10, 3),
1227 (20,)
1228 ]:
1229 self.assertEqual(list(copy.copy(islice(range(100), *args))),
1230 list(range(*args)))
1231 self.assertEqual(list(copy.deepcopy(islice(range(100), *args))),
1232 list(range(*args)))
Serhiy Storchakabad12572014-12-15 14:03:42 +02001233 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1234 self.pickletest(proto, islice(range(100), *args))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001235
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001236 # Issue #21321: check source iterator is not referenced
1237 # from islice() after the latter has been exhausted
1238 it = (x for x in (1, 2))
1239 wr = weakref.ref(it)
1240 it = islice(it, 1)
1241 self.assertIsNotNone(wr())
1242 list(it) # exhaust the iterator
Benjamin Peterson8e163512014-08-24 18:07:28 -05001243 support.gc_collect()
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001244 self.assertIsNone(wr())
1245
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001246 def test_takewhile(self):
1247 data = [1, 3, 5, 20, 2, 4, 6, 8]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001248 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001249 self.assertEqual(list(takewhile(underten, [])), [])
1250 self.assertRaises(TypeError, takewhile)
1251 self.assertRaises(TypeError, takewhile, operator.pow)
1252 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001253 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
1254 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001255 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
1256 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +00001257 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001258
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001259 # check copy, deepcopy, pickle
1260 self.assertEqual(list(copy.copy(takewhile(underten, data))), [1, 3, 5])
1261 self.assertEqual(list(copy.deepcopy(takewhile(underten, data))),
1262 [1, 3, 5])
Serhiy Storchakabad12572014-12-15 14:03:42 +02001263 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1264 self.pickletest(proto, takewhile(underten, data))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001265
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001266 def test_dropwhile(self):
1267 data = [1, 3, 5, 20, 2, 4, 6, 8]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001268 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001269 self.assertEqual(list(dropwhile(underten, [])), [])
1270 self.assertRaises(TypeError, dropwhile)
1271 self.assertRaises(TypeError, dropwhile, operator.pow)
1272 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001273 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
1274 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001275
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001276 # check copy, deepcopy, pickle
1277 self.assertEqual(list(copy.copy(dropwhile(underten, data))), [20, 2, 4, 6, 8])
1278 self.assertEqual(list(copy.deepcopy(dropwhile(underten, data))),
1279 [20, 2, 4, 6, 8])
Serhiy Storchakabad12572014-12-15 14:03:42 +02001280 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1281 self.pickletest(proto, dropwhile(underten, data))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001282
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001283 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +00001284 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001285
1286 a, b = tee([]) # test empty iterator
1287 self.assertEqual(list(a), [])
1288 self.assertEqual(list(b), [])
1289
1290 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +00001291 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001292
1293 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +00001294 self.assertEqual(list(a), list(range(n)))
1295 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001296
1297 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +00001298 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001299 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001300 del a
Guido van Rossum805365e2007-05-07 22:24:25 +00001301 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001302
1303 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +00001304 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001305 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001306 del b
Guido van Rossum805365e2007-05-07 22:24:25 +00001307 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001308
Guido van Rossum805365e2007-05-07 22:24:25 +00001309 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001310 order = [0]*n + [1]*n
1311 random.shuffle(order)
1312 lists = ([], [])
1313 its = tee(irange(n))
1314 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +00001315 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001316 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +00001317 self.assertEqual(lists[0], list(range(n)))
1318 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001319
Raymond Hettingerad983e72003-11-12 14:32:26 +00001320 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001321 self.assertRaises(TypeError, tee)
1322 self.assertRaises(TypeError, tee, 3)
1323 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +00001324 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001325
Raymond Hettingerad983e72003-11-12 14:32:26 +00001326 # tee object should be instantiable
1327 a, b = tee('abc')
1328 c = type(a)('def')
1329 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +00001330
Raymond Hettingerad983e72003-11-12 14:32:26 +00001331 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +00001332 a, b, c = tee(range(2000), 3)
1333 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001334 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +00001335 self.assertEqual(list(b), list(range(2000)))
1336 self.assertEqual([next(c), next(c)], list(range(2)))
1337 self.assertEqual(list(a), list(range(100,2000)))
1338 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +00001339
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001340 # test values of n
1341 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +00001342 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +00001343 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001344 result = tee('abc', n)
1345 self.assertEqual(type(result), tuple)
1346 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001347 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001348
Raymond Hettingerad983e72003-11-12 14:32:26 +00001349 # tee pass-through to copyable iterator
1350 a, b = tee('abc')
1351 c, d = tee(a)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +00001352 self.assertTrue(a is c)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001353
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001354 # test tee_new
1355 t1, t2 = tee('abc')
1356 tnew = type(t1)
1357 self.assertRaises(TypeError, tnew)
1358 self.assertRaises(TypeError, tnew, 10)
1359 t3 = tnew(t1)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +00001360 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +00001361
Raymond Hettingera9f60922004-10-17 16:40:14 +00001362 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +00001363 a, b = tee(range(10))
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001364 p = weakref.proxy(a)
Raymond Hettingera9f60922004-10-17 16:40:14 +00001365 self.assertEqual(getattr(p, '__class__'), type(b))
1366 del a
1367 self.assertRaises(ReferenceError, getattr, p, '__class__')
1368
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001369 ans = list('abc')
1370 long_ans = list(range(10000))
1371
1372 # check copy
1373 a, b = tee('abc')
1374 self.assertEqual(list(copy.copy(a)), ans)
1375 self.assertEqual(list(copy.copy(b)), ans)
1376 a, b = tee(list(range(10000)))
1377 self.assertEqual(list(copy.copy(a)), long_ans)
1378 self.assertEqual(list(copy.copy(b)), long_ans)
1379
1380 # check partially consumed copy
1381 a, b = tee('abc')
1382 take(2, a)
1383 take(1, b)
1384 self.assertEqual(list(copy.copy(a)), ans[2:])
1385 self.assertEqual(list(copy.copy(b)), ans[1:])
1386 self.assertEqual(list(a), ans[2:])
1387 self.assertEqual(list(b), ans[1:])
1388 a, b = tee(range(10000))
1389 take(100, a)
1390 take(60, b)
1391 self.assertEqual(list(copy.copy(a)), long_ans[100:])
1392 self.assertEqual(list(copy.copy(b)), long_ans[60:])
1393 self.assertEqual(list(a), long_ans[100:])
1394 self.assertEqual(list(b), long_ans[60:])
1395
1396 # check deepcopy
1397 a, b = tee('abc')
1398 self.assertEqual(list(copy.deepcopy(a)), ans)
1399 self.assertEqual(list(copy.deepcopy(b)), ans)
1400 self.assertEqual(list(a), ans)
1401 self.assertEqual(list(b), ans)
1402 a, b = tee(range(10000))
1403 self.assertEqual(list(copy.deepcopy(a)), long_ans)
1404 self.assertEqual(list(copy.deepcopy(b)), long_ans)
1405 self.assertEqual(list(a), long_ans)
1406 self.assertEqual(list(b), long_ans)
1407
1408 # check partially consumed deepcopy
1409 a, b = tee('abc')
1410 take(2, a)
1411 take(1, b)
1412 self.assertEqual(list(copy.deepcopy(a)), ans[2:])
1413 self.assertEqual(list(copy.deepcopy(b)), ans[1:])
1414 self.assertEqual(list(a), ans[2:])
1415 self.assertEqual(list(b), ans[1:])
1416 a, b = tee(range(10000))
1417 take(100, a)
1418 take(60, b)
1419 self.assertEqual(list(copy.deepcopy(a)), long_ans[100:])
1420 self.assertEqual(list(copy.deepcopy(b)), long_ans[60:])
1421 self.assertEqual(list(a), long_ans[100:])
1422 self.assertEqual(list(b), long_ans[60:])
1423
1424 # check pickle
Serhiy Storchakabad12572014-12-15 14:03:42 +02001425 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1426 self.pickletest(proto, iter(tee('abc')))
1427 a, b = tee('abc')
1428 self.pickletest(proto, a, compare=ans)
1429 self.pickletest(proto, b, compare=ans)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001430
Serhiy Storchakaa3e91282013-01-25 13:19:31 +02001431 # Issue 13454: Crash when deleting backward iterator from tee()
1432 def test_tee_del_backward(self):
Serhiy Storchaka5bb893c2013-01-26 11:52:06 +02001433 forward, backward = tee(repeat(None, 20000000))
Serhiy Storchaka9db55002015-03-28 20:38:37 +02001434 try:
1435 any(forward) # exhaust the iterator
1436 del backward
1437 except:
1438 del forward, backward
1439 raise
Serhiy Storchakaa3e91282013-01-25 13:19:31 +02001440
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001441 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001442 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +00001443
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001444 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +00001445 self.assertRaises(StopIteration, next, f([]))
1446 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001447
Georg Brandla18af4e2007-04-21 15:47:16 +00001448 self.assertRaises(StopIteration, next, islice([], None))
1449 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001450
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001451 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +00001452 self.assertRaises(StopIteration, next, p)
1453 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001454 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +00001455 self.assertRaises(StopIteration, next, p)
1456 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001457
Georg Brandla18af4e2007-04-21 15:47:16 +00001458 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001459
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001460 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +00001461 self.assertRaises(StopIteration, next, f(lambda x:x, []))
1462 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001463
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001464class TestExamples(unittest.TestCase):
1465
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001466 def test_accumulate(self):
Raymond Hettinger482ba772010-12-01 22:48:00 +00001467 self.assertEqual(list(accumulate([1,2,3,4,5])), [1, 3, 6, 10, 15])
1468
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001469 def test_accumulate_reducible(self):
1470 # check copy, deepcopy, pickle
1471 data = [1, 2, 3, 4, 5]
1472 accumulated = [1, 3, 6, 10, 15]
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001473
Serhiy Storchakabad12572014-12-15 14:03:42 +02001474 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1475 it = accumulate(data)
1476 self.assertEqual(list(pickle.loads(pickle.dumps(it, proto))), accumulated[:])
1477 self.assertEqual(next(it), 1)
1478 self.assertEqual(list(pickle.loads(pickle.dumps(it, proto))), accumulated[1:])
1479 it = accumulate(data)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001480 self.assertEqual(next(it), 1)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001481 self.assertEqual(list(copy.deepcopy(it)), accumulated[1:])
1482 self.assertEqual(list(copy.copy(it)), accumulated[1:])
1483
Serhiy Storchakad5516252016-03-06 14:00:45 +02001484 def test_accumulate_reducible_none(self):
1485 # Issue #25718: total is None
1486 it = accumulate([None, None, None], operator.is_)
1487 self.assertEqual(next(it), None)
1488 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1489 it_copy = pickle.loads(pickle.dumps(it, proto))
1490 self.assertEqual(list(it_copy), [True, False])
1491 self.assertEqual(list(copy.deepcopy(it)), [True, False])
1492 self.assertEqual(list(copy.copy(it)), [True, False])
1493
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001494 def test_chain(self):
1495 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
1496
1497 def test_chain_from_iterable(self):
1498 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
1499
1500 def test_combinations(self):
1501 self.assertEqual(list(combinations('ABCD', 2)),
1502 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
1503 self.assertEqual(list(combinations(range(4), 3)),
1504 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
1505
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001506 def test_combinations_with_replacement(self):
1507 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
1508 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
1509
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001510 def test_compress(self):
1511 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
1512
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001513 def test_count(self):
1514 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
1515
1516 def test_cycle(self):
1517 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
1518
1519 def test_dropwhile(self):
1520 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
1521
1522 def test_groupby(self):
1523 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
1524 list('ABCDAB'))
1525 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
1526 [list('AAAA'), list('BBB'), list('CC'), list('D')])
1527
1528 def test_filter(self):
1529 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
1530
1531 def test_filterfalse(self):
1532 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
1533
1534 def test_map(self):
1535 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
1536
1537 def test_islice(self):
1538 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
1539 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
1540 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
1541 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1542
1543 def test_zip(self):
1544 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
1545
1546 def test_zip_longest(self):
1547 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
1548 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
1549
1550 def test_permutations(self):
1551 self.assertEqual(list(permutations('ABCD', 2)),
1552 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
1553 self.assertEqual(list(permutations(range(3))),
1554 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
1555
1556 def test_product(self):
1557 self.assertEqual(list(product('ABCD', 'xy')),
1558 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
1559 self.assertEqual(list(product(range(2), repeat=3)),
1560 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
1561 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
1562
1563 def test_repeat(self):
1564 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
1565
1566 def test_stapmap(self):
1567 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
1568 [32, 9, 1000])
1569
1570 def test_takewhile(self):
1571 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
1572
1573
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001574class TestGC(unittest.TestCase):
1575
1576 def makecycle(self, iterator, container):
1577 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +00001578 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001579 del container, iterator
1580
Raymond Hettinger482ba772010-12-01 22:48:00 +00001581 def test_accumulate(self):
1582 a = []
1583 self.makecycle(accumulate([1,2,a,3]), a)
1584
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001585 def test_chain(self):
1586 a = []
1587 self.makecycle(chain(a), a)
1588
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001589 def test_chain_from_iterable(self):
1590 a = []
1591 self.makecycle(chain.from_iterable([a]), a)
1592
1593 def test_combinations(self):
1594 a = []
1595 self.makecycle(combinations([1,2,a,3], 3), a)
1596
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001597 def test_combinations_with_replacement(self):
1598 a = []
1599 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1600
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001601 def test_compress(self):
1602 a = []
1603 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1604
Raymond Hettingerd6280f42009-02-16 20:50:56 +00001605 def test_count(self):
1606 a = []
1607 Int = type('Int', (int,), dict(x=a))
1608 self.makecycle(count(Int(0), Int(1)), a)
1609
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001610 def test_cycle(self):
1611 a = []
1612 self.makecycle(cycle([a]*2), a)
1613
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001614 def test_dropwhile(self):
1615 a = []
1616 self.makecycle(dropwhile(bool, [0, a, a]), a)
1617
1618 def test_groupby(self):
1619 a = []
1620 self.makecycle(groupby([a]*2, lambda x:x), a)
1621
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001622 def test_issue2246(self):
1623 # Issue 2246 -- the _grouper iterator was not included in GC
1624 n = 10
1625 keyfunc = lambda x: x
1626 for i, j in groupby(range(n), key=keyfunc):
1627 keyfunc.__dict__.setdefault('x',[]).append(j)
1628
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001629 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001630 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001631 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001632
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001633 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001634 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001635 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001636
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001637 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001638 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001639 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001640
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001641 def test_zip_longest(self):
1642 a = []
1643 self.makecycle(zip_longest([a]*2, [a]*3), a)
1644 b = [a, None]
1645 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1646
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001647 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001648 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001649 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001650
1651 def test_islice(self):
1652 a = []
1653 self.makecycle(islice([a]*2, None), a)
1654
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001655 def test_permutations(self):
1656 a = []
1657 self.makecycle(permutations([1,2,a,3], 3), a)
1658
1659 def test_product(self):
1660 a = []
1661 self.makecycle(product([1,2,a,3], repeat=3), a)
1662
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001663 def test_repeat(self):
1664 a = []
1665 self.makecycle(repeat(a), a)
1666
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001667 def test_starmap(self):
1668 a = []
1669 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1670
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001671 def test_takewhile(self):
1672 a = []
1673 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1674
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001675def R(seqn):
1676 'Regular generator'
1677 for i in seqn:
1678 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001679
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001680class G:
1681 'Sequence using __getitem__'
1682 def __init__(self, seqn):
1683 self.seqn = seqn
1684 def __getitem__(self, i):
1685 return self.seqn[i]
1686
1687class I:
1688 'Sequence using iterator protocol'
1689 def __init__(self, seqn):
1690 self.seqn = seqn
1691 self.i = 0
1692 def __iter__(self):
1693 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001694 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001695 if self.i >= len(self.seqn): raise StopIteration
1696 v = self.seqn[self.i]
1697 self.i += 1
1698 return v
1699
1700class Ig:
1701 'Sequence using iterator protocol defined with a generator'
1702 def __init__(self, seqn):
1703 self.seqn = seqn
1704 self.i = 0
1705 def __iter__(self):
1706 for val in self.seqn:
1707 yield val
1708
1709class X:
1710 'Missing __getitem__ and __iter__'
1711 def __init__(self, seqn):
1712 self.seqn = seqn
1713 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001714 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001715 if self.i >= len(self.seqn): raise StopIteration
1716 v = self.seqn[self.i]
1717 self.i += 1
1718 return v
1719
1720class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001721 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001722 def __init__(self, seqn):
1723 self.seqn = seqn
1724 self.i = 0
1725 def __iter__(self):
1726 return self
1727
1728class E:
1729 'Test propagation of exceptions'
1730 def __init__(self, seqn):
1731 self.seqn = seqn
1732 self.i = 0
1733 def __iter__(self):
1734 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001735 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001736 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001737
1738class S:
1739 'Test immediate stop'
1740 def __init__(self, seqn):
1741 pass
1742 def __iter__(self):
1743 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001744 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001745 raise StopIteration
1746
1747def L(seqn):
1748 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001749 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001750
1751
1752class TestVariousIteratorArgs(unittest.TestCase):
1753
Raymond Hettinger482ba772010-12-01 22:48:00 +00001754 def test_accumulate(self):
1755 s = [1,2,3,4,5]
1756 r = [1,3,6,10,15]
1757 n = len(s)
1758 for g in (G, I, Ig, L, R):
1759 self.assertEqual(list(accumulate(g(s))), r)
1760 self.assertEqual(list(accumulate(S(s))), [])
1761 self.assertRaises(TypeError, accumulate, X(s))
1762 self.assertRaises(TypeError, accumulate, N(s))
1763 self.assertRaises(ZeroDivisionError, list, accumulate(E(s)))
1764
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001765 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001766 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001767 for g in (G, I, Ig, S, L, R):
1768 self.assertEqual(list(chain(g(s))), list(g(s)))
1769 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001770 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001771 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001772 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1773
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001774 def test_compress(self):
1775 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1776 n = len(s)
1777 for g in (G, I, Ig, S, L, R):
1778 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1779 self.assertRaises(TypeError, compress, X(s), repeat(1))
1780 self.assertRaises(TypeError, compress, N(s), repeat(1))
1781 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1782
Christian Heimesc3f30c42008-02-22 16:37:40 +00001783 def test_product(self):
1784 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1785 self.assertRaises(TypeError, product, X(s))
1786 self.assertRaises(TypeError, product, N(s))
1787 self.assertRaises(ZeroDivisionError, product, E(s))
1788
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001789 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001790 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001791 for g in (G, I, Ig, S, L, R):
1792 tgtlen = len(s) * 3
1793 expected = list(g(s))*3
1794 actual = list(islice(cycle(g(s)), tgtlen))
1795 self.assertEqual(actual, expected)
1796 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001797 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001798 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1799
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001800 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001801 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001802 for g in (G, I, Ig, S, L, R):
1803 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1804 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001805 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001806 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1807
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001808 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001809 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001810 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001811 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001812 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001813 self.assertRaises(TypeError, filter, isEven, X(s))
1814 self.assertRaises(TypeError, filter, isEven, N(s))
1815 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001816
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001817 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001818 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001819 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001820 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001821 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001822 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1823 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1824 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001825
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001826 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001827 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001828 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001829 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1830 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1831 self.assertRaises(TypeError, zip, X(s))
1832 self.assertRaises(TypeError, zip, N(s))
1833 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001834
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001835 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001836 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001837 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001838 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1839 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1840 self.assertRaises(TypeError, zip_longest, X(s))
1841 self.assertRaises(TypeError, zip_longest, N(s))
1842 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001843
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001844 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001845 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001846 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001847 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001848 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001849 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001850 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001851 self.assertRaises(TypeError, map, onearg, X(s))
1852 self.assertRaises(TypeError, map, onearg, N(s))
1853 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001854
1855 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001856 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001857 for g in (G, I, Ig, S, L, R):
1858 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1859 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001860 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001861 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1862
1863 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001864 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001865 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001866 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001867 self.assertEqual(list(starmap(operator.pow, g(ss))),
1868 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001869 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001870 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001871 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1872
1873 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001874 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001875 for g in (G, I, Ig, S, L, R):
1876 tgt = []
1877 for elem in g(s):
1878 if not isEven(elem): break
1879 tgt.append(elem)
1880 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1881 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001882 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001883 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1884
1885 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001886 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001887 for g in (G, I, Ig, S, L, R):
1888 tgt = []
1889 for elem in g(s):
1890 if not tgt and isOdd(elem): continue
1891 tgt.append(elem)
1892 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1893 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001894 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001895 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1896
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001897 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001898 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001899 for g in (G, I, Ig, S, L, R):
1900 it1, it2 = tee(g(s))
1901 self.assertEqual(list(it1), list(g(s)))
1902 self.assertEqual(list(it2), list(g(s)))
1903 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001904 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001905 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1906
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001907class LengthTransparency(unittest.TestCase):
1908
1909 def test_repeat(self):
Armin Ronacheraa9a79d2012-10-06 14:03:24 +02001910 self.assertEqual(operator.length_hint(repeat(None, 50)), 50)
Raymond Hettinger97d35552014-06-24 21:36:58 -07001911 self.assertEqual(operator.length_hint(repeat(None, 0)), 0)
Armin Ronacheraa9a79d2012-10-06 14:03:24 +02001912 self.assertEqual(operator.length_hint(repeat(None), 12), 12)
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001913
Raymond Hettinger97d35552014-06-24 21:36:58 -07001914 def test_repeat_with_negative_times(self):
1915 self.assertEqual(operator.length_hint(repeat(None, -1)), 0)
1916 self.assertEqual(operator.length_hint(repeat(None, -2)), 0)
1917 self.assertEqual(operator.length_hint(repeat(None, times=-1)), 0)
1918 self.assertEqual(operator.length_hint(repeat(None, times=-2)), 0)
1919
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001920class RegressionTests(unittest.TestCase):
1921
1922 def test_sf_793826(self):
1923 # Fix Armin Rigo's successful efforts to wreak havoc
1924
1925 def mutatingtuple(tuple1, f, tuple2):
1926 # this builds a tuple t which is a copy of tuple1,
1927 # then calls f(t), then mutates t to be equal to tuple2
1928 # (needs len(tuple1) == len(tuple2)).
1929 def g(value, first=[1]):
1930 if first:
1931 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00001932 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001933 return value
1934 items = list(tuple2)
1935 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001936 gen = map(g, items)
1937 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00001938 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001939
1940 def f(t):
1941 global T
1942 T = t
1943 first[:] = list(T)
1944
1945 first = []
1946 mutatingtuple((1,2,3), f, (4,5,6))
1947 second = list(T)
1948 self.assertEqual(first, second)
1949
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001950
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001951 def test_sf_950057(self):
1952 # Make sure that chain() and cycle() catch exceptions immediately
1953 # rather than when shifting between input sources
1954
1955 def gen1():
1956 hist.append(0)
1957 yield 1
1958 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001959 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001960 hist.append(2)
1961
1962 def gen2(x):
1963 hist.append(3)
1964 yield 2
1965 hist.append(4)
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001966
1967 hist = []
1968 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1969 self.assertEqual(hist, [0,1])
1970
1971 hist = []
1972 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1973 self.assertEqual(hist, [0,1])
1974
1975 hist = []
1976 self.assertRaises(AssertionError, list, cycle(gen1()))
1977 self.assertEqual(hist, [0,1])
1978
Thomas Woutersb2137042007-02-01 18:02:27 +00001979class SubclassWithKwargsTest(unittest.TestCase):
1980 def test_keywords_in_subclass(self):
1981 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001982 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001983 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00001984 class Subclass(cls):
1985 def __init__(self, newarg=None, *args):
1986 cls.__init__(self, *args)
1987 try:
1988 Subclass(newarg=1)
1989 except TypeError as err:
1990 # we expect type errors because of wrong argument count
Ezio Melottib58e0bd2010-01-23 15:40:09 +00001991 self.assertNotIn("does not take keyword arguments", err.args[0])
Thomas Woutersb2137042007-02-01 18:02:27 +00001992
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02001993@support.cpython_only
1994class SizeofTest(unittest.TestCase):
1995 def setUp(self):
1996 self.ssize_t = struct.calcsize('n')
1997
1998 check_sizeof = support.check_sizeof
1999
2000 def test_product_sizeof(self):
2001 basesize = support.calcobjsize('3Pi')
2002 check = self.check_sizeof
2003 check(product('ab', '12'), basesize + 2 * self.ssize_t)
2004 check(product(*(('abc',) * 10)), basesize + 10 * self.ssize_t)
2005
2006 def test_combinations_sizeof(self):
2007 basesize = support.calcobjsize('3Pni')
2008 check = self.check_sizeof
2009 check(combinations('abcd', 3), basesize + 3 * self.ssize_t)
2010 check(combinations(range(10), 4), basesize + 4 * self.ssize_t)
2011
2012 def test_combinations_with_replacement_sizeof(self):
2013 cwr = combinations_with_replacement
2014 basesize = support.calcobjsize('3Pni')
2015 check = self.check_sizeof
2016 check(cwr('abcd', 3), basesize + 3 * self.ssize_t)
2017 check(cwr(range(10), 4), basesize + 4 * self.ssize_t)
2018
2019 def test_permutations_sizeof(self):
2020 basesize = support.calcobjsize('4Pni')
2021 check = self.check_sizeof
2022 check(permutations('abcd'),
2023 basesize + 4 * self.ssize_t + 4 * self.ssize_t)
2024 check(permutations('abcd', 3),
2025 basesize + 4 * self.ssize_t + 3 * self.ssize_t)
2026 check(permutations('abcde', 3),
2027 basesize + 5 * self.ssize_t + 3 * self.ssize_t)
2028 check(permutations(range(10), 4),
2029 basesize + 10 * self.ssize_t + 4 * self.ssize_t)
2030
Thomas Woutersb2137042007-02-01 18:02:27 +00002031
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002032libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002033
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002034
2035>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002036>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00002037... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00002038...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002039Check 1200 is for $120.15
2040Check 1201 is for $764.05
2041Check 1202 is for $823.14
2042
2043>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002044>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00002045... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00002046...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000020471
20488
204927
2050
2051>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00002052>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00002053... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00002054...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002055Alex
2056Laura
2057Martin
2058Walter
2059Samuele
2060
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002061>>> from operator import itemgetter
2062>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00002063>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002064>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00002065... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00002066...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000020671 ['a', 'c', 'e']
20682 ['b', 'd', 'f']
20693 ['g']
2070
Raymond Hettinger734fb572004-01-20 20:04:40 +00002071# Find runs of consecutive numbers using groupby. The key to the solution
2072# is differencing with a range so that consecutive numbers all appear in
2073# same group.
2074>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00002075>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00002076... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00002077...
Raymond Hettinger734fb572004-01-20 20:04:40 +00002078[1]
2079[4, 5, 6]
2080[10]
2081[15, 16, 17, 18]
2082[22]
2083[25, 26, 27, 28]
2084
Georg Brandl3dbca812008-07-23 16:10:53 +00002085>>> def take(n, iterable):
2086... "Return first n items of the iterable as a list"
2087... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00002088
Georg Brandl3dbca812008-07-23 16:10:53 +00002089>>> def enumerate(iterable, start=0):
2090... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002091
Georg Brandl3dbca812008-07-23 16:10:53 +00002092>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002093... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00002094... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002095
Raymond Hettingercdf8ba32009-02-19 04:45:07 +00002096>>> def nth(iterable, n, default=None):
2097... "Returns the nth item or a default value"
2098... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002099
Raymond Hettingere525ee32016-03-06 18:11:38 -08002100>>> def all_equal(iterable):
2101... "Returns True if all the elements are equal to each other"
2102... g = groupby(iterable)
2103... return next(g, True) and not next(g, False)
2104
Georg Brandl3dbca812008-07-23 16:10:53 +00002105>>> def quantify(iterable, pred=bool):
2106... "Count how many times the predicate is true"
2107... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00002108
Georg Brandl3dbca812008-07-23 16:10:53 +00002109>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002110... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00002111... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002112
Georg Brandl3dbca812008-07-23 16:10:53 +00002113>>> def ncycles(iterable, n):
Ezio Melotti13925002011-03-16 11:05:33 +02002114... "Returns the sequence elements n times"
Georg Brandl3dbca812008-07-23 16:10:53 +00002115... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002116
2117>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002118... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002119
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002120>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002121... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002122
2123>>> def repeatfunc(func, times=None, *args):
2124... "Repeat calls to func with specified arguments."
2125... " Example: repeatfunc(random.random)"
2126... if times is None:
2127... return starmap(func, repeat(args))
2128... else:
2129... return starmap(func, repeat(args, times))
2130
Raymond Hettingerd591f662003-10-26 15:34:50 +00002131>>> def pairwise(iterable):
2132... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
2133... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00002134... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00002135... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00002136... except StopIteration:
2137... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002138... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002139
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002140>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00002141... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002142... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00002143... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002144
2145>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00002146... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002147... # Recipe credited to George Sakkis
2148... pending = len(iterables)
2149... nexts = cycle(iter(it).__next__ for it in iterables)
2150... while pending:
2151... try:
2152... for next in nexts:
2153... yield next()
2154... except StopIteration:
2155... pending -= 1
2156... nexts = cycle(islice(nexts, pending))
2157
2158>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00002159... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
2160... s = list(iterable)
2161... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002162
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00002163>>> def unique_everseen(iterable, key=None):
2164... "List unique elements, preserving order. Remember all elements ever seen."
2165... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
2166... # unique_everseen('ABBCcAD', str.lower) --> A B C D
2167... seen = set()
2168... seen_add = seen.add
2169... if key is None:
2170... for element in iterable:
2171... if element not in seen:
2172... seen_add(element)
2173... yield element
2174... else:
2175... for element in iterable:
2176... k = key(element)
2177... if k not in seen:
2178... seen_add(k)
2179... yield element
2180
2181>>> def unique_justseen(iterable, key=None):
2182... "List unique elements, preserving order. Remember only the element just seen."
2183... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
2184... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
2185... return map(next, map(itemgetter(1), groupby(iterable, key)))
2186
Raymond Hettinger31b26f62014-04-02 03:16:42 -07002187>>> def first_true(iterable, default=False, pred=None):
2188... '''Returns the first true value in the iterable.
2189...
2190... If no true value is found, returns *default*
2191...
2192... If *pred* is not None, returns the first item
2193... for which pred(item) is true.
2194...
2195... '''
2196... # first_true([a,b,c], x) --> a or b or c or x
2197... # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x
2198... return next(filter(pred, iterable), default)
2199
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002200This is not part of the examples but it tests to make sure the definitions
2201perform as purported.
2202
Raymond Hettingera098b332003-09-08 23:58:40 +00002203>>> take(10, count())
2204[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
2205
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002206>>> list(enumerate('abc'))
2207[(0, 'a'), (1, 'b'), (2, 'c')]
2208
2209>>> list(islice(tabulate(lambda x: 2*x), 4))
2210[0, 2, 4, 6]
2211
2212>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00002213'd'
2214
2215>>> nth('abcde', 9) is None
2216True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002217
Raymond Hettingere525ee32016-03-06 18:11:38 -08002218>>> [all_equal(s) for s in ('', 'A', 'AAAA', 'AAAB', 'AAABA')]
2219[True, True, True, False, False]
2220
Guido van Rossum805365e2007-05-07 22:24:25 +00002221>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000222250
2223
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002224>>> a = [[1, 2, 3], [4, 5, 6]]
2225>>> flatten(a)
2226[1, 2, 3, 4, 5, 6]
2227
2228>>> list(repeatfunc(pow, 5, 2, 3))
2229[8, 8, 8, 8, 8]
2230
2231>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002232>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002233[0, 0, 0, 0, 0]
2234
Raymond Hettingerd591f662003-10-26 15:34:50 +00002235>>> list(pairwise('abcd'))
2236[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002237
Raymond Hettingerd591f662003-10-26 15:34:50 +00002238>>> list(pairwise([]))
2239[]
2240
2241>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00002242[]
2243
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002244>>> list(islice(padnone('abc'), 0, 6))
2245['a', 'b', 'c', None, None, None]
2246
2247>>> list(ncycles('abc', 3))
2248['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
2249
2250>>> dotproduct([1,2,3], [4,5,6])
225132
2252
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002253>>> list(grouper(3, 'abcdefg', 'x'))
2254[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
2255
2256>>> list(roundrobin('abc', 'd', 'ef'))
2257['a', 'd', 'e', 'b', 'f', 'c']
2258
Raymond Hettingerace67332009-01-26 02:23:50 +00002259>>> list(powerset([1,2,3]))
2260[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002261
Raymond Hettinger191e8502009-01-27 13:29:43 +00002262>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
2263True
2264
2265>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
2266True
2267
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00002268>>> list(unique_everseen('AAAABBBCCDAABBB'))
2269['A', 'B', 'C', 'D']
2270
2271>>> list(unique_everseen('ABBCcAD', str.lower))
2272['A', 'B', 'C', 'D']
2273
2274>>> list(unique_justseen('AAAABBBCCDAABBB'))
2275['A', 'B', 'C', 'D', 'A', 'B']
2276
2277>>> list(unique_justseen('ABBCcAD', str.lower))
2278['A', 'B', 'C', 'A', 'D']
2279
Raymond Hettinger31b26f62014-04-02 03:16:42 -07002280>>> first_true('ABC0DEF1', '9', str.isdigit)
2281'0'
2282
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002283"""
2284
2285__test__ = {'libreftest' : libreftest}
2286
2287def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002288 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00002289 RegressionTests, LengthTransparency,
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002290 SubclassWithKwargsTest, TestExamples,
2291 SizeofTest)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002292 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002293
2294 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002295 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00002296 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00002297 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00002298 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002299 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00002300 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00002301 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00002302 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002303
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002304 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002305 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002306
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002307if __name__ == "__main__":
2308 test_main(verbose=True)