blob: 98b8c83731899a97e554e09773303f5c21357707 [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
Lisa Roach9718b592018-09-23 17:34:59 -0700150 self.pickletest(proto, accumulate(range(10), initial=7))
151 self.assertEqual(list(accumulate([10, 5, 1], initial=None)), [10, 15, 16])
152 self.assertEqual(list(accumulate([10, 5, 1], initial=100)), [100, 110, 115, 116])
153 self.assertEqual(list(accumulate([], initial=100)), [100])
154 with self.assertRaises(TypeError):
155 list(accumulate([10, 20], 100))
Raymond Hettinger5d446132011-03-27 18:52:10 -0700156
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000157 def test_chain(self):
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000158
159 def chain2(*iterables):
160 'Pure python version in the docs'
161 for it in iterables:
162 for element in it:
163 yield element
164
165 for c in (chain, chain2):
166 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
167 self.assertEqual(list(c('abc')), list('abc'))
168 self.assertEqual(list(c('')), [])
169 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
170 self.assertRaises(TypeError, list,c(2, 3))
Christian Heimesf16baeb2008-02-29 14:57:44 +0000171
172 def test_chain_from_iterable(self):
173 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
174 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
175 self.assertEqual(list(chain.from_iterable([''])), [])
176 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
177 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000178
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000179 def test_chain_reducible(self):
Serhiy Storchakabad12572014-12-15 14:03:42 +0200180 for oper in [copy.deepcopy] + picklecopiers:
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000181 it = chain('abc', 'def')
182 self.assertEqual(list(oper(it)), list('abcdef'))
183 self.assertEqual(next(it), 'a')
184 self.assertEqual(list(oper(it)), list('bcdef'))
185
186 self.assertEqual(list(oper(chain(''))), [])
187 self.assertEqual(take(4, oper(chain('abc', 'def'))), list('abcd'))
188 self.assertRaises(TypeError, list, oper(chain(2, 3)))
Serhiy Storchakabad12572014-12-15 14:03:42 +0200189 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
190 self.pickletest(proto, chain('abc', 'def'), compare=list('abcdef'))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000191
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300192 def test_chain_setstate(self):
193 self.assertRaises(TypeError, chain().__setstate__, ())
194 self.assertRaises(TypeError, chain().__setstate__, [])
195 self.assertRaises(TypeError, chain().__setstate__, 0)
196 self.assertRaises(TypeError, chain().__setstate__, ([],))
197 self.assertRaises(TypeError, chain().__setstate__, (iter([]), []))
198 it = chain()
199 it.__setstate__((iter(['abc', 'def']),))
200 self.assertEqual(list(it), ['a', 'b', 'c', 'd', 'e', 'f'])
201 it = chain()
202 it.__setstate__((iter(['abc', 'def']), iter(['ghi'])))
203 self.assertEqual(list(it), ['ghi', 'a', 'b', 'c', 'd', 'e', 'f'])
204
Christian Heimes380f7f22008-02-28 11:19:05 +0000205 def test_combinations(self):
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000206 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Christian Heimes380f7f22008-02-28 11:19:05 +0000207 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Christian Heimes78644762008-03-04 23:39:23 +0000208 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Christian Heimes380f7f22008-02-28 11:19:05 +0000209 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000210
Serhiy Storchakabad12572014-12-15 14:03:42 +0200211 for op in [lambda a:a] + picklecopiers:
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000212 self.assertEqual(list(op(combinations('abc', 32))), []) # r > n
213
214 self.assertEqual(list(op(combinations('ABCD', 2))),
215 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
216 testIntermediate = combinations('ABCD', 2)
217 next(testIntermediate)
218 self.assertEqual(list(op(testIntermediate)),
219 [('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
220
221 self.assertEqual(list(op(combinations(range(4), 3))),
222 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
223 testIntermediate = combinations(range(4), 3)
224 next(testIntermediate)
225 self.assertEqual(list(op(testIntermediate)),
226 [(0,1,3), (0,2,3), (1,2,3)])
227
Christian Heimes78644762008-03-04 23:39:23 +0000228
229 def combinations1(iterable, r):
230 'Pure python version shown in the docs'
231 pool = tuple(iterable)
232 n = len(pool)
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000233 if r > n:
234 return
Christian Heimes78644762008-03-04 23:39:23 +0000235 indices = list(range(r))
236 yield tuple(pool[i] for i in indices)
237 while 1:
238 for i in reversed(range(r)):
239 if indices[i] != i + n - r:
240 break
241 else:
242 return
243 indices[i] += 1
244 for j in range(i+1, r):
245 indices[j] = indices[j-1] + 1
246 yield tuple(pool[i] for i in indices)
247
248 def combinations2(iterable, r):
249 'Pure python version shown in the docs'
250 pool = tuple(iterable)
251 n = len(pool)
252 for indices in permutations(range(n), r):
253 if sorted(indices) == list(indices):
254 yield tuple(pool[i] for i in indices)
255
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000256 def combinations3(iterable, r):
257 'Pure python version from cwr()'
258 pool = tuple(iterable)
259 n = len(pool)
260 for indices in combinations_with_replacement(range(n), r):
261 if len(set(indices)) == r:
262 yield tuple(pool[i] for i in indices)
263
Christian Heimes78644762008-03-04 23:39:23 +0000264 for n in range(7):
Christian Heimes380f7f22008-02-28 11:19:05 +0000265 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000266 for r in range(n+2):
Christian Heimes380f7f22008-02-28 11:19:05 +0000267 result = list(combinations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000268 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 +0000269 self.assertEqual(len(result), len(set(result))) # no repeats
270 self.assertEqual(result, sorted(result)) # lexicographic order
271 for c in result:
272 self.assertEqual(len(c), r) # r-length combinations
273 self.assertEqual(len(set(c)), r) # no duplicate elements
274 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000275 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000276 self.assertEqual(list(c),
277 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000278 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000279 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000280 self.assertEqual(result, list(combinations3(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000281
Serhiy Storchakabad12572014-12-15 14:03:42 +0200282 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
283 self.pickletest(proto, combinations(values, r)) # test pickling
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000284
Benjamin Peterson4b40eeb2015-02-01 20:59:00 -0500285 @support.bigaddrspacetest
286 def test_combinations_overflow(self):
Serhiy Storchakadee948b2015-02-03 01:34:09 +0200287 with self.assertRaises((OverflowError, MemoryError)):
Benjamin Peterson4b40eeb2015-02-01 20:59:00 -0500288 combinations("AA", 2**29)
289
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000290 # Test implementation detail: tuple re-use
Alex Gaynore151d212011-07-17 16:21:30 -0700291 @support.impl_detail("tuple reuse is specific to CPython")
292 def test_combinations_tuple_reuse(self):
Christian Heimes78644762008-03-04 23:39:23 +0000293 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
294 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
295
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000296 def test_combinations_with_replacement(self):
297 cwr = combinations_with_replacement
298 self.assertRaises(TypeError, cwr, 'abc') # missing r argument
299 self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
300 self.assertRaises(TypeError, cwr, None) # pool is not iterable
301 self.assertRaises(ValueError, cwr, 'abc', -2) # r is negative
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000302
Serhiy Storchakabad12572014-12-15 14:03:42 +0200303 for op in [lambda a:a] + picklecopiers:
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000304 self.assertEqual(list(op(cwr('ABC', 2))),
305 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
306 testIntermediate = cwr('ABC', 2)
307 next(testIntermediate)
308 self.assertEqual(list(op(testIntermediate)),
309 [('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
310
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000311
312 def cwr1(iterable, r):
313 'Pure python version shown in the docs'
314 # number items returned: (n+r-1)! / r! / (n-1)! when n>0
315 pool = tuple(iterable)
316 n = len(pool)
317 if not n and r:
318 return
319 indices = [0] * r
320 yield tuple(pool[i] for i in indices)
321 while 1:
322 for i in reversed(range(r)):
323 if indices[i] != n - 1:
324 break
325 else:
326 return
327 indices[i:] = [indices[i] + 1] * (r - i)
328 yield tuple(pool[i] for i in indices)
329
330 def cwr2(iterable, r):
331 'Pure python version shown in the docs'
332 pool = tuple(iterable)
333 n = len(pool)
334 for indices in product(range(n), repeat=r):
335 if sorted(indices) == list(indices):
336 yield tuple(pool[i] for i in indices)
337
338 def numcombs(n, r):
339 if not n:
340 return 0 if r else 1
341 return fact(n+r-1) / fact(r)/ fact(n-1)
342
343 for n in range(7):
344 values = [5*x-12 for x in range(n)]
345 for r in range(n+2):
346 result = list(cwr(values, r))
347
348 self.assertEqual(len(result), numcombs(n, r)) # right number of combs
349 self.assertEqual(len(result), len(set(result))) # no repeats
350 self.assertEqual(result, sorted(result)) # lexicographic order
351
352 regular_combs = list(combinations(values, r)) # compare to combs without replacement
353 if n == 0 or r <= 1:
Ezio Melottib3aedd42010-11-20 19:04:17 +0000354 self.assertEqual(result, regular_combs) # cases that should be identical
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000355 else:
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000356 self.assertTrue(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000357
358 for c in result:
359 self.assertEqual(len(c), r) # r-length combinations
360 noruns = [k for k,v in groupby(c)] # combo without consecutive repeats
361 self.assertEqual(len(noruns), len(set(noruns))) # no repeats other than consecutive
362 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000363 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000364 self.assertEqual(noruns,
365 [e for e in values if e in c]) # comb is a subsequence of the input iterable
366 self.assertEqual(result, list(cwr1(values, r))) # matches first pure python version
367 self.assertEqual(result, list(cwr2(values, r))) # matches second pure python version
368
Serhiy Storchakabad12572014-12-15 14:03:42 +0200369 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
370 self.pickletest(proto, cwr(values,r)) # test pickling
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000371
Benjamin Peterson6f082292015-02-01 21:10:47 -0500372 @support.bigaddrspacetest
373 def test_combinations_with_replacement_overflow(self):
Serhiy Storchakadee948b2015-02-03 01:34:09 +0200374 with self.assertRaises((OverflowError, MemoryError)):
Benjamin Peterson6f082292015-02-01 21:10:47 -0500375 combinations_with_replacement("AA", 2**30)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000376
Benjamin Peterson6f082292015-02-01 21:10:47 -0500377 # Test implementation detail: tuple re-use
Alex Gaynore151d212011-07-17 16:21:30 -0700378 @support.impl_detail("tuple reuse is specific to CPython")
379 def test_combinations_with_replacement_tuple_reuse(self):
380 cwr = combinations_with_replacement
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000381 self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
382 self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
383
Christian Heimes78644762008-03-04 23:39:23 +0000384 def test_permutations(self):
385 self.assertRaises(TypeError, permutations) # too few arguments
386 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000387 self.assertRaises(TypeError, permutations, None) # pool is not iterable
388 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000389 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000390 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Christian Heimes78644762008-03-04 23:39:23 +0000391 self.assertEqual(list(permutations(range(3), 2)),
392 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
393
394 def permutations1(iterable, r=None):
395 'Pure python version shown in the docs'
396 pool = tuple(iterable)
397 n = len(pool)
398 r = n if r is None else r
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000399 if r > n:
400 return
Christian Heimes78644762008-03-04 23:39:23 +0000401 indices = list(range(n))
402 cycles = list(range(n-r+1, n+1))[::-1]
403 yield tuple(pool[i] for i in indices[:r])
404 while n:
405 for i in reversed(range(r)):
406 cycles[i] -= 1
407 if cycles[i] == 0:
408 indices[i:] = indices[i+1:] + indices[i:i+1]
409 cycles[i] = n - i
410 else:
411 j = cycles[i]
412 indices[i], indices[-j] = indices[-j], indices[i]
413 yield tuple(pool[i] for i in indices[:r])
414 break
415 else:
416 return
417
418 def permutations2(iterable, r=None):
419 'Pure python version shown in the docs'
420 pool = tuple(iterable)
421 n = len(pool)
422 r = n if r is None else r
423 for indices in product(range(n), repeat=r):
424 if len(set(indices)) == r:
425 yield tuple(pool[i] for i in indices)
426
427 for n in range(7):
428 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000429 for r in range(n+2):
Christian Heimes78644762008-03-04 23:39:23 +0000430 result = list(permutations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000431 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 +0000432 self.assertEqual(len(result), len(set(result))) # no repeats
433 self.assertEqual(result, sorted(result)) # lexicographic order
434 for p in result:
435 self.assertEqual(len(p), r) # r-length permutations
436 self.assertEqual(len(set(p)), r) # no duplicate elements
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000437 self.assertTrue(all(e in values for e in p)) # elements taken from input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000438 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000439 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000440 if r == n:
441 self.assertEqual(result, list(permutations(values, None))) # test r as None
442 self.assertEqual(result, list(permutations(values))) # test default r
443
Serhiy Storchakabad12572014-12-15 14:03:42 +0200444 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
445 self.pickletest(proto, permutations(values, r)) # test pickling
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000446
Benjamin Peterson0eaabf12015-02-01 21:34:07 -0500447 @support.bigaddrspacetest
448 def test_permutations_overflow(self):
Serhiy Storchakadee948b2015-02-03 01:34:09 +0200449 with self.assertRaises((OverflowError, MemoryError)):
Benjamin Peterson0eaabf12015-02-01 21:34:07 -0500450 permutations("A", 2**30)
Benjamin Peterson0eaabf12015-02-01 21:34:07 -0500451
Zachary Waredca807b2014-04-24 13:22:05 -0500452 @support.impl_detail("tuple reuse is specific to CPython")
Alex Gaynore151d212011-07-17 16:21:30 -0700453 def test_permutations_tuple_reuse(self):
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000454 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Christian Heimes78644762008-03-04 23:39:23 +0000455 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Christian Heimes380f7f22008-02-28 11:19:05 +0000456
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000457 def test_combinatorics(self):
458 # Test relationships between product(), permutations(),
459 # combinations() and combinations_with_replacement().
460
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000461 for n in range(6):
462 s = 'ABCDEFG'[:n]
463 for r in range(8):
464 prod = list(product(s, repeat=r))
465 cwr = list(combinations_with_replacement(s, r))
466 perm = list(permutations(s, r))
467 comb = list(combinations(s, r))
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000468
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000469 # Check size
Ezio Melottib3aedd42010-11-20 19:04:17 +0000470 self.assertEqual(len(prod), n**r)
471 self.assertEqual(len(cwr), (fact(n+r-1) / fact(r)/ fact(n-1)) if n else (not r))
472 self.assertEqual(len(perm), 0 if r>n else fact(n) / fact(n-r))
473 self.assertEqual(len(comb), 0 if r>n else fact(n) / fact(r) / fact(n-r))
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000474
475 # Check lexicographic order without repeated tuples
Ezio Melottib3aedd42010-11-20 19:04:17 +0000476 self.assertEqual(prod, sorted(set(prod)))
477 self.assertEqual(cwr, sorted(set(cwr)))
478 self.assertEqual(perm, sorted(set(perm)))
479 self.assertEqual(comb, sorted(set(comb)))
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000480
481 # Check interrelationships
Ezio Melottib3aedd42010-11-20 19:04:17 +0000482 self.assertEqual(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
483 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 +0000484 self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
485 self.assertEqual(comb, [t for t in cwr if len(set(t))==r]) # comb: cwrs without dups
486 self.assertEqual(comb, list(filter(set(cwr).__contains__, perm))) # comb: perm that is a cwr
487 self.assertEqual(comb, list(filter(set(perm).__contains__, cwr))) # comb: cwr that is a perm
488 self.assertEqual(comb, sorted(set(cwr) & set(perm))) # comb: both a cwr and a perm
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000489
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000490 def test_compress(self):
Raymond Hettinger15a49502009-02-19 02:17:09 +0000491 self.assertEqual(list(compress(data='ABCDEF', selectors=[1,0,1,0,1,1])), list('ACEF'))
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000492 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
493 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
494 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
495 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
496 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
497 n = 10000
498 data = chain.from_iterable(repeat(range(6), n))
499 selectors = chain.from_iterable(repeat((0, 1)))
500 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
501 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
502 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
503 self.assertRaises(TypeError, compress, range(6)) # too few args
504 self.assertRaises(TypeError, compress, range(6), None) # too many args
505
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000506 # check copy, deepcopy, pickle
Serhiy Storchakabad12572014-12-15 14:03:42 +0200507 for op in [lambda a:copy.copy(a), lambda a:copy.deepcopy(a)] + picklecopiers:
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000508 for data, selectors, result1, result2 in [
509 ('ABCDEF', [1,0,1,0,1,1], 'ACEF', 'CEF'),
510 ('ABCDEF', [0,0,0,0,0,0], '', ''),
511 ('ABCDEF', [1,1,1,1,1,1], 'ABCDEF', 'BCDEF'),
512 ('ABCDEF', [1,0,1], 'AC', 'C'),
513 ('ABC', [0,1,1,1,1,1], 'BC', 'C'),
514 ]:
515
516 self.assertEqual(list(op(compress(data=data, selectors=selectors))), list(result1))
517 self.assertEqual(list(op(compress(data, selectors))), list(result1))
518 testIntermediate = compress(data, selectors)
519 if result1:
520 next(testIntermediate)
521 self.assertEqual(list(op(testIntermediate)), list(result2))
522
523
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000524 def test_count(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000525 self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
526 self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
527 self.assertEqual(take(2, lzip('abc',count(3))), [('a', 3), ('b', 4)])
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000528 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
529 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000530 self.assertRaises(TypeError, count, 2, 3, 4)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000531 self.assertRaises(TypeError, count, 'a')
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300532 self.assertEqual(take(10, count(maxsize-5)),
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000533 list(range(maxsize-5, maxsize+5)))
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300534 self.assertEqual(take(10, count(-maxsize-5)),
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000535 list(range(-maxsize-5, -maxsize+5)))
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300536 self.assertEqual(take(3, count(3.25)), [3.25, 4.25, 5.25])
537 self.assertEqual(take(3, count(3.25-4j)), [3.25-4j, 4.25-4j, 5.25-4j])
538 self.assertEqual(take(3, count(Decimal('1.1'))),
539 [Decimal('1.1'), Decimal('2.1'), Decimal('3.1')])
540 self.assertEqual(take(3, count(Fraction(2, 3))),
541 [Fraction(2, 3), Fraction(5, 3), Fraction(8, 3)])
542 BIGINT = 1<<1000
543 self.assertEqual(take(3, count(BIGINT)), [BIGINT, BIGINT+1, BIGINT+2])
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000544 c = count(3)
545 self.assertEqual(repr(c), 'count(3)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000546 next(c)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000547 self.assertEqual(repr(c), 'count(4)')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000548 c = count(-9)
549 self.assertEqual(repr(c), 'count(-9)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000550 next(c)
551 self.assertEqual(next(c), -8)
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300552 self.assertEqual(repr(count(10.25)), 'count(10.25)')
553 self.assertEqual(repr(count(10.0)), 'count(10.0)')
554 self.assertEqual(type(next(count(10.0))), float)
Christian Heimesa37d4c62007-12-04 23:02:19 +0000555 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 +0300556 # Test repr
557 r1 = repr(count(i))
558 r2 = 'count(%r)'.__mod__(i)
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000559 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000560
Raymond Hettingera3e1ad22009-11-30 22:02:31 +0000561 # check copy, deepcopy, pickle
562 for value in -3, 3, maxsize-5, maxsize+5:
563 c = count(value)
564 self.assertEqual(next(copy.copy(c)), value)
565 self.assertEqual(next(copy.deepcopy(c)), value)
Serhiy Storchakabad12572014-12-15 14:03:42 +0200566 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
567 self.pickletest(proto, count(value))
Raymond Hettingera3e1ad22009-11-30 22:02:31 +0000568
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +0000569 #check proper internal error handling for large "step' sizes
570 count(1, maxsize+5); sys.exc_info()
571
Raymond Hettinger30729212009-02-12 06:28:27 +0000572 def test_count_with_stride(self):
573 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingereb13fdd2009-02-21 22:30:12 +0000574 self.assertEqual(lzip('abc',count(start=2,step=3)),
575 [('a', 2), ('b', 5), ('c', 8)])
576 self.assertEqual(lzip('abc',count(step=-1)),
577 [('a', 0), ('b', -1), ('c', -2)])
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300578 self.assertRaises(TypeError, count, 'a', 'b')
Raymond Hettinger30729212009-02-12 06:28:27 +0000579 self.assertEqual(lzip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
580 self.assertEqual(lzip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000581 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000582 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
583 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300584 self.assertEqual(take(3, count(10, maxsize+5)),
585 list(range(10, 10+3*(maxsize+5), maxsize+5)))
586 self.assertEqual(take(3, count(2, 1.25)), [2, 3.25, 4.5])
Raymond Hettinger30729212009-02-12 06:28:27 +0000587 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettinger8a882a72009-02-12 12:08:18 +0000588 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
589 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerde09acf2009-02-12 12:53:53 +0000590 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
591 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300592 BIGINT = 1<<1000
593 self.assertEqual(take(3, count(step=BIGINT)), [0, BIGINT, 2*BIGINT])
Raymond Hettinger30729212009-02-12 06:28:27 +0000594 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
595 c = count(3, 5)
596 self.assertEqual(repr(c), 'count(3, 5)')
597 next(c)
598 self.assertEqual(repr(c), 'count(8, 5)')
599 c = count(-9, 0)
600 self.assertEqual(repr(c), 'count(-9, 0)')
601 next(c)
602 self.assertEqual(repr(c), 'count(-9, 0)')
603 c = count(-9, -3)
604 self.assertEqual(repr(c), 'count(-9, -3)')
605 next(c)
606 self.assertEqual(repr(c), 'count(-12, -3)')
607 self.assertEqual(repr(c), 'count(-12, -3)')
608 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
609 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
610 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 +0300611 self.assertEqual(repr(count(10, 1.00)), 'count(10, 1.0)')
612 c = count(10, 1.0)
613 self.assertEqual(type(next(c)), int)
614 self.assertEqual(type(next(c)), float)
Raymond Hettinger30729212009-02-12 06:28:27 +0000615 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
616 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 +0300617 # Test repr
618 r1 = repr(count(i, j))
Raymond Hettinger30729212009-02-12 06:28:27 +0000619 if j == 1:
Serhiy Storchaka95949422013-08-27 19:40:23 +0300620 r2 = ('count(%r)' % i)
Raymond Hettinger30729212009-02-12 06:28:27 +0000621 else:
Serhiy Storchaka95949422013-08-27 19:40:23 +0300622 r2 = ('count(%r, %r)' % (i, j))
Raymond Hettinger30729212009-02-12 06:28:27 +0000623 self.assertEqual(r1, r2)
Serhiy Storchakabad12572014-12-15 14:03:42 +0200624 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
625 self.pickletest(proto, count(i, j))
Raymond Hettinger30729212009-02-12 06:28:27 +0000626
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000627 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000628 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000629 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000630 self.assertRaises(TypeError, cycle)
631 self.assertRaises(TypeError, cycle, 5)
632 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000633
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000634 # check copy, deepcopy, pickle
635 c = cycle('abc')
636 self.assertEqual(next(c), 'a')
637 #simple copy currently not supported, because __reduce__ returns
638 #an internal iterator
639 #self.assertEqual(take(10, copy.copy(c)), list('bcabcabcab'))
640 self.assertEqual(take(10, copy.deepcopy(c)), list('bcabcabcab'))
Serhiy Storchakabad12572014-12-15 14:03:42 +0200641 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
642 self.assertEqual(take(10, pickle.loads(pickle.dumps(c, proto))),
643 list('bcabcabcab'))
644 next(c)
645 self.assertEqual(take(10, pickle.loads(pickle.dumps(c, proto))),
646 list('cabcabcabc'))
647 next(c)
648 next(c)
649 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
650 self.pickletest(proto, cycle('abc'))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000651
Raymond Hettingera166ce52015-08-15 14:45:49 -0700652 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
653 # test with partial consumed input iterable
654 it = iter('abcde')
655 c = cycle(it)
Raymond Hettinger1cadf762015-08-15 14:47:27 -0700656 _ = [next(c) for i in range(2)] # consume 2 of 5 inputs
Raymond Hettingera166ce52015-08-15 14:45:49 -0700657 p = pickle.dumps(c, proto)
658 d = pickle.loads(p) # rebuild the cycle object
659 self.assertEqual(take(20, d), list('cdeabcdeabcdeabcdeab'))
660
661 # test with completely consumed input iterable
662 it = iter('abcde')
663 c = cycle(it)
Raymond Hettinger1cadf762015-08-15 14:47:27 -0700664 _ = [next(c) for i in range(7)] # consume 7 of 5 inputs
Raymond Hettingera166ce52015-08-15 14:45:49 -0700665 p = pickle.dumps(c, proto)
666 d = pickle.loads(p) # rebuild the cycle object
667 self.assertEqual(take(20, d), list('cdeabcdeabcdeabcdeab'))
668
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700669 def test_cycle_setstate(self):
670 # Verify both modes for restoring state
671
672 # Mode 0 is efficient. It uses an incompletely consumed input
673 # iterator to build a cycle object and then passes in state with
674 # a list of previously consumed values. There is no data
Martin Panterf1579822016-05-26 06:03:33 +0000675 # overlap between the two.
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700676 c = cycle('defg')
677 c.__setstate__((list('abc'), 0))
678 self.assertEqual(take(20, c), list('defgabcdefgabcdefgab'))
679
680 # Mode 1 is inefficient. It starts with a cycle object built
681 # from an iterator over the remaining elements in a partial
682 # cycle and then passes in state with all of the previously
683 # seen values (this overlaps values included in the iterator).
684 c = cycle('defg')
685 c.__setstate__((list('abcdefg'), 1))
686 self.assertEqual(take(20, c), list('defgabcdefgabcdefgab'))
687
688 # The first argument to setstate needs to be a tuple
Serhiy Storchaka8f0f2052016-10-02 09:13:14 +0300689 with self.assertRaises(TypeError):
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700690 cycle('defg').__setstate__([list('abcdefg'), 0])
691
692 # The first argument in the setstate tuple must be a list
693 with self.assertRaises(TypeError):
694 c = cycle('defg')
Serhiy Storchaka8f0f2052016-10-02 09:13:14 +0300695 c.__setstate__((tuple('defg'), 0))
696 take(20, c)
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700697
Serhiy Storchaka8f0f2052016-10-02 09:13:14 +0300698 # The second argument in the setstate tuple must be an int
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700699 with self.assertRaises(TypeError):
700 cycle('defg').__setstate__((list('abcdefg'), 'x'))
701
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300702 self.assertRaises(TypeError, cycle('').__setstate__, ())
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300703 self.assertRaises(TypeError, cycle('').__setstate__, ([],))
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300704
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000705 def test_groupby(self):
706 # Check whether it accepts arguments correctly
707 self.assertEqual([], list(groupby([])))
708 self.assertEqual([], list(groupby([], key=id)))
709 self.assertRaises(TypeError, list, groupby('abc', []))
710 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000711 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000712
713 # Check normal input
714 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
715 (2,15,22), (3,16,23), (3,17,23)]
716 dup = []
717 for k, g in groupby(s, lambda r:r[0]):
718 for elem in g:
719 self.assertEqual(k, elem[0])
720 dup.append(elem)
721 self.assertEqual(s, dup)
722
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000723 # Check normal pickled
Serhiy Storchakabad12572014-12-15 14:03:42 +0200724 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
725 dup = []
726 for k, g in pickle.loads(pickle.dumps(groupby(s, testR), proto)):
727 for elem in g:
728 self.assertEqual(k, elem[0])
729 dup.append(elem)
730 self.assertEqual(s, dup)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000731
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000732 # Check nested case
733 dup = []
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000734 for k, g in groupby(s, testR):
735 for ik, ig in groupby(g, testR2):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000736 for elem in ig:
737 self.assertEqual(k, elem[0])
738 self.assertEqual(ik, elem[2])
739 dup.append(elem)
740 self.assertEqual(s, dup)
741
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000742 # Check nested and pickled
Serhiy Storchakabad12572014-12-15 14:03:42 +0200743 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
744 dup = []
745 for k, g in pickle.loads(pickle.dumps(groupby(s, testR), proto)):
746 for ik, ig in pickle.loads(pickle.dumps(groupby(g, testR2), proto)):
747 for elem in ig:
748 self.assertEqual(k, elem[0])
749 self.assertEqual(ik, elem[2])
750 dup.append(elem)
751 self.assertEqual(s, dup)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000752
753
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000754 # Check case where inner iterator is not used
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000755 keys = [k for k, g in groupby(s, testR)]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000756 expectedkeys = set([r[0] for r in s])
757 self.assertEqual(set(keys), expectedkeys)
758 self.assertEqual(len(keys), len(expectedkeys))
759
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300760 # Check case where inner iterator is used after advancing the groupby
761 # iterator
762 s = list(zip('AABBBAAAA', range(9)))
763 it = groupby(s, testR)
764 _, g1 = next(it)
765 _, g2 = next(it)
766 _, g3 = next(it)
767 self.assertEqual(list(g1), [])
768 self.assertEqual(list(g2), [])
769 self.assertEqual(next(g3), ('A', 5))
770 list(it) # exhaust the groupby iterator
771 self.assertEqual(list(g3), [])
772
773 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
774 it = groupby(s, testR)
775 _, g = next(it)
776 next(it)
777 next(it)
778 self.assertEqual(list(pickle.loads(pickle.dumps(g, proto))), [])
779
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000780 # Exercise pipes and filters style
781 s = 'abracadabra'
782 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000783 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000784 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
785 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000786 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000787 self.assertEqual(r, ['a', 'b', 'r'])
788 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000789 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000790 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
791 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000792 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000793 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
794
Georg Brandla18af4e2007-04-21 15:47:16 +0000795 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000796 class ExpectedError(Exception):
797 pass
798 def delayed_raise(n=0):
799 for i in range(n):
800 yield 'yo'
801 raise ExpectedError
802 def gulp(iterable, keyp=None, func=list):
803 return [func(g) for k, g in groupby(iterable, keyp)]
804
Georg Brandla18af4e2007-04-21 15:47:16 +0000805 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000806 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000807 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000808 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
809
Serhiy Storchakaa60c2fe2015-03-12 21:56:08 +0200810 # __eq__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000811 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000812 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000813 raise ExpectedError
814 s = [DummyCmp(), DummyCmp(), None]
815
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000816 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000817 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000818 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000819 self.assertRaises(ExpectedError, gulp, s)
820
821 # keyfunc failure
822 def keyfunc(obj):
823 if keyfunc.skip > 0:
824 keyfunc.skip -= 1
825 return obj
826 else:
827 raise ExpectedError
828
829 # keyfunc failure on outer object
830 keyfunc.skip = 0
831 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
832 keyfunc.skip = 1
833 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
834
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000835 def test_filter(self):
836 self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
837 self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])
838 self.assertEqual(list(filter(bool, [0,1,0,2,0])), [1,2])
839 self.assertEqual(take(4, filter(isEven, count())), [0,2,4,6])
840 self.assertRaises(TypeError, filter)
841 self.assertRaises(TypeError, filter, lambda x:x)
842 self.assertRaises(TypeError, filter, lambda x:x, range(6), 7)
843 self.assertRaises(TypeError, filter, isEven, 3)
844 self.assertRaises(TypeError, next, filter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000845
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000846 # check copy, deepcopy, pickle
847 ans = [0,2,4]
848
849 c = filter(isEven, range(6))
850 self.assertEqual(list(copy.copy(c)), ans)
851 c = filter(isEven, range(6))
852 self.assertEqual(list(copy.deepcopy(c)), ans)
Serhiy Storchakabad12572014-12-15 14:03:42 +0200853 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
854 c = filter(isEven, range(6))
855 self.assertEqual(list(pickle.loads(pickle.dumps(c, proto))), ans)
856 next(c)
857 self.assertEqual(list(pickle.loads(pickle.dumps(c, proto))), ans[1:])
858 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
859 c = filter(isEven, range(6))
860 self.pickletest(proto, c)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000861
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000862 def test_filterfalse(self):
863 self.assertEqual(list(filterfalse(isEven, range(6))), [1,3,5])
864 self.assertEqual(list(filterfalse(None, [0,1,0,2,0])), [0,0,0])
865 self.assertEqual(list(filterfalse(bool, [0,1,0,2,0])), [0,0,0])
866 self.assertEqual(take(4, filterfalse(isEven, count())), [1,3,5,7])
867 self.assertRaises(TypeError, filterfalse)
868 self.assertRaises(TypeError, filterfalse, lambda x:x)
869 self.assertRaises(TypeError, filterfalse, lambda x:x, range(6), 7)
870 self.assertRaises(TypeError, filterfalse, isEven, 3)
871 self.assertRaises(TypeError, next, filterfalse(range(6), range(6)))
Serhiy Storchakabad12572014-12-15 14:03:42 +0200872 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
873 self.pickletest(proto, filterfalse(isEven, range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000874
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000875 def test_zip(self):
876 # XXX This is rather silly now that builtin zip() calls zip()...
877 ans = [(x,y) for x, y in zip('abc',count())]
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000878 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000879 self.assertEqual(list(zip('abc', range(6))), lzip('abc', range(6)))
880 self.assertEqual(list(zip('abcdef', range(3))), lzip('abcdef', range(3)))
881 self.assertEqual(take(3,zip('abcdef', count())), lzip('abcdef', range(3)))
882 self.assertEqual(list(zip('abcdef')), lzip('abcdef'))
883 self.assertEqual(list(zip()), lzip())
884 self.assertRaises(TypeError, zip, 3)
885 self.assertRaises(TypeError, zip, range(3), 3)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000886 self.assertEqual([tuple(list(pair)) for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000887 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000888 self.assertEqual([pair for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000889 lzip('abc', 'def'))
Alex Gaynore151d212011-07-17 16:21:30 -0700890
891 @support.impl_detail("tuple reuse is specific to CPython")
892 def test_zip_tuple_reuse(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000893 ids = list(map(id, zip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000894 self.assertEqual(min(ids), max(ids))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000895 ids = list(map(id, list(zip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000896 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000897
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000898 # check copy, deepcopy, pickle
899 ans = [(x,y) for x, y in copy.copy(zip('abc',count()))]
900 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
901
902 ans = [(x,y) for x, y in copy.deepcopy(zip('abc',count()))]
903 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
904
Serhiy Storchakabad12572014-12-15 14:03:42 +0200905 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
906 ans = [(x,y) for x, y in pickle.loads(pickle.dumps(zip('abc',count()), proto))]
907 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000908
Serhiy Storchakabad12572014-12-15 14:03:42 +0200909 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
910 testIntermediate = zip('abc',count())
911 next(testIntermediate)
912 ans = [(x,y) for x, y in pickle.loads(pickle.dumps(testIntermediate, proto))]
913 self.assertEqual(ans, [('b', 1), ('c', 2)])
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000914
Serhiy Storchakabad12572014-12-15 14:03:42 +0200915 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
916 self.pickletest(proto, zip('abc', count()))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000917
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000918 def test_ziplongest(self):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000919 for args in [
920 ['abc', range(6)],
921 [range(6), 'abc'],
922 [range(1000), range(2000,2100), range(3000,3050)],
923 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
924 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
925 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000926 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
927 for i in range(max(map(len, args)))]
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000928 self.assertEqual(list(zip_longest(*args)), target)
929 self.assertEqual(list(zip_longest(*args, **{})), target)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000930 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 +0000931 self.assertEqual(list(zip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000932
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000933 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 +0000934
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000935 self.assertEqual(list(zip_longest()), list(zip()))
936 self.assertEqual(list(zip_longest([])), list(zip([])))
937 self.assertEqual(list(zip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000938
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000939 self.assertEqual(list(zip_longest('abc', 'defg', **{})),
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000940 list(zip(list('abc')+[None], 'defg'))) # empty keyword dict
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000941 self.assertRaises(TypeError, zip_longest, 3)
942 self.assertRaises(TypeError, zip_longest, range(3), 3)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000943
944 for stmt in [
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000945 "zip_longest('abc', fv=1)",
946 "zip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000947 ]:
948 try:
949 eval(stmt, globals(), locals())
950 except TypeError:
951 pass
952 else:
953 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000954
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000955 self.assertEqual([tuple(list(pair)) for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000956 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000957 self.assertEqual([pair for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000958 list(zip('abc', 'def')))
Alex Gaynore151d212011-07-17 16:21:30 -0700959
960 @support.impl_detail("tuple reuse is specific to CPython")
961 def test_zip_longest_tuple_reuse(self):
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000962 ids = list(map(id, zip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000963 self.assertEqual(min(ids), max(ids))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000964 ids = list(map(id, list(zip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000965 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
966
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000967 def test_zip_longest_pickling(self):
Serhiy Storchakabad12572014-12-15 14:03:42 +0200968 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
969 self.pickletest(proto, zip_longest("abc", "def"))
970 self.pickletest(proto, zip_longest("abc", "defgh"))
971 self.pickletest(proto, zip_longest("abc", "defgh", fillvalue=1))
972 self.pickletest(proto, zip_longest("", "defgh"))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000973
Sergey Fedoseev6a650aa2019-08-30 09:25:48 +0500974 def test_zip_longest_bad_iterable(self):
975 exception = TypeError()
976
977 class BadIterable:
978 def __iter__(self):
979 raise exception
980
981 with self.assertRaises(TypeError) as cm:
982 zip_longest(BadIterable())
983
984 self.assertIs(cm.exception, exception)
985
Raymond Hettingerfc438512009-11-01 20:55:33 +0000986 def test_bug_7244(self):
987
988 class Repeater:
989 # this class is similar to itertools.repeat
990 def __init__(self, o, t, e):
991 self.o = o
992 self.t = int(t)
993 self.e = e
994 def __iter__(self): # its iterator is itself
995 return self
996 def __next__(self):
997 if self.t > 0:
998 self.t -= 1
999 return self.o
1000 else:
1001 raise self.e
1002
1003 # Formerly this code in would fail in debug mode
1004 # with Undetected Error and Stop Iteration
1005 r1 = Repeater(1, 3, StopIteration)
1006 r2 = Repeater(2, 4, StopIteration)
1007 def run(r1, r2):
1008 result = []
1009 for i, j in zip_longest(r1, r2, fillvalue=0):
1010 with support.captured_output('stdout'):
1011 print((i, j))
1012 result.append((i, j))
1013 return result
1014 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
1015
1016 # Formerly, the RuntimeError would be lost
1017 # and StopIteration would stop as expected
1018 r1 = Repeater(1, 3, RuntimeError)
1019 r2 = Repeater(2, 4, StopIteration)
1020 it = zip_longest(r1, r2, fillvalue=0)
1021 self.assertEqual(next(it), (1, 2))
1022 self.assertEqual(next(it), (1, 2))
1023 self.assertEqual(next(it), (1, 2))
1024 self.assertRaises(RuntimeError, next, it)
1025
Christian Heimesc3f30c42008-02-22 16:37:40 +00001026 def test_product(self):
1027 for args, result in [
Christian Heimes78644762008-03-04 23:39:23 +00001028 ([], [()]), # zero iterables
Christian Heimesc3f30c42008-02-22 16:37:40 +00001029 (['ab'], [('a',), ('b',)]), # one iterable
1030 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
1031 ([range(0), range(2), range(3)], []), # first iterable with zero length
1032 ([range(2), range(0), range(3)], []), # middle iterable with zero length
1033 ([range(2), range(3), range(0)], []), # last iterable with zero length
1034 ]:
1035 self.assertEqual(list(product(*args)), result)
Christian Heimesf16baeb2008-02-29 14:57:44 +00001036 for r in range(4):
1037 self.assertEqual(list(product(*(args*r))),
1038 list(product(*args, **dict(repeat=r))))
Christian Heimesc3f30c42008-02-22 16:37:40 +00001039 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
1040 self.assertRaises(TypeError, product, range(6), None)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001041
1042 def product1(*args, **kwds):
1043 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
1044 n = len(pools)
1045 if n == 0:
1046 yield ()
1047 return
1048 if any(len(pool) == 0 for pool in pools):
1049 return
1050 indices = [0] * n
1051 yield tuple(pool[i] for pool, i in zip(pools, indices))
1052 while 1:
1053 for i in reversed(range(n)): # right to left
1054 if indices[i] == len(pools[i]) - 1:
1055 continue
1056 indices[i] += 1
1057 for j in range(i+1, n):
1058 indices[j] = 0
1059 yield tuple(pool[i] for pool, i in zip(pools, indices))
1060 break
1061 else:
1062 return
1063
1064 def product2(*args, **kwds):
1065 'Pure python version used in docs'
1066 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
1067 result = [[]]
1068 for pool in pools:
1069 result = [x+[y] for x in result for y in pool]
1070 for prod in result:
1071 yield tuple(prod)
1072
Christian Heimesc3f30c42008-02-22 16:37:40 +00001073 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
1074 set('abcdefg'), range(11), tuple(range(13))]
1075 for i in range(100):
1076 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Christian Heimes78644762008-03-04 23:39:23 +00001077 expected_len = prod(map(len, args))
1078 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001079 self.assertEqual(list(product(*args)), list(product1(*args)))
1080 self.assertEqual(list(product(*args)), list(product2(*args)))
Christian Heimesc3f30c42008-02-22 16:37:40 +00001081 args = map(iter, args)
Christian Heimes78644762008-03-04 23:39:23 +00001082 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesc3f30c42008-02-22 16:37:40 +00001083
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05001084 @support.bigaddrspacetest
1085 def test_product_overflow(self):
Serhiy Storchakadee948b2015-02-03 01:34:09 +02001086 with self.assertRaises((OverflowError, MemoryError)):
1087 product(*(['ab']*2**5), repeat=2**25)
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05001088
Alex Gaynore151d212011-07-17 16:21:30 -07001089 @support.impl_detail("tuple reuse is specific to CPython")
1090 def test_product_tuple_reuse(self):
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001091 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
1092 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +00001093
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001094 def test_product_pickling(self):
1095 # check copy, deepcopy, pickle
1096 for args, result in [
1097 ([], [()]), # zero iterables
1098 (['ab'], [('a',), ('b',)]), # one iterable
1099 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
1100 ([range(0), range(2), range(3)], []), # first iterable with zero length
1101 ([range(2), range(0), range(3)], []), # middle iterable with zero length
1102 ([range(2), range(3), range(0)], []), # last iterable with zero length
1103 ]:
1104 self.assertEqual(list(copy.copy(product(*args))), result)
1105 self.assertEqual(list(copy.deepcopy(product(*args))), result)
Serhiy Storchakabad12572014-12-15 14:03:42 +02001106 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1107 self.pickletest(proto, product(*args))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001108
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00001109 def test_product_issue_25021(self):
1110 # test that indices are properly clamped to the length of the tuples
1111 p = product((1, 2),(3,))
1112 p.__setstate__((0, 0x1000)) # will access tuple element 1 if not clamped
1113 self.assertEqual(next(p), (2, 3))
1114 # test that empty tuple in the list will result in an immediate StopIteration
1115 p = product((1, 2), (), (3,))
1116 p.__setstate__((0, 0, 0x1000)) # will access tuple element 1 if not clamped
1117 self.assertRaises(StopIteration, next, p)
1118
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001119 def test_repeat(self):
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00001120 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
Guido van Rossum805365e2007-05-07 22:24:25 +00001121 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001122 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001123 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +00001124 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001125 self.assertEqual(list(repeat('a', 0)), [])
1126 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001127 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001128 self.assertRaises(TypeError, repeat, None, 3, 4)
1129 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001130 r = repeat(1+0j)
1131 self.assertEqual(repr(r), 'repeat((1+0j))')
1132 r = repeat(1+0j, 5)
1133 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
1134 list(r)
1135 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001136
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001137 # check copy, deepcopy, pickle
1138 c = repeat(object='a', times=10)
1139 self.assertEqual(next(c), 'a')
1140 self.assertEqual(take(2, copy.copy(c)), list('a' * 2))
1141 self.assertEqual(take(2, copy.deepcopy(c)), list('a' * 2))
Serhiy Storchakabad12572014-12-15 14:03:42 +02001142 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1143 self.pickletest(proto, repeat(object='a', times=10))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001144
Raymond Hettinger97d35552014-06-24 21:36:58 -07001145 def test_repeat_with_negative_times(self):
1146 self.assertEqual(repr(repeat('a', -1)), "repeat('a', 0)")
1147 self.assertEqual(repr(repeat('a', -2)), "repeat('a', 0)")
1148 self.assertEqual(repr(repeat('a', times=-1)), "repeat('a', 0)")
1149 self.assertEqual(repr(repeat('a', times=-2)), "repeat('a', 0)")
1150
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001151 def test_map(self):
1152 self.assertEqual(list(map(operator.pow, range(3), range(1,7))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001153 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001154 self.assertEqual(list(map(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001155 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001156 self.assertEqual(list(map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +00001157 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001158 self.assertEqual(take(2,map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +00001159 [('a',0),('b',1)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001160 self.assertEqual(list(map(operator.pow, [])), [])
1161 self.assertRaises(TypeError, map)
1162 self.assertRaises(TypeError, list, map(None, range(3), range(3)))
1163 self.assertRaises(TypeError, map, operator.neg)
1164 self.assertRaises(TypeError, next, map(10, range(5)))
1165 self.assertRaises(ValueError, next, map(errfunc, [4], [5]))
1166 self.assertRaises(TypeError, next, map(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001167
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001168 # check copy, deepcopy, pickle
1169 ans = [('a',0),('b',1),('c',2)]
1170
1171 c = map(tupleize, 'abc', count())
1172 self.assertEqual(list(copy.copy(c)), ans)
1173
1174 c = map(tupleize, 'abc', count())
1175 self.assertEqual(list(copy.deepcopy(c)), ans)
1176
Serhiy Storchakabad12572014-12-15 14:03:42 +02001177 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1178 c = map(tupleize, 'abc', count())
1179 self.pickletest(proto, c)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001180
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001181 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001182 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
1183 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001184 self.assertEqual(take(3, starmap(operator.pow, zip(count(), count(1)))),
Raymond Hettinger02420702003-06-29 20:36:23 +00001185 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001186 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +00001187 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
1188 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001189 self.assertRaises(TypeError, starmap)
1190 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001191 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
1192 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
1193 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001194
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001195 # check copy, deepcopy, pickle
1196 ans = [0**1, 1**2, 2**3]
1197
1198 c = starmap(operator.pow, zip(range(3), range(1,7)))
1199 self.assertEqual(list(copy.copy(c)), ans)
1200
1201 c = starmap(operator.pow, zip(range(3), range(1,7)))
1202 self.assertEqual(list(copy.deepcopy(c)), ans)
1203
Serhiy Storchakabad12572014-12-15 14:03:42 +02001204 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1205 c = starmap(operator.pow, zip(range(3), range(1,7)))
1206 self.pickletest(proto, c)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001207
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001208 def test_islice(self):
1209 for args in [ # islice(args) should agree with range(args)
1210 (10, 20, 3),
1211 (10, 3, 20),
1212 (10, 20),
Cheryl Sabellada1734c2018-03-26 21:29:33 -04001213 (10, 10),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001214 (10, 3),
1215 (20,)
1216 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +00001217 self.assertEqual(list(islice(range(100), *args)),
1218 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001219
1220 for args, tgtargs in [ # Stop when seqn is exhausted
1221 ((10, 110, 3), ((10, 100, 3))),
1222 ((10, 110), ((10, 100))),
1223 ((110,), (100,))
1224 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +00001225 self.assertEqual(list(islice(range(100), *args)),
1226 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001227
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001228 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +00001229 self.assertEqual(list(islice(range(10), None)), list(range(10)))
1230 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
1231 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
1232 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
1233 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001234
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +00001235 # Test number of items consumed SF #1171417
1236 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +00001237 self.assertEqual(list(islice(it, 3)), list(range(3)))
1238 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +00001239
Cheryl Sabellada1734c2018-03-26 21:29:33 -04001240 it = iter(range(10))
1241 self.assertEqual(list(islice(it, 3, 3)), [])
1242 self.assertEqual(list(it), list(range(3, 10)))
1243
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001244 # Test invalid arguments
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001245 ra = range(10)
1246 self.assertRaises(TypeError, islice, ra)
1247 self.assertRaises(TypeError, islice, ra, 1, 2, 3, 4)
1248 self.assertRaises(ValueError, islice, ra, -5, 10, 1)
1249 self.assertRaises(ValueError, islice, ra, 1, -5, -1)
1250 self.assertRaises(ValueError, islice, ra, 1, 10, -1)
1251 self.assertRaises(ValueError, islice, ra, 1, 10, 0)
1252 self.assertRaises(ValueError, islice, ra, 'a')
1253 self.assertRaises(ValueError, islice, ra, 'a', 1)
1254 self.assertRaises(ValueError, islice, ra, 1, 'a')
1255 self.assertRaises(ValueError, islice, ra, 'a', 1, 1)
1256 self.assertRaises(ValueError, islice, ra, 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +00001257 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001258
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001259 # Issue #10323: Less islice in a predictable state
1260 c = count()
1261 self.assertEqual(list(islice(c, 1, 3, 50)), [1])
1262 self.assertEqual(next(c), 3)
1263
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001264 # check copy, deepcopy, pickle
1265 for args in [ # islice(args) should agree with range(args)
1266 (10, 20, 3),
1267 (10, 3, 20),
1268 (10, 20),
1269 (10, 3),
1270 (20,)
1271 ]:
1272 self.assertEqual(list(copy.copy(islice(range(100), *args))),
1273 list(range(*args)))
1274 self.assertEqual(list(copy.deepcopy(islice(range(100), *args))),
1275 list(range(*args)))
Serhiy Storchakabad12572014-12-15 14:03:42 +02001276 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1277 self.pickletest(proto, islice(range(100), *args))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001278
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001279 # Issue #21321: check source iterator is not referenced
1280 # from islice() after the latter has been exhausted
1281 it = (x for x in (1, 2))
1282 wr = weakref.ref(it)
1283 it = islice(it, 1)
1284 self.assertIsNotNone(wr())
1285 list(it) # exhaust the iterator
Benjamin Peterson8e163512014-08-24 18:07:28 -05001286 support.gc_collect()
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001287 self.assertIsNone(wr())
1288
Will Roberts0ecdc522017-06-08 08:03:04 +02001289 # Issue #30537: islice can accept integer-like objects as
1290 # arguments
1291 class IntLike(object):
1292 def __init__(self, val):
1293 self.val = val
1294 def __index__(self):
1295 return self.val
1296 self.assertEqual(list(islice(range(100), IntLike(10))), list(range(10)))
1297 self.assertEqual(list(islice(range(100), IntLike(10), IntLike(50))),
1298 list(range(10, 50)))
1299 self.assertEqual(list(islice(range(100), IntLike(10), IntLike(50), IntLike(5))),
1300 list(range(10,50,5)))
1301
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001302 def test_takewhile(self):
1303 data = [1, 3, 5, 20, 2, 4, 6, 8]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001304 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001305 self.assertEqual(list(takewhile(underten, [])), [])
1306 self.assertRaises(TypeError, takewhile)
1307 self.assertRaises(TypeError, takewhile, operator.pow)
1308 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001309 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
1310 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001311 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
1312 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +00001313 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001314
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001315 # check copy, deepcopy, pickle
1316 self.assertEqual(list(copy.copy(takewhile(underten, data))), [1, 3, 5])
1317 self.assertEqual(list(copy.deepcopy(takewhile(underten, data))),
1318 [1, 3, 5])
Serhiy Storchakabad12572014-12-15 14:03:42 +02001319 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1320 self.pickletest(proto, takewhile(underten, data))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001321
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001322 def test_dropwhile(self):
1323 data = [1, 3, 5, 20, 2, 4, 6, 8]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001324 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001325 self.assertEqual(list(dropwhile(underten, [])), [])
1326 self.assertRaises(TypeError, dropwhile)
1327 self.assertRaises(TypeError, dropwhile, operator.pow)
1328 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001329 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
1330 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001331
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001332 # check copy, deepcopy, pickle
1333 self.assertEqual(list(copy.copy(dropwhile(underten, data))), [20, 2, 4, 6, 8])
1334 self.assertEqual(list(copy.deepcopy(dropwhile(underten, data))),
1335 [20, 2, 4, 6, 8])
Serhiy Storchakabad12572014-12-15 14:03:42 +02001336 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1337 self.pickletest(proto, dropwhile(underten, data))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001338
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001339 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +00001340 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001341
1342 a, b = tee([]) # test empty iterator
1343 self.assertEqual(list(a), [])
1344 self.assertEqual(list(b), [])
1345
1346 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +00001347 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001348
1349 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +00001350 self.assertEqual(list(a), list(range(n)))
1351 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001352
1353 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +00001354 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001355 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001356 del a
Guido van Rossum805365e2007-05-07 22:24:25 +00001357 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001358
1359 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +00001360 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001361 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001362 del b
Guido van Rossum805365e2007-05-07 22:24:25 +00001363 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001364
Guido van Rossum805365e2007-05-07 22:24:25 +00001365 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001366 order = [0]*n + [1]*n
1367 random.shuffle(order)
1368 lists = ([], [])
1369 its = tee(irange(n))
1370 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +00001371 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001372 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +00001373 self.assertEqual(lists[0], list(range(n)))
1374 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001375
Raymond Hettingerad983e72003-11-12 14:32:26 +00001376 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001377 self.assertRaises(TypeError, tee)
1378 self.assertRaises(TypeError, tee, 3)
1379 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +00001380 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001381
Raymond Hettingerad983e72003-11-12 14:32:26 +00001382 # tee object should be instantiable
1383 a, b = tee('abc')
1384 c = type(a)('def')
1385 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +00001386
Raymond Hettingerad983e72003-11-12 14:32:26 +00001387 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +00001388 a, b, c = tee(range(2000), 3)
1389 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001390 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +00001391 self.assertEqual(list(b), list(range(2000)))
1392 self.assertEqual([next(c), next(c)], list(range(2)))
1393 self.assertEqual(list(a), list(range(100,2000)))
1394 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +00001395
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001396 # test values of n
1397 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +00001398 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +00001399 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001400 result = tee('abc', n)
1401 self.assertEqual(type(result), tuple)
1402 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001403 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001404
Raymond Hettingerad983e72003-11-12 14:32:26 +00001405 # tee pass-through to copyable iterator
1406 a, b = tee('abc')
1407 c, d = tee(a)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +00001408 self.assertTrue(a is c)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001409
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001410 # test tee_new
1411 t1, t2 = tee('abc')
1412 tnew = type(t1)
1413 self.assertRaises(TypeError, tnew)
1414 self.assertRaises(TypeError, tnew, 10)
1415 t3 = tnew(t1)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +00001416 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +00001417
Raymond Hettingera9f60922004-10-17 16:40:14 +00001418 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +00001419 a, b = tee(range(10))
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001420 p = weakref.proxy(a)
Raymond Hettingera9f60922004-10-17 16:40:14 +00001421 self.assertEqual(getattr(p, '__class__'), type(b))
1422 del a
1423 self.assertRaises(ReferenceError, getattr, p, '__class__')
1424
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001425 ans = list('abc')
1426 long_ans = list(range(10000))
1427
1428 # check copy
1429 a, b = tee('abc')
1430 self.assertEqual(list(copy.copy(a)), ans)
1431 self.assertEqual(list(copy.copy(b)), ans)
1432 a, b = tee(list(range(10000)))
1433 self.assertEqual(list(copy.copy(a)), long_ans)
1434 self.assertEqual(list(copy.copy(b)), long_ans)
1435
1436 # check partially consumed copy
1437 a, b = tee('abc')
1438 take(2, a)
1439 take(1, b)
1440 self.assertEqual(list(copy.copy(a)), ans[2:])
1441 self.assertEqual(list(copy.copy(b)), ans[1:])
1442 self.assertEqual(list(a), ans[2:])
1443 self.assertEqual(list(b), ans[1:])
1444 a, b = tee(range(10000))
1445 take(100, a)
1446 take(60, b)
1447 self.assertEqual(list(copy.copy(a)), long_ans[100:])
1448 self.assertEqual(list(copy.copy(b)), long_ans[60:])
1449 self.assertEqual(list(a), long_ans[100:])
1450 self.assertEqual(list(b), long_ans[60:])
1451
1452 # check deepcopy
1453 a, b = tee('abc')
1454 self.assertEqual(list(copy.deepcopy(a)), ans)
1455 self.assertEqual(list(copy.deepcopy(b)), ans)
1456 self.assertEqual(list(a), ans)
1457 self.assertEqual(list(b), ans)
1458 a, b = tee(range(10000))
1459 self.assertEqual(list(copy.deepcopy(a)), long_ans)
1460 self.assertEqual(list(copy.deepcopy(b)), long_ans)
1461 self.assertEqual(list(a), long_ans)
1462 self.assertEqual(list(b), long_ans)
1463
1464 # check partially consumed deepcopy
1465 a, b = tee('abc')
1466 take(2, a)
1467 take(1, b)
1468 self.assertEqual(list(copy.deepcopy(a)), ans[2:])
1469 self.assertEqual(list(copy.deepcopy(b)), ans[1:])
1470 self.assertEqual(list(a), ans[2:])
1471 self.assertEqual(list(b), ans[1:])
1472 a, b = tee(range(10000))
1473 take(100, a)
1474 take(60, b)
1475 self.assertEqual(list(copy.deepcopy(a)), long_ans[100:])
1476 self.assertEqual(list(copy.deepcopy(b)), long_ans[60:])
1477 self.assertEqual(list(a), long_ans[100:])
1478 self.assertEqual(list(b), long_ans[60:])
1479
1480 # check pickle
Serhiy Storchakabad12572014-12-15 14:03:42 +02001481 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1482 self.pickletest(proto, iter(tee('abc')))
1483 a, b = tee('abc')
1484 self.pickletest(proto, a, compare=ans)
1485 self.pickletest(proto, b, compare=ans)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001486
Serhiy Storchakaa3e91282013-01-25 13:19:31 +02001487 # Issue 13454: Crash when deleting backward iterator from tee()
1488 def test_tee_del_backward(self):
Serhiy Storchaka5bb893c2013-01-26 11:52:06 +02001489 forward, backward = tee(repeat(None, 20000000))
Serhiy Storchaka9db55002015-03-28 20:38:37 +02001490 try:
1491 any(forward) # exhaust the iterator
1492 del backward
1493 except:
1494 del forward, backward
1495 raise
Serhiy Storchakaa3e91282013-01-25 13:19:31 +02001496
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001497 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001498 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +00001499
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001500 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +00001501 self.assertRaises(StopIteration, next, f([]))
1502 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001503
Georg Brandla18af4e2007-04-21 15:47:16 +00001504 self.assertRaises(StopIteration, next, islice([], None))
1505 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001506
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001507 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +00001508 self.assertRaises(StopIteration, next, p)
1509 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001510 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +00001511 self.assertRaises(StopIteration, next, p)
1512 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001513
Georg Brandla18af4e2007-04-21 15:47:16 +00001514 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001515
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001516 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +00001517 self.assertRaises(StopIteration, next, f(lambda x:x, []))
1518 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001519
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001520class TestExamples(unittest.TestCase):
1521
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001522 def test_accumulate(self):
Raymond Hettinger482ba772010-12-01 22:48:00 +00001523 self.assertEqual(list(accumulate([1,2,3,4,5])), [1, 3, 6, 10, 15])
1524
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001525 def test_accumulate_reducible(self):
1526 # check copy, deepcopy, pickle
1527 data = [1, 2, 3, 4, 5]
1528 accumulated = [1, 3, 6, 10, 15]
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001529
Serhiy Storchakabad12572014-12-15 14:03:42 +02001530 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1531 it = accumulate(data)
1532 self.assertEqual(list(pickle.loads(pickle.dumps(it, proto))), accumulated[:])
1533 self.assertEqual(next(it), 1)
1534 self.assertEqual(list(pickle.loads(pickle.dumps(it, proto))), accumulated[1:])
1535 it = accumulate(data)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001536 self.assertEqual(next(it), 1)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001537 self.assertEqual(list(copy.deepcopy(it)), accumulated[1:])
1538 self.assertEqual(list(copy.copy(it)), accumulated[1:])
1539
Serhiy Storchakad5516252016-03-06 14:00:45 +02001540 def test_accumulate_reducible_none(self):
1541 # Issue #25718: total is None
1542 it = accumulate([None, None, None], operator.is_)
1543 self.assertEqual(next(it), None)
1544 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1545 it_copy = pickle.loads(pickle.dumps(it, proto))
1546 self.assertEqual(list(it_copy), [True, False])
1547 self.assertEqual(list(copy.deepcopy(it)), [True, False])
1548 self.assertEqual(list(copy.copy(it)), [True, False])
1549
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001550 def test_chain(self):
1551 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
1552
1553 def test_chain_from_iterable(self):
1554 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
1555
1556 def test_combinations(self):
1557 self.assertEqual(list(combinations('ABCD', 2)),
1558 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
1559 self.assertEqual(list(combinations(range(4), 3)),
1560 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
1561
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001562 def test_combinations_with_replacement(self):
1563 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
1564 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
1565
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001566 def test_compress(self):
1567 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
1568
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001569 def test_count(self):
1570 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
1571
1572 def test_cycle(self):
1573 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
1574
1575 def test_dropwhile(self):
1576 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
1577
1578 def test_groupby(self):
1579 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
1580 list('ABCDAB'))
1581 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
1582 [list('AAAA'), list('BBB'), list('CC'), list('D')])
1583
1584 def test_filter(self):
1585 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
1586
1587 def test_filterfalse(self):
1588 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
1589
1590 def test_map(self):
1591 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
1592
1593 def test_islice(self):
1594 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
1595 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
1596 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
1597 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1598
1599 def test_zip(self):
1600 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
1601
1602 def test_zip_longest(self):
1603 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
1604 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
1605
1606 def test_permutations(self):
1607 self.assertEqual(list(permutations('ABCD', 2)),
1608 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
1609 self.assertEqual(list(permutations(range(3))),
1610 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
1611
1612 def test_product(self):
1613 self.assertEqual(list(product('ABCD', 'xy')),
1614 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
1615 self.assertEqual(list(product(range(2), repeat=3)),
1616 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
1617 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
1618
1619 def test_repeat(self):
1620 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
1621
1622 def test_stapmap(self):
1623 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
1624 [32, 9, 1000])
1625
1626 def test_takewhile(self):
1627 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
1628
1629
Cheryl Sabellada1734c2018-03-26 21:29:33 -04001630class TestPurePythonRoughEquivalents(unittest.TestCase):
1631
1632 @staticmethod
1633 def islice(iterable, *args):
1634 s = slice(*args)
1635 start, stop, step = s.start or 0, s.stop or sys.maxsize, s.step or 1
1636 it = iter(range(start, stop, step))
1637 try:
1638 nexti = next(it)
1639 except StopIteration:
1640 # Consume *iterable* up to the *start* position.
1641 for i, element in zip(range(start), iterable):
1642 pass
1643 return
1644 try:
1645 for i, element in enumerate(iterable):
1646 if i == nexti:
1647 yield element
1648 nexti = next(it)
1649 except StopIteration:
1650 # Consume to *stop*.
1651 for i, element in zip(range(i + 1, stop), iterable):
1652 pass
1653
1654 def test_islice_recipe(self):
1655 self.assertEqual(list(self.islice('ABCDEFG', 2)), list('AB'))
1656 self.assertEqual(list(self.islice('ABCDEFG', 2, 4)), list('CD'))
1657 self.assertEqual(list(self.islice('ABCDEFG', 2, None)), list('CDEFG'))
1658 self.assertEqual(list(self.islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1659 # Test items consumed.
1660 it = iter(range(10))
1661 self.assertEqual(list(self.islice(it, 3)), list(range(3)))
1662 self.assertEqual(list(it), list(range(3, 10)))
1663 it = iter(range(10))
1664 self.assertEqual(list(self.islice(it, 3, 3)), [])
1665 self.assertEqual(list(it), list(range(3, 10)))
1666 # Test that slice finishes in predictable state.
1667 c = count()
1668 self.assertEqual(list(self.islice(c, 1, 3, 50)), [1])
1669 self.assertEqual(next(c), 3)
1670
1671
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001672class TestGC(unittest.TestCase):
1673
1674 def makecycle(self, iterator, container):
1675 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +00001676 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001677 del container, iterator
1678
Raymond Hettinger482ba772010-12-01 22:48:00 +00001679 def test_accumulate(self):
1680 a = []
1681 self.makecycle(accumulate([1,2,a,3]), a)
1682
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001683 def test_chain(self):
1684 a = []
1685 self.makecycle(chain(a), a)
1686
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001687 def test_chain_from_iterable(self):
1688 a = []
1689 self.makecycle(chain.from_iterable([a]), a)
1690
1691 def test_combinations(self):
1692 a = []
1693 self.makecycle(combinations([1,2,a,3], 3), a)
1694
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001695 def test_combinations_with_replacement(self):
1696 a = []
1697 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1698
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001699 def test_compress(self):
1700 a = []
1701 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1702
Raymond Hettingerd6280f42009-02-16 20:50:56 +00001703 def test_count(self):
1704 a = []
1705 Int = type('Int', (int,), dict(x=a))
1706 self.makecycle(count(Int(0), Int(1)), a)
1707
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001708 def test_cycle(self):
1709 a = []
1710 self.makecycle(cycle([a]*2), a)
1711
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001712 def test_dropwhile(self):
1713 a = []
1714 self.makecycle(dropwhile(bool, [0, a, a]), a)
1715
1716 def test_groupby(self):
1717 a = []
1718 self.makecycle(groupby([a]*2, lambda x:x), a)
1719
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001720 def test_issue2246(self):
1721 # Issue 2246 -- the _grouper iterator was not included in GC
1722 n = 10
1723 keyfunc = lambda x: x
1724 for i, j in groupby(range(n), key=keyfunc):
1725 keyfunc.__dict__.setdefault('x',[]).append(j)
1726
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001727 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001728 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001729 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001730
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001731 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001732 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001733 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001734
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001735 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001736 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001737 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001738
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001739 def test_zip_longest(self):
1740 a = []
1741 self.makecycle(zip_longest([a]*2, [a]*3), a)
1742 b = [a, None]
1743 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1744
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001745 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001746 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001747 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001748
1749 def test_islice(self):
1750 a = []
1751 self.makecycle(islice([a]*2, None), a)
1752
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001753 def test_permutations(self):
1754 a = []
1755 self.makecycle(permutations([1,2,a,3], 3), a)
1756
1757 def test_product(self):
1758 a = []
1759 self.makecycle(product([1,2,a,3], repeat=3), a)
1760
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001761 def test_repeat(self):
1762 a = []
1763 self.makecycle(repeat(a), a)
1764
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001765 def test_starmap(self):
1766 a = []
1767 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1768
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001769 def test_takewhile(self):
1770 a = []
1771 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1772
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001773def R(seqn):
1774 'Regular generator'
1775 for i in seqn:
1776 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001777
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001778class G:
1779 'Sequence using __getitem__'
1780 def __init__(self, seqn):
1781 self.seqn = seqn
1782 def __getitem__(self, i):
1783 return self.seqn[i]
1784
1785class I:
1786 'Sequence using iterator protocol'
1787 def __init__(self, seqn):
1788 self.seqn = seqn
1789 self.i = 0
1790 def __iter__(self):
1791 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001792 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001793 if self.i >= len(self.seqn): raise StopIteration
1794 v = self.seqn[self.i]
1795 self.i += 1
1796 return v
1797
1798class Ig:
1799 'Sequence using iterator protocol defined with a generator'
1800 def __init__(self, seqn):
1801 self.seqn = seqn
1802 self.i = 0
1803 def __iter__(self):
1804 for val in self.seqn:
1805 yield val
1806
1807class X:
1808 'Missing __getitem__ and __iter__'
1809 def __init__(self, seqn):
1810 self.seqn = seqn
1811 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001812 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001813 if self.i >= len(self.seqn): raise StopIteration
1814 v = self.seqn[self.i]
1815 self.i += 1
1816 return v
1817
1818class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001819 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001820 def __init__(self, seqn):
1821 self.seqn = seqn
1822 self.i = 0
1823 def __iter__(self):
1824 return self
1825
1826class E:
1827 'Test propagation of exceptions'
1828 def __init__(self, seqn):
1829 self.seqn = seqn
1830 self.i = 0
1831 def __iter__(self):
1832 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001833 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001834 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001835
1836class S:
1837 'Test immediate stop'
1838 def __init__(self, seqn):
1839 pass
1840 def __iter__(self):
1841 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001842 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001843 raise StopIteration
1844
1845def L(seqn):
1846 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001847 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001848
1849
1850class TestVariousIteratorArgs(unittest.TestCase):
1851
Raymond Hettinger482ba772010-12-01 22:48:00 +00001852 def test_accumulate(self):
1853 s = [1,2,3,4,5]
1854 r = [1,3,6,10,15]
1855 n = len(s)
1856 for g in (G, I, Ig, L, R):
1857 self.assertEqual(list(accumulate(g(s))), r)
1858 self.assertEqual(list(accumulate(S(s))), [])
1859 self.assertRaises(TypeError, accumulate, X(s))
1860 self.assertRaises(TypeError, accumulate, N(s))
1861 self.assertRaises(ZeroDivisionError, list, accumulate(E(s)))
1862
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001863 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001864 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001865 for g in (G, I, Ig, S, L, R):
1866 self.assertEqual(list(chain(g(s))), list(g(s)))
1867 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001868 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001869 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001870 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1871
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001872 def test_compress(self):
1873 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1874 n = len(s)
1875 for g in (G, I, Ig, S, L, R):
1876 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1877 self.assertRaises(TypeError, compress, X(s), repeat(1))
1878 self.assertRaises(TypeError, compress, N(s), repeat(1))
1879 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1880
Christian Heimesc3f30c42008-02-22 16:37:40 +00001881 def test_product(self):
1882 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1883 self.assertRaises(TypeError, product, X(s))
1884 self.assertRaises(TypeError, product, N(s))
1885 self.assertRaises(ZeroDivisionError, product, E(s))
1886
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001887 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001888 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001889 for g in (G, I, Ig, S, L, R):
1890 tgtlen = len(s) * 3
1891 expected = list(g(s))*3
1892 actual = list(islice(cycle(g(s)), tgtlen))
1893 self.assertEqual(actual, expected)
1894 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001895 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001896 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1897
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001898 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001899 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001900 for g in (G, I, Ig, S, L, R):
1901 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1902 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001903 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001904 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1905
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001906 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001907 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001908 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001909 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001910 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001911 self.assertRaises(TypeError, filter, isEven, X(s))
1912 self.assertRaises(TypeError, filter, isEven, N(s))
1913 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001914
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001915 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001916 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001917 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001918 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001919 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001920 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1921 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1922 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001923
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001924 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001925 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001926 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001927 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1928 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1929 self.assertRaises(TypeError, zip, X(s))
1930 self.assertRaises(TypeError, zip, N(s))
1931 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001932
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001933 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001934 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001935 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001936 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1937 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1938 self.assertRaises(TypeError, zip_longest, X(s))
1939 self.assertRaises(TypeError, zip_longest, N(s))
1940 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001941
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001942 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001943 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001944 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001945 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001946 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001947 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001948 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001949 self.assertRaises(TypeError, map, onearg, X(s))
1950 self.assertRaises(TypeError, map, onearg, N(s))
1951 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001952
1953 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001954 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001955 for g in (G, I, Ig, S, L, R):
1956 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1957 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001958 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001959 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1960
1961 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001962 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001963 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001964 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001965 self.assertEqual(list(starmap(operator.pow, g(ss))),
1966 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001967 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001968 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001969 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1970
1971 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001972 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001973 for g in (G, I, Ig, S, L, R):
1974 tgt = []
1975 for elem in g(s):
1976 if not isEven(elem): break
1977 tgt.append(elem)
1978 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1979 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001980 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001981 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1982
1983 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001984 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001985 for g in (G, I, Ig, S, L, R):
1986 tgt = []
1987 for elem in g(s):
1988 if not tgt and isOdd(elem): continue
1989 tgt.append(elem)
1990 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1991 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001992 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001993 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1994
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001995 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001996 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001997 for g in (G, I, Ig, S, L, R):
1998 it1, it2 = tee(g(s))
1999 self.assertEqual(list(it1), list(g(s)))
2000 self.assertEqual(list(it2), list(g(s)))
2001 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00002002 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002003 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
2004
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002005class LengthTransparency(unittest.TestCase):
2006
2007 def test_repeat(self):
Armin Ronacheraa9a79d2012-10-06 14:03:24 +02002008 self.assertEqual(operator.length_hint(repeat(None, 50)), 50)
Raymond Hettinger97d35552014-06-24 21:36:58 -07002009 self.assertEqual(operator.length_hint(repeat(None, 0)), 0)
Armin Ronacheraa9a79d2012-10-06 14:03:24 +02002010 self.assertEqual(operator.length_hint(repeat(None), 12), 12)
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002011
Raymond Hettinger97d35552014-06-24 21:36:58 -07002012 def test_repeat_with_negative_times(self):
2013 self.assertEqual(operator.length_hint(repeat(None, -1)), 0)
2014 self.assertEqual(operator.length_hint(repeat(None, -2)), 0)
2015 self.assertEqual(operator.length_hint(repeat(None, times=-1)), 0)
2016 self.assertEqual(operator.length_hint(repeat(None, times=-2)), 0)
2017
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002018class RegressionTests(unittest.TestCase):
2019
2020 def test_sf_793826(self):
2021 # Fix Armin Rigo's successful efforts to wreak havoc
2022
2023 def mutatingtuple(tuple1, f, tuple2):
2024 # this builds a tuple t which is a copy of tuple1,
2025 # then calls f(t), then mutates t to be equal to tuple2
2026 # (needs len(tuple1) == len(tuple2)).
2027 def g(value, first=[1]):
2028 if first:
2029 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00002030 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002031 return value
2032 items = list(tuple2)
2033 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002034 gen = map(g, items)
2035 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00002036 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002037
2038 def f(t):
2039 global T
2040 T = t
2041 first[:] = list(T)
2042
2043 first = []
2044 mutatingtuple((1,2,3), f, (4,5,6))
2045 second = list(T)
2046 self.assertEqual(first, second)
2047
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002048
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00002049 def test_sf_950057(self):
2050 # Make sure that chain() and cycle() catch exceptions immediately
2051 # rather than when shifting between input sources
2052
2053 def gen1():
2054 hist.append(0)
2055 yield 1
2056 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00002057 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00002058 hist.append(2)
2059
2060 def gen2(x):
2061 hist.append(3)
2062 yield 2
2063 hist.append(4)
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00002064
2065 hist = []
2066 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
2067 self.assertEqual(hist, [0,1])
2068
2069 hist = []
2070 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
2071 self.assertEqual(hist, [0,1])
2072
2073 hist = []
2074 self.assertRaises(AssertionError, list, cycle(gen1()))
2075 self.assertEqual(hist, [0,1])
2076
Neil Schemenauer52a48e62019-07-30 11:08:18 -07002077 @support.skip_if_pgo_task
T. Wouters5466d4a2017-03-30 09:58:35 -07002078 def test_long_chain_of_empty_iterables(self):
2079 # Make sure itertools.chain doesn't run into recursion limits when
2080 # dealing with long chains of empty iterables. Even with a high
2081 # number this would probably only fail in Py_DEBUG mode.
2082 it = chain.from_iterable(() for unused in range(10000000))
2083 with self.assertRaises(StopIteration):
2084 next(it)
2085
Serhiy Storchakac740e4f2017-09-26 21:47:56 +03002086 def test_issue30347_1(self):
2087 def f(n):
2088 if n == 5:
2089 list(b)
2090 return n != 6
2091 for (k, b) in groupby(range(10), f):
2092 list(b) # shouldn't crash
2093
2094 def test_issue30347_2(self):
2095 class K:
2096 def __init__(self, v):
2097 pass
2098 def __eq__(self, other):
2099 nonlocal i
2100 i += 1
2101 if i == 1:
2102 next(g, None)
2103 return True
2104 i = 0
2105 g = next(groupby(range(10), K))[1]
2106 for j in range(2):
2107 next(g, None) # shouldn't crash
2108
2109
Thomas Woutersb2137042007-02-01 18:02:27 +00002110class SubclassWithKwargsTest(unittest.TestCase):
2111 def test_keywords_in_subclass(self):
2112 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002113 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002114 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00002115 class Subclass(cls):
2116 def __init__(self, newarg=None, *args):
2117 cls.__init__(self, *args)
2118 try:
2119 Subclass(newarg=1)
2120 except TypeError as err:
2121 # we expect type errors because of wrong argument count
Serhiy Storchaka5eb788b2017-06-06 18:45:22 +03002122 self.assertNotIn("keyword arguments", err.args[0])
Thomas Woutersb2137042007-02-01 18:02:27 +00002123
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002124@support.cpython_only
2125class SizeofTest(unittest.TestCase):
2126 def setUp(self):
2127 self.ssize_t = struct.calcsize('n')
2128
2129 check_sizeof = support.check_sizeof
2130
2131 def test_product_sizeof(self):
2132 basesize = support.calcobjsize('3Pi')
2133 check = self.check_sizeof
2134 check(product('ab', '12'), basesize + 2 * self.ssize_t)
2135 check(product(*(('abc',) * 10)), basesize + 10 * self.ssize_t)
2136
2137 def test_combinations_sizeof(self):
2138 basesize = support.calcobjsize('3Pni')
2139 check = self.check_sizeof
2140 check(combinations('abcd', 3), basesize + 3 * self.ssize_t)
2141 check(combinations(range(10), 4), basesize + 4 * self.ssize_t)
2142
2143 def test_combinations_with_replacement_sizeof(self):
2144 cwr = combinations_with_replacement
2145 basesize = support.calcobjsize('3Pni')
2146 check = self.check_sizeof
2147 check(cwr('abcd', 3), basesize + 3 * self.ssize_t)
2148 check(cwr(range(10), 4), basesize + 4 * self.ssize_t)
2149
2150 def test_permutations_sizeof(self):
2151 basesize = support.calcobjsize('4Pni')
2152 check = self.check_sizeof
2153 check(permutations('abcd'),
2154 basesize + 4 * self.ssize_t + 4 * self.ssize_t)
2155 check(permutations('abcd', 3),
2156 basesize + 4 * self.ssize_t + 3 * self.ssize_t)
2157 check(permutations('abcde', 3),
2158 basesize + 5 * self.ssize_t + 3 * self.ssize_t)
2159 check(permutations(range(10), 4),
2160 basesize + 10 * self.ssize_t + 4 * self.ssize_t)
2161
Thomas Woutersb2137042007-02-01 18:02:27 +00002162
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002163libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002164
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002165
2166>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002167>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00002168... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00002169...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002170Check 1200 is for $120.15
2171Check 1201 is for $764.05
2172Check 1202 is for $823.14
2173
2174>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002175>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00002176... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00002177...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000021781
21798
218027
2181
2182>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00002183>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00002184... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00002185...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002186Alex
2187Laura
2188Martin
2189Walter
2190Samuele
2191
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002192>>> from operator import itemgetter
2193>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00002194>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002195>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00002196... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00002197...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000021981 ['a', 'c', 'e']
21992 ['b', 'd', 'f']
22003 ['g']
2201
Raymond Hettinger734fb572004-01-20 20:04:40 +00002202# Find runs of consecutive numbers using groupby. The key to the solution
2203# is differencing with a range so that consecutive numbers all appear in
2204# same group.
2205>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00002206>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00002207... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00002208...
Raymond Hettinger734fb572004-01-20 20:04:40 +00002209[1]
2210[4, 5, 6]
2211[10]
2212[15, 16, 17, 18]
2213[22]
2214[25, 26, 27, 28]
2215
Georg Brandl3dbca812008-07-23 16:10:53 +00002216>>> def take(n, iterable):
2217... "Return first n items of the iterable as a list"
2218... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00002219
Raymond Hettinger9265dd72018-04-08 08:44:20 -07002220>>> def prepend(value, iterator):
2221... "Prepend a single value in front of an iterator"
2222... # prepend(1, [2, 3, 4]) -> 1 2 3 4
2223... return chain([value], iterator)
2224
Georg Brandl3dbca812008-07-23 16:10:53 +00002225>>> def enumerate(iterable, start=0):
2226... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002227
Georg Brandl3dbca812008-07-23 16:10:53 +00002228>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002229... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00002230... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002231
Cheryl Sabellada1734c2018-03-26 21:29:33 -04002232>>> import collections
2233>>> def consume(iterator, n=None):
2234... "Advance the iterator n-steps ahead. If n is None, consume entirely."
2235... # Use functions that consume iterators at C speed.
2236... if n is None:
2237... # feed the entire iterator into a zero-length deque
2238... collections.deque(iterator, maxlen=0)
2239... else:
2240... # advance to the empty slice starting at position n
2241... next(islice(iterator, n, n), None)
2242
Raymond Hettingercdf8ba32009-02-19 04:45:07 +00002243>>> def nth(iterable, n, default=None):
2244... "Returns the nth item or a default value"
2245... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002246
Raymond Hettingere525ee32016-03-06 18:11:38 -08002247>>> def all_equal(iterable):
2248... "Returns True if all the elements are equal to each other"
2249... g = groupby(iterable)
2250... return next(g, True) and not next(g, False)
2251
Georg Brandl3dbca812008-07-23 16:10:53 +00002252>>> def quantify(iterable, pred=bool):
2253... "Count how many times the predicate is true"
2254... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00002255
Georg Brandl3dbca812008-07-23 16:10:53 +00002256>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002257... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00002258... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002259
Georg Brandl3dbca812008-07-23 16:10:53 +00002260>>> def ncycles(iterable, n):
Ezio Melotti13925002011-03-16 11:05:33 +02002261... "Returns the sequence elements n times"
Georg Brandl3dbca812008-07-23 16:10:53 +00002262... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002263
2264>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002265... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002266
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002267>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002268... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002269
2270>>> def repeatfunc(func, times=None, *args):
2271... "Repeat calls to func with specified arguments."
2272... " Example: repeatfunc(random.random)"
2273... if times is None:
2274... return starmap(func, repeat(args))
2275... else:
2276... return starmap(func, repeat(args, times))
2277
Raymond Hettingerd591f662003-10-26 15:34:50 +00002278>>> def pairwise(iterable):
2279... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
2280... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00002281... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00002282... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00002283... except StopIteration:
2284... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002285... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002286
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002287>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00002288... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002289... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00002290... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002291
2292>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00002293... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002294... # Recipe credited to George Sakkis
2295... pending = len(iterables)
2296... nexts = cycle(iter(it).__next__ for it in iterables)
2297... while pending:
2298... try:
2299... for next in nexts:
2300... yield next()
2301... except StopIteration:
2302... pending -= 1
2303... nexts = cycle(islice(nexts, pending))
2304
2305>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00002306... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
2307... s = list(iterable)
2308... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002309
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00002310>>> def unique_everseen(iterable, key=None):
2311... "List unique elements, preserving order. Remember all elements ever seen."
2312... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
2313... # unique_everseen('ABBCcAD', str.lower) --> A B C D
2314... seen = set()
2315... seen_add = seen.add
2316... if key is None:
2317... for element in iterable:
2318... if element not in seen:
2319... seen_add(element)
2320... yield element
2321... else:
2322... for element in iterable:
2323... k = key(element)
2324... if k not in seen:
2325... seen_add(k)
2326... yield element
2327
2328>>> def unique_justseen(iterable, key=None):
2329... "List unique elements, preserving order. Remember only the element just seen."
2330... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
2331... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
2332... return map(next, map(itemgetter(1), groupby(iterable, key)))
2333
Raymond Hettinger31b26f62014-04-02 03:16:42 -07002334>>> def first_true(iterable, default=False, pred=None):
2335... '''Returns the first true value in the iterable.
2336...
2337... If no true value is found, returns *default*
2338...
2339... If *pred* is not None, returns the first item
2340... for which pred(item) is true.
2341...
2342... '''
2343... # first_true([a,b,c], x) --> a or b or c or x
2344... # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x
2345... return next(filter(pred, iterable), default)
2346
Raymond Hettingerd37258d2018-01-13 10:35:40 -08002347>>> def nth_combination(iterable, r, index):
2348... 'Equivalent to list(combinations(iterable, r))[index]'
2349... pool = tuple(iterable)
2350... n = len(pool)
2351... if r < 0 or r > n:
2352... raise ValueError
2353... c = 1
2354... k = min(r, n-r)
2355... for i in range(1, k+1):
2356... c = c * (n - k + i) // i
2357... if index < 0:
2358... index += c
2359... if index < 0 or index >= c:
2360... raise IndexError
2361... result = []
2362... while r:
2363... c, n, r = c*r//n, n-1, r-1
2364... while index >= c:
2365... index -= c
2366... c, n = c*(n-r)//n, n-1
2367... result.append(pool[-1-n])
2368... return tuple(result)
2369
2370
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002371This is not part of the examples but it tests to make sure the definitions
2372perform as purported.
2373
Raymond Hettingera098b332003-09-08 23:58:40 +00002374>>> take(10, count())
2375[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
2376
Raymond Hettinger9265dd72018-04-08 08:44:20 -07002377>>> list(prepend(1, [2, 3, 4]))
2378[1, 2, 3, 4]
2379
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002380>>> list(enumerate('abc'))
2381[(0, 'a'), (1, 'b'), (2, 'c')]
2382
2383>>> list(islice(tabulate(lambda x: 2*x), 4))
2384[0, 2, 4, 6]
2385
Cheryl Sabellada1734c2018-03-26 21:29:33 -04002386>>> it = iter(range(10))
2387>>> consume(it, 3)
2388>>> next(it)
23893
2390>>> consume(it)
2391>>> next(it, 'Done')
2392'Done'
2393
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002394>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00002395'd'
2396
2397>>> nth('abcde', 9) is None
2398True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002399
Raymond Hettingere525ee32016-03-06 18:11:38 -08002400>>> [all_equal(s) for s in ('', 'A', 'AAAA', 'AAAB', 'AAABA')]
2401[True, True, True, False, False]
2402
Guido van Rossum805365e2007-05-07 22:24:25 +00002403>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000240450
2405
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002406>>> a = [[1, 2, 3], [4, 5, 6]]
2407>>> flatten(a)
2408[1, 2, 3, 4, 5, 6]
2409
2410>>> list(repeatfunc(pow, 5, 2, 3))
2411[8, 8, 8, 8, 8]
2412
2413>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002414>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002415[0, 0, 0, 0, 0]
2416
Raymond Hettingerd591f662003-10-26 15:34:50 +00002417>>> list(pairwise('abcd'))
2418[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002419
Raymond Hettingerd591f662003-10-26 15:34:50 +00002420>>> list(pairwise([]))
2421[]
2422
2423>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00002424[]
2425
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002426>>> list(islice(padnone('abc'), 0, 6))
2427['a', 'b', 'c', None, None, None]
2428
2429>>> list(ncycles('abc', 3))
2430['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
2431
2432>>> dotproduct([1,2,3], [4,5,6])
243332
2434
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002435>>> list(grouper(3, 'abcdefg', 'x'))
2436[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
2437
2438>>> list(roundrobin('abc', 'd', 'ef'))
2439['a', 'd', 'e', 'b', 'f', 'c']
2440
Raymond Hettingerace67332009-01-26 02:23:50 +00002441>>> list(powerset([1,2,3]))
2442[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002443
Raymond Hettinger191e8502009-01-27 13:29:43 +00002444>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
2445True
2446
2447>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
2448True
2449
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00002450>>> list(unique_everseen('AAAABBBCCDAABBB'))
2451['A', 'B', 'C', 'D']
2452
2453>>> list(unique_everseen('ABBCcAD', str.lower))
2454['A', 'B', 'C', 'D']
2455
2456>>> list(unique_justseen('AAAABBBCCDAABBB'))
2457['A', 'B', 'C', 'D', 'A', 'B']
2458
2459>>> list(unique_justseen('ABBCcAD', str.lower))
2460['A', 'B', 'C', 'A', 'D']
2461
Raymond Hettinger31b26f62014-04-02 03:16:42 -07002462>>> first_true('ABC0DEF1', '9', str.isdigit)
2463'0'
2464
Raymond Hettingerd37258d2018-01-13 10:35:40 -08002465>>> population = 'ABCDEFGH'
2466>>> for r in range(len(population) + 1):
2467... seq = list(combinations(population, r))
2468... for i in range(len(seq)):
2469... assert nth_combination(population, r, i) == seq[i]
2470... for i in range(-len(seq), 0):
2471... assert nth_combination(population, r, i) == seq[i]
2472
2473
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002474"""
2475
2476__test__ = {'libreftest' : libreftest}
2477
2478def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002479 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00002480 RegressionTests, LengthTransparency,
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002481 SubclassWithKwargsTest, TestExamples,
Cheryl Sabellada1734c2018-03-26 21:29:33 -04002482 TestPurePythonRoughEquivalents,
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002483 SizeofTest)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002484 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002485
2486 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002487 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00002488 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00002489 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00002490 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002491 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00002492 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00002493 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00002494 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002495
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002496 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002497 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002498
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002499if __name__ == "__main__":
2500 test_main(verbose=True)