blob: ea060a98a5eef497e4cc129544c24eadb38842b3 [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
Raymond Hettingerfc438512009-11-01 20:55:33 +0000974 def test_bug_7244(self):
975
976 class Repeater:
977 # this class is similar to itertools.repeat
978 def __init__(self, o, t, e):
979 self.o = o
980 self.t = int(t)
981 self.e = e
982 def __iter__(self): # its iterator is itself
983 return self
984 def __next__(self):
985 if self.t > 0:
986 self.t -= 1
987 return self.o
988 else:
989 raise self.e
990
991 # Formerly this code in would fail in debug mode
992 # with Undetected Error and Stop Iteration
993 r1 = Repeater(1, 3, StopIteration)
994 r2 = Repeater(2, 4, StopIteration)
995 def run(r1, r2):
996 result = []
997 for i, j in zip_longest(r1, r2, fillvalue=0):
998 with support.captured_output('stdout'):
999 print((i, j))
1000 result.append((i, j))
1001 return result
1002 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
1003
1004 # Formerly, the RuntimeError would be lost
1005 # and StopIteration would stop as expected
1006 r1 = Repeater(1, 3, RuntimeError)
1007 r2 = Repeater(2, 4, StopIteration)
1008 it = zip_longest(r1, r2, fillvalue=0)
1009 self.assertEqual(next(it), (1, 2))
1010 self.assertEqual(next(it), (1, 2))
1011 self.assertEqual(next(it), (1, 2))
1012 self.assertRaises(RuntimeError, next, it)
1013
Christian Heimesc3f30c42008-02-22 16:37:40 +00001014 def test_product(self):
1015 for args, result in [
Christian Heimes78644762008-03-04 23:39:23 +00001016 ([], [()]), # zero iterables
Christian Heimesc3f30c42008-02-22 16:37:40 +00001017 (['ab'], [('a',), ('b',)]), # one iterable
1018 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
1019 ([range(0), range(2), range(3)], []), # first iterable with zero length
1020 ([range(2), range(0), range(3)], []), # middle iterable with zero length
1021 ([range(2), range(3), range(0)], []), # last iterable with zero length
1022 ]:
1023 self.assertEqual(list(product(*args)), result)
Christian Heimesf16baeb2008-02-29 14:57:44 +00001024 for r in range(4):
1025 self.assertEqual(list(product(*(args*r))),
1026 list(product(*args, **dict(repeat=r))))
Christian Heimesc3f30c42008-02-22 16:37:40 +00001027 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
1028 self.assertRaises(TypeError, product, range(6), None)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001029
1030 def product1(*args, **kwds):
1031 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
1032 n = len(pools)
1033 if n == 0:
1034 yield ()
1035 return
1036 if any(len(pool) == 0 for pool in pools):
1037 return
1038 indices = [0] * n
1039 yield tuple(pool[i] for pool, i in zip(pools, indices))
1040 while 1:
1041 for i in reversed(range(n)): # right to left
1042 if indices[i] == len(pools[i]) - 1:
1043 continue
1044 indices[i] += 1
1045 for j in range(i+1, n):
1046 indices[j] = 0
1047 yield tuple(pool[i] for pool, i in zip(pools, indices))
1048 break
1049 else:
1050 return
1051
1052 def product2(*args, **kwds):
1053 'Pure python version used in docs'
1054 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
1055 result = [[]]
1056 for pool in pools:
1057 result = [x+[y] for x in result for y in pool]
1058 for prod in result:
1059 yield tuple(prod)
1060
Christian Heimesc3f30c42008-02-22 16:37:40 +00001061 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
1062 set('abcdefg'), range(11), tuple(range(13))]
1063 for i in range(100):
1064 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Christian Heimes78644762008-03-04 23:39:23 +00001065 expected_len = prod(map(len, args))
1066 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001067 self.assertEqual(list(product(*args)), list(product1(*args)))
1068 self.assertEqual(list(product(*args)), list(product2(*args)))
Christian Heimesc3f30c42008-02-22 16:37:40 +00001069 args = map(iter, args)
Christian Heimes78644762008-03-04 23:39:23 +00001070 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesc3f30c42008-02-22 16:37:40 +00001071
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05001072 @support.bigaddrspacetest
1073 def test_product_overflow(self):
Serhiy Storchakadee948b2015-02-03 01:34:09 +02001074 with self.assertRaises((OverflowError, MemoryError)):
1075 product(*(['ab']*2**5), repeat=2**25)
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05001076
Alex Gaynore151d212011-07-17 16:21:30 -07001077 @support.impl_detail("tuple reuse is specific to CPython")
1078 def test_product_tuple_reuse(self):
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001079 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
1080 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +00001081
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001082 def test_product_pickling(self):
1083 # check copy, deepcopy, pickle
1084 for args, result in [
1085 ([], [()]), # zero iterables
1086 (['ab'], [('a',), ('b',)]), # one iterable
1087 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
1088 ([range(0), range(2), range(3)], []), # first iterable with zero length
1089 ([range(2), range(0), range(3)], []), # middle iterable with zero length
1090 ([range(2), range(3), range(0)], []), # last iterable with zero length
1091 ]:
1092 self.assertEqual(list(copy.copy(product(*args))), result)
1093 self.assertEqual(list(copy.deepcopy(product(*args))), result)
Serhiy Storchakabad12572014-12-15 14:03:42 +02001094 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1095 self.pickletest(proto, product(*args))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001096
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00001097 def test_product_issue_25021(self):
1098 # test that indices are properly clamped to the length of the tuples
1099 p = product((1, 2),(3,))
1100 p.__setstate__((0, 0x1000)) # will access tuple element 1 if not clamped
1101 self.assertEqual(next(p), (2, 3))
1102 # test that empty tuple in the list will result in an immediate StopIteration
1103 p = product((1, 2), (), (3,))
1104 p.__setstate__((0, 0, 0x1000)) # will access tuple element 1 if not clamped
1105 self.assertRaises(StopIteration, next, p)
1106
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001107 def test_repeat(self):
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00001108 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
Guido van Rossum805365e2007-05-07 22:24:25 +00001109 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001110 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001111 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +00001112 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001113 self.assertEqual(list(repeat('a', 0)), [])
1114 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001115 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001116 self.assertRaises(TypeError, repeat, None, 3, 4)
1117 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001118 r = repeat(1+0j)
1119 self.assertEqual(repr(r), 'repeat((1+0j))')
1120 r = repeat(1+0j, 5)
1121 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
1122 list(r)
1123 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001124
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001125 # check copy, deepcopy, pickle
1126 c = repeat(object='a', times=10)
1127 self.assertEqual(next(c), 'a')
1128 self.assertEqual(take(2, copy.copy(c)), list('a' * 2))
1129 self.assertEqual(take(2, copy.deepcopy(c)), list('a' * 2))
Serhiy Storchakabad12572014-12-15 14:03:42 +02001130 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1131 self.pickletest(proto, repeat(object='a', times=10))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001132
Raymond Hettinger97d35552014-06-24 21:36:58 -07001133 def test_repeat_with_negative_times(self):
1134 self.assertEqual(repr(repeat('a', -1)), "repeat('a', 0)")
1135 self.assertEqual(repr(repeat('a', -2)), "repeat('a', 0)")
1136 self.assertEqual(repr(repeat('a', times=-1)), "repeat('a', 0)")
1137 self.assertEqual(repr(repeat('a', times=-2)), "repeat('a', 0)")
1138
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001139 def test_map(self):
1140 self.assertEqual(list(map(operator.pow, range(3), range(1,7))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001141 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001142 self.assertEqual(list(map(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001143 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001144 self.assertEqual(list(map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +00001145 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001146 self.assertEqual(take(2,map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +00001147 [('a',0),('b',1)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001148 self.assertEqual(list(map(operator.pow, [])), [])
1149 self.assertRaises(TypeError, map)
1150 self.assertRaises(TypeError, list, map(None, range(3), range(3)))
1151 self.assertRaises(TypeError, map, operator.neg)
1152 self.assertRaises(TypeError, next, map(10, range(5)))
1153 self.assertRaises(ValueError, next, map(errfunc, [4], [5]))
1154 self.assertRaises(TypeError, next, map(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001155
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001156 # check copy, deepcopy, pickle
1157 ans = [('a',0),('b',1),('c',2)]
1158
1159 c = map(tupleize, 'abc', count())
1160 self.assertEqual(list(copy.copy(c)), ans)
1161
1162 c = map(tupleize, 'abc', count())
1163 self.assertEqual(list(copy.deepcopy(c)), ans)
1164
Serhiy Storchakabad12572014-12-15 14:03:42 +02001165 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1166 c = map(tupleize, 'abc', count())
1167 self.pickletest(proto, c)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001168
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001169 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001170 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
1171 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001172 self.assertEqual(take(3, starmap(operator.pow, zip(count(), count(1)))),
Raymond Hettinger02420702003-06-29 20:36:23 +00001173 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001174 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +00001175 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
1176 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001177 self.assertRaises(TypeError, starmap)
1178 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001179 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
1180 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
1181 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001182
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001183 # check copy, deepcopy, pickle
1184 ans = [0**1, 1**2, 2**3]
1185
1186 c = starmap(operator.pow, zip(range(3), range(1,7)))
1187 self.assertEqual(list(copy.copy(c)), ans)
1188
1189 c = starmap(operator.pow, zip(range(3), range(1,7)))
1190 self.assertEqual(list(copy.deepcopy(c)), ans)
1191
Serhiy Storchakabad12572014-12-15 14:03:42 +02001192 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1193 c = starmap(operator.pow, zip(range(3), range(1,7)))
1194 self.pickletest(proto, c)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001195
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001196 def test_islice(self):
1197 for args in [ # islice(args) should agree with range(args)
1198 (10, 20, 3),
1199 (10, 3, 20),
1200 (10, 20),
Cheryl Sabellada1734c2018-03-26 21:29:33 -04001201 (10, 10),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001202 (10, 3),
1203 (20,)
1204 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +00001205 self.assertEqual(list(islice(range(100), *args)),
1206 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001207
1208 for args, tgtargs in [ # Stop when seqn is exhausted
1209 ((10, 110, 3), ((10, 100, 3))),
1210 ((10, 110), ((10, 100))),
1211 ((110,), (100,))
1212 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +00001213 self.assertEqual(list(islice(range(100), *args)),
1214 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001215
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001216 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +00001217 self.assertEqual(list(islice(range(10), None)), list(range(10)))
1218 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
1219 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
1220 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
1221 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001222
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +00001223 # Test number of items consumed SF #1171417
1224 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +00001225 self.assertEqual(list(islice(it, 3)), list(range(3)))
1226 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +00001227
Cheryl Sabellada1734c2018-03-26 21:29:33 -04001228 it = iter(range(10))
1229 self.assertEqual(list(islice(it, 3, 3)), [])
1230 self.assertEqual(list(it), list(range(3, 10)))
1231
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001232 # Test invalid arguments
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001233 ra = range(10)
1234 self.assertRaises(TypeError, islice, ra)
1235 self.assertRaises(TypeError, islice, ra, 1, 2, 3, 4)
1236 self.assertRaises(ValueError, islice, ra, -5, 10, 1)
1237 self.assertRaises(ValueError, islice, ra, 1, -5, -1)
1238 self.assertRaises(ValueError, islice, ra, 1, 10, -1)
1239 self.assertRaises(ValueError, islice, ra, 1, 10, 0)
1240 self.assertRaises(ValueError, islice, ra, 'a')
1241 self.assertRaises(ValueError, islice, ra, 'a', 1)
1242 self.assertRaises(ValueError, islice, ra, 1, 'a')
1243 self.assertRaises(ValueError, islice, ra, 'a', 1, 1)
1244 self.assertRaises(ValueError, islice, ra, 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +00001245 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001246
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001247 # Issue #10323: Less islice in a predictable state
1248 c = count()
1249 self.assertEqual(list(islice(c, 1, 3, 50)), [1])
1250 self.assertEqual(next(c), 3)
1251
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001252 # check copy, deepcopy, pickle
1253 for args in [ # islice(args) should agree with range(args)
1254 (10, 20, 3),
1255 (10, 3, 20),
1256 (10, 20),
1257 (10, 3),
1258 (20,)
1259 ]:
1260 self.assertEqual(list(copy.copy(islice(range(100), *args))),
1261 list(range(*args)))
1262 self.assertEqual(list(copy.deepcopy(islice(range(100), *args))),
1263 list(range(*args)))
Serhiy Storchakabad12572014-12-15 14:03:42 +02001264 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1265 self.pickletest(proto, islice(range(100), *args))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001266
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001267 # Issue #21321: check source iterator is not referenced
1268 # from islice() after the latter has been exhausted
1269 it = (x for x in (1, 2))
1270 wr = weakref.ref(it)
1271 it = islice(it, 1)
1272 self.assertIsNotNone(wr())
1273 list(it) # exhaust the iterator
Benjamin Peterson8e163512014-08-24 18:07:28 -05001274 support.gc_collect()
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001275 self.assertIsNone(wr())
1276
Will Roberts0ecdc522017-06-08 08:03:04 +02001277 # Issue #30537: islice can accept integer-like objects as
1278 # arguments
1279 class IntLike(object):
1280 def __init__(self, val):
1281 self.val = val
1282 def __index__(self):
1283 return self.val
1284 self.assertEqual(list(islice(range(100), IntLike(10))), list(range(10)))
1285 self.assertEqual(list(islice(range(100), IntLike(10), IntLike(50))),
1286 list(range(10, 50)))
1287 self.assertEqual(list(islice(range(100), IntLike(10), IntLike(50), IntLike(5))),
1288 list(range(10,50,5)))
1289
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001290 def test_takewhile(self):
1291 data = [1, 3, 5, 20, 2, 4, 6, 8]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001292 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001293 self.assertEqual(list(takewhile(underten, [])), [])
1294 self.assertRaises(TypeError, takewhile)
1295 self.assertRaises(TypeError, takewhile, operator.pow)
1296 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001297 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
1298 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001299 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
1300 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +00001301 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001302
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001303 # check copy, deepcopy, pickle
1304 self.assertEqual(list(copy.copy(takewhile(underten, data))), [1, 3, 5])
1305 self.assertEqual(list(copy.deepcopy(takewhile(underten, data))),
1306 [1, 3, 5])
Serhiy Storchakabad12572014-12-15 14:03:42 +02001307 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1308 self.pickletest(proto, takewhile(underten, data))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001309
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001310 def test_dropwhile(self):
1311 data = [1, 3, 5, 20, 2, 4, 6, 8]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001312 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001313 self.assertEqual(list(dropwhile(underten, [])), [])
1314 self.assertRaises(TypeError, dropwhile)
1315 self.assertRaises(TypeError, dropwhile, operator.pow)
1316 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001317 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
1318 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001319
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001320 # check copy, deepcopy, pickle
1321 self.assertEqual(list(copy.copy(dropwhile(underten, data))), [20, 2, 4, 6, 8])
1322 self.assertEqual(list(copy.deepcopy(dropwhile(underten, data))),
1323 [20, 2, 4, 6, 8])
Serhiy Storchakabad12572014-12-15 14:03:42 +02001324 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1325 self.pickletest(proto, dropwhile(underten, data))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001326
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001327 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +00001328 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001329
1330 a, b = tee([]) # test empty iterator
1331 self.assertEqual(list(a), [])
1332 self.assertEqual(list(b), [])
1333
1334 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +00001335 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001336
1337 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +00001338 self.assertEqual(list(a), list(range(n)))
1339 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001340
1341 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +00001342 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001343 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001344 del a
Guido van Rossum805365e2007-05-07 22:24:25 +00001345 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001346
1347 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +00001348 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001349 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001350 del b
Guido van Rossum805365e2007-05-07 22:24:25 +00001351 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001352
Guido van Rossum805365e2007-05-07 22:24:25 +00001353 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001354 order = [0]*n + [1]*n
1355 random.shuffle(order)
1356 lists = ([], [])
1357 its = tee(irange(n))
1358 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +00001359 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001360 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +00001361 self.assertEqual(lists[0], list(range(n)))
1362 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001363
Raymond Hettingerad983e72003-11-12 14:32:26 +00001364 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001365 self.assertRaises(TypeError, tee)
1366 self.assertRaises(TypeError, tee, 3)
1367 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +00001368 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001369
Raymond Hettingerad983e72003-11-12 14:32:26 +00001370 # tee object should be instantiable
1371 a, b = tee('abc')
1372 c = type(a)('def')
1373 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +00001374
Raymond Hettingerad983e72003-11-12 14:32:26 +00001375 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +00001376 a, b, c = tee(range(2000), 3)
1377 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001378 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +00001379 self.assertEqual(list(b), list(range(2000)))
1380 self.assertEqual([next(c), next(c)], list(range(2)))
1381 self.assertEqual(list(a), list(range(100,2000)))
1382 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +00001383
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001384 # test values of n
1385 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +00001386 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +00001387 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001388 result = tee('abc', n)
1389 self.assertEqual(type(result), tuple)
1390 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001391 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001392
Raymond Hettingerad983e72003-11-12 14:32:26 +00001393 # tee pass-through to copyable iterator
1394 a, b = tee('abc')
1395 c, d = tee(a)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +00001396 self.assertTrue(a is c)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001397
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001398 # test tee_new
1399 t1, t2 = tee('abc')
1400 tnew = type(t1)
1401 self.assertRaises(TypeError, tnew)
1402 self.assertRaises(TypeError, tnew, 10)
1403 t3 = tnew(t1)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +00001404 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +00001405
Raymond Hettingera9f60922004-10-17 16:40:14 +00001406 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +00001407 a, b = tee(range(10))
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001408 p = weakref.proxy(a)
Raymond Hettingera9f60922004-10-17 16:40:14 +00001409 self.assertEqual(getattr(p, '__class__'), type(b))
1410 del a
1411 self.assertRaises(ReferenceError, getattr, p, '__class__')
1412
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001413 ans = list('abc')
1414 long_ans = list(range(10000))
1415
1416 # check copy
1417 a, b = tee('abc')
1418 self.assertEqual(list(copy.copy(a)), ans)
1419 self.assertEqual(list(copy.copy(b)), ans)
1420 a, b = tee(list(range(10000)))
1421 self.assertEqual(list(copy.copy(a)), long_ans)
1422 self.assertEqual(list(copy.copy(b)), long_ans)
1423
1424 # check partially consumed copy
1425 a, b = tee('abc')
1426 take(2, a)
1427 take(1, b)
1428 self.assertEqual(list(copy.copy(a)), ans[2:])
1429 self.assertEqual(list(copy.copy(b)), ans[1:])
1430 self.assertEqual(list(a), ans[2:])
1431 self.assertEqual(list(b), ans[1:])
1432 a, b = tee(range(10000))
1433 take(100, a)
1434 take(60, b)
1435 self.assertEqual(list(copy.copy(a)), long_ans[100:])
1436 self.assertEqual(list(copy.copy(b)), long_ans[60:])
1437 self.assertEqual(list(a), long_ans[100:])
1438 self.assertEqual(list(b), long_ans[60:])
1439
1440 # check deepcopy
1441 a, b = tee('abc')
1442 self.assertEqual(list(copy.deepcopy(a)), ans)
1443 self.assertEqual(list(copy.deepcopy(b)), ans)
1444 self.assertEqual(list(a), ans)
1445 self.assertEqual(list(b), ans)
1446 a, b = tee(range(10000))
1447 self.assertEqual(list(copy.deepcopy(a)), long_ans)
1448 self.assertEqual(list(copy.deepcopy(b)), long_ans)
1449 self.assertEqual(list(a), long_ans)
1450 self.assertEqual(list(b), long_ans)
1451
1452 # check partially consumed deepcopy
1453 a, b = tee('abc')
1454 take(2, a)
1455 take(1, b)
1456 self.assertEqual(list(copy.deepcopy(a)), ans[2:])
1457 self.assertEqual(list(copy.deepcopy(b)), ans[1:])
1458 self.assertEqual(list(a), ans[2:])
1459 self.assertEqual(list(b), ans[1:])
1460 a, b = tee(range(10000))
1461 take(100, a)
1462 take(60, b)
1463 self.assertEqual(list(copy.deepcopy(a)), long_ans[100:])
1464 self.assertEqual(list(copy.deepcopy(b)), long_ans[60:])
1465 self.assertEqual(list(a), long_ans[100:])
1466 self.assertEqual(list(b), long_ans[60:])
1467
1468 # check pickle
Serhiy Storchakabad12572014-12-15 14:03:42 +02001469 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1470 self.pickletest(proto, iter(tee('abc')))
1471 a, b = tee('abc')
1472 self.pickletest(proto, a, compare=ans)
1473 self.pickletest(proto, b, compare=ans)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001474
Serhiy Storchakaa3e91282013-01-25 13:19:31 +02001475 # Issue 13454: Crash when deleting backward iterator from tee()
1476 def test_tee_del_backward(self):
Serhiy Storchaka5bb893c2013-01-26 11:52:06 +02001477 forward, backward = tee(repeat(None, 20000000))
Serhiy Storchaka9db55002015-03-28 20:38:37 +02001478 try:
1479 any(forward) # exhaust the iterator
1480 del backward
1481 except:
1482 del forward, backward
1483 raise
Serhiy Storchakaa3e91282013-01-25 13:19:31 +02001484
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001485 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001486 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +00001487
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001488 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +00001489 self.assertRaises(StopIteration, next, f([]))
1490 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001491
Georg Brandla18af4e2007-04-21 15:47:16 +00001492 self.assertRaises(StopIteration, next, islice([], None))
1493 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001494
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001495 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +00001496 self.assertRaises(StopIteration, next, p)
1497 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001498 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +00001499 self.assertRaises(StopIteration, next, p)
1500 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001501
Georg Brandla18af4e2007-04-21 15:47:16 +00001502 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001503
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001504 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +00001505 self.assertRaises(StopIteration, next, f(lambda x:x, []))
1506 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001507
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001508class TestExamples(unittest.TestCase):
1509
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001510 def test_accumulate(self):
Raymond Hettinger482ba772010-12-01 22:48:00 +00001511 self.assertEqual(list(accumulate([1,2,3,4,5])), [1, 3, 6, 10, 15])
1512
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001513 def test_accumulate_reducible(self):
1514 # check copy, deepcopy, pickle
1515 data = [1, 2, 3, 4, 5]
1516 accumulated = [1, 3, 6, 10, 15]
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001517
Serhiy Storchakabad12572014-12-15 14:03:42 +02001518 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1519 it = accumulate(data)
1520 self.assertEqual(list(pickle.loads(pickle.dumps(it, proto))), accumulated[:])
1521 self.assertEqual(next(it), 1)
1522 self.assertEqual(list(pickle.loads(pickle.dumps(it, proto))), accumulated[1:])
1523 it = accumulate(data)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001524 self.assertEqual(next(it), 1)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001525 self.assertEqual(list(copy.deepcopy(it)), accumulated[1:])
1526 self.assertEqual(list(copy.copy(it)), accumulated[1:])
1527
Serhiy Storchakad5516252016-03-06 14:00:45 +02001528 def test_accumulate_reducible_none(self):
1529 # Issue #25718: total is None
1530 it = accumulate([None, None, None], operator.is_)
1531 self.assertEqual(next(it), None)
1532 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1533 it_copy = pickle.loads(pickle.dumps(it, proto))
1534 self.assertEqual(list(it_copy), [True, False])
1535 self.assertEqual(list(copy.deepcopy(it)), [True, False])
1536 self.assertEqual(list(copy.copy(it)), [True, False])
1537
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001538 def test_chain(self):
1539 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
1540
1541 def test_chain_from_iterable(self):
1542 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
1543
1544 def test_combinations(self):
1545 self.assertEqual(list(combinations('ABCD', 2)),
1546 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
1547 self.assertEqual(list(combinations(range(4), 3)),
1548 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
1549
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001550 def test_combinations_with_replacement(self):
1551 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
1552 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
1553
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001554 def test_compress(self):
1555 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
1556
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001557 def test_count(self):
1558 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
1559
1560 def test_cycle(self):
1561 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
1562
1563 def test_dropwhile(self):
1564 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
1565
1566 def test_groupby(self):
1567 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
1568 list('ABCDAB'))
1569 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
1570 [list('AAAA'), list('BBB'), list('CC'), list('D')])
1571
1572 def test_filter(self):
1573 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
1574
1575 def test_filterfalse(self):
1576 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
1577
1578 def test_map(self):
1579 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
1580
1581 def test_islice(self):
1582 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
1583 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
1584 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
1585 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1586
1587 def test_zip(self):
1588 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
1589
1590 def test_zip_longest(self):
1591 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
1592 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
1593
1594 def test_permutations(self):
1595 self.assertEqual(list(permutations('ABCD', 2)),
1596 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
1597 self.assertEqual(list(permutations(range(3))),
1598 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
1599
1600 def test_product(self):
1601 self.assertEqual(list(product('ABCD', 'xy')),
1602 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
1603 self.assertEqual(list(product(range(2), repeat=3)),
1604 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
1605 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
1606
1607 def test_repeat(self):
1608 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
1609
1610 def test_stapmap(self):
1611 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
1612 [32, 9, 1000])
1613
1614 def test_takewhile(self):
1615 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
1616
1617
Cheryl Sabellada1734c2018-03-26 21:29:33 -04001618class TestPurePythonRoughEquivalents(unittest.TestCase):
1619
1620 @staticmethod
1621 def islice(iterable, *args):
1622 s = slice(*args)
1623 start, stop, step = s.start or 0, s.stop or sys.maxsize, s.step or 1
1624 it = iter(range(start, stop, step))
1625 try:
1626 nexti = next(it)
1627 except StopIteration:
1628 # Consume *iterable* up to the *start* position.
1629 for i, element in zip(range(start), iterable):
1630 pass
1631 return
1632 try:
1633 for i, element in enumerate(iterable):
1634 if i == nexti:
1635 yield element
1636 nexti = next(it)
1637 except StopIteration:
1638 # Consume to *stop*.
1639 for i, element in zip(range(i + 1, stop), iterable):
1640 pass
1641
1642 def test_islice_recipe(self):
1643 self.assertEqual(list(self.islice('ABCDEFG', 2)), list('AB'))
1644 self.assertEqual(list(self.islice('ABCDEFG', 2, 4)), list('CD'))
1645 self.assertEqual(list(self.islice('ABCDEFG', 2, None)), list('CDEFG'))
1646 self.assertEqual(list(self.islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1647 # Test items consumed.
1648 it = iter(range(10))
1649 self.assertEqual(list(self.islice(it, 3)), list(range(3)))
1650 self.assertEqual(list(it), list(range(3, 10)))
1651 it = iter(range(10))
1652 self.assertEqual(list(self.islice(it, 3, 3)), [])
1653 self.assertEqual(list(it), list(range(3, 10)))
1654 # Test that slice finishes in predictable state.
1655 c = count()
1656 self.assertEqual(list(self.islice(c, 1, 3, 50)), [1])
1657 self.assertEqual(next(c), 3)
1658
1659
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001660class TestGC(unittest.TestCase):
1661
1662 def makecycle(self, iterator, container):
1663 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +00001664 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001665 del container, iterator
1666
Raymond Hettinger482ba772010-12-01 22:48:00 +00001667 def test_accumulate(self):
1668 a = []
1669 self.makecycle(accumulate([1,2,a,3]), a)
1670
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001671 def test_chain(self):
1672 a = []
1673 self.makecycle(chain(a), a)
1674
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001675 def test_chain_from_iterable(self):
1676 a = []
1677 self.makecycle(chain.from_iterable([a]), a)
1678
1679 def test_combinations(self):
1680 a = []
1681 self.makecycle(combinations([1,2,a,3], 3), a)
1682
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001683 def test_combinations_with_replacement(self):
1684 a = []
1685 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1686
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001687 def test_compress(self):
1688 a = []
1689 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1690
Raymond Hettingerd6280f42009-02-16 20:50:56 +00001691 def test_count(self):
1692 a = []
1693 Int = type('Int', (int,), dict(x=a))
1694 self.makecycle(count(Int(0), Int(1)), a)
1695
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001696 def test_cycle(self):
1697 a = []
1698 self.makecycle(cycle([a]*2), a)
1699
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001700 def test_dropwhile(self):
1701 a = []
1702 self.makecycle(dropwhile(bool, [0, a, a]), a)
1703
1704 def test_groupby(self):
1705 a = []
1706 self.makecycle(groupby([a]*2, lambda x:x), a)
1707
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001708 def test_issue2246(self):
1709 # Issue 2246 -- the _grouper iterator was not included in GC
1710 n = 10
1711 keyfunc = lambda x: x
1712 for i, j in groupby(range(n), key=keyfunc):
1713 keyfunc.__dict__.setdefault('x',[]).append(j)
1714
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001715 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001716 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001717 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001718
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001719 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001720 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001721 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001722
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001723 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001724 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001725 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001726
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001727 def test_zip_longest(self):
1728 a = []
1729 self.makecycle(zip_longest([a]*2, [a]*3), a)
1730 b = [a, None]
1731 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1732
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001733 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001734 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001735 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001736
1737 def test_islice(self):
1738 a = []
1739 self.makecycle(islice([a]*2, None), a)
1740
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001741 def test_permutations(self):
1742 a = []
1743 self.makecycle(permutations([1,2,a,3], 3), a)
1744
1745 def test_product(self):
1746 a = []
1747 self.makecycle(product([1,2,a,3], repeat=3), a)
1748
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001749 def test_repeat(self):
1750 a = []
1751 self.makecycle(repeat(a), a)
1752
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001753 def test_starmap(self):
1754 a = []
1755 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1756
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001757 def test_takewhile(self):
1758 a = []
1759 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1760
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001761def R(seqn):
1762 'Regular generator'
1763 for i in seqn:
1764 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001765
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001766class G:
1767 'Sequence using __getitem__'
1768 def __init__(self, seqn):
1769 self.seqn = seqn
1770 def __getitem__(self, i):
1771 return self.seqn[i]
1772
1773class I:
1774 'Sequence using iterator protocol'
1775 def __init__(self, seqn):
1776 self.seqn = seqn
1777 self.i = 0
1778 def __iter__(self):
1779 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001780 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001781 if self.i >= len(self.seqn): raise StopIteration
1782 v = self.seqn[self.i]
1783 self.i += 1
1784 return v
1785
1786class Ig:
1787 'Sequence using iterator protocol defined with a generator'
1788 def __init__(self, seqn):
1789 self.seqn = seqn
1790 self.i = 0
1791 def __iter__(self):
1792 for val in self.seqn:
1793 yield val
1794
1795class X:
1796 'Missing __getitem__ and __iter__'
1797 def __init__(self, seqn):
1798 self.seqn = seqn
1799 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001800 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001801 if self.i >= len(self.seqn): raise StopIteration
1802 v = self.seqn[self.i]
1803 self.i += 1
1804 return v
1805
1806class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001807 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001808 def __init__(self, seqn):
1809 self.seqn = seqn
1810 self.i = 0
1811 def __iter__(self):
1812 return self
1813
1814class E:
1815 'Test propagation of exceptions'
1816 def __init__(self, seqn):
1817 self.seqn = seqn
1818 self.i = 0
1819 def __iter__(self):
1820 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001821 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001822 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001823
1824class S:
1825 'Test immediate stop'
1826 def __init__(self, seqn):
1827 pass
1828 def __iter__(self):
1829 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001830 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001831 raise StopIteration
1832
1833def L(seqn):
1834 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001835 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001836
1837
1838class TestVariousIteratorArgs(unittest.TestCase):
1839
Raymond Hettinger482ba772010-12-01 22:48:00 +00001840 def test_accumulate(self):
1841 s = [1,2,3,4,5]
1842 r = [1,3,6,10,15]
1843 n = len(s)
1844 for g in (G, I, Ig, L, R):
1845 self.assertEqual(list(accumulate(g(s))), r)
1846 self.assertEqual(list(accumulate(S(s))), [])
1847 self.assertRaises(TypeError, accumulate, X(s))
1848 self.assertRaises(TypeError, accumulate, N(s))
1849 self.assertRaises(ZeroDivisionError, list, accumulate(E(s)))
1850
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001851 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001852 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001853 for g in (G, I, Ig, S, L, R):
1854 self.assertEqual(list(chain(g(s))), list(g(s)))
1855 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001856 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001857 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001858 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1859
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001860 def test_compress(self):
1861 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1862 n = len(s)
1863 for g in (G, I, Ig, S, L, R):
1864 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1865 self.assertRaises(TypeError, compress, X(s), repeat(1))
1866 self.assertRaises(TypeError, compress, N(s), repeat(1))
1867 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1868
Christian Heimesc3f30c42008-02-22 16:37:40 +00001869 def test_product(self):
1870 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1871 self.assertRaises(TypeError, product, X(s))
1872 self.assertRaises(TypeError, product, N(s))
1873 self.assertRaises(ZeroDivisionError, product, E(s))
1874
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001875 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001876 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001877 for g in (G, I, Ig, S, L, R):
1878 tgtlen = len(s) * 3
1879 expected = list(g(s))*3
1880 actual = list(islice(cycle(g(s)), tgtlen))
1881 self.assertEqual(actual, expected)
1882 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001883 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001884 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1885
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001886 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001887 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001888 for g in (G, I, Ig, S, L, R):
1889 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1890 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001891 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001892 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1893
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001894 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001895 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001896 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001897 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001898 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001899 self.assertRaises(TypeError, filter, isEven, X(s))
1900 self.assertRaises(TypeError, filter, isEven, N(s))
1901 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001902
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001903 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001904 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001905 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001906 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001907 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001908 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1909 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1910 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001911
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001912 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001913 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001914 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001915 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1916 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1917 self.assertRaises(TypeError, zip, X(s))
1918 self.assertRaises(TypeError, zip, N(s))
1919 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001920
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001921 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001922 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001923 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001924 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1925 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1926 self.assertRaises(TypeError, zip_longest, X(s))
1927 self.assertRaises(TypeError, zip_longest, N(s))
1928 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001929
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001930 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001931 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001932 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001933 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001934 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001935 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001936 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001937 self.assertRaises(TypeError, map, onearg, X(s))
1938 self.assertRaises(TypeError, map, onearg, N(s))
1939 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001940
1941 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001942 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001943 for g in (G, I, Ig, S, L, R):
1944 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1945 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001946 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001947 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1948
1949 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001950 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001951 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001952 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001953 self.assertEqual(list(starmap(operator.pow, g(ss))),
1954 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001955 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001956 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001957 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1958
1959 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001960 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001961 for g in (G, I, Ig, S, L, R):
1962 tgt = []
1963 for elem in g(s):
1964 if not isEven(elem): break
1965 tgt.append(elem)
1966 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1967 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001968 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001969 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1970
1971 def test_dropwhile(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 tgt and isOdd(elem): continue
1977 tgt.append(elem)
1978 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1979 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001980 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001981 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1982
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001983 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001984 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001985 for g in (G, I, Ig, S, L, R):
1986 it1, it2 = tee(g(s))
1987 self.assertEqual(list(it1), list(g(s)))
1988 self.assertEqual(list(it2), list(g(s)))
1989 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001990 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001991 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1992
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001993class LengthTransparency(unittest.TestCase):
1994
1995 def test_repeat(self):
Armin Ronacheraa9a79d2012-10-06 14:03:24 +02001996 self.assertEqual(operator.length_hint(repeat(None, 50)), 50)
Raymond Hettinger97d35552014-06-24 21:36:58 -07001997 self.assertEqual(operator.length_hint(repeat(None, 0)), 0)
Armin Ronacheraa9a79d2012-10-06 14:03:24 +02001998 self.assertEqual(operator.length_hint(repeat(None), 12), 12)
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001999
Raymond Hettinger97d35552014-06-24 21:36:58 -07002000 def test_repeat_with_negative_times(self):
2001 self.assertEqual(operator.length_hint(repeat(None, -1)), 0)
2002 self.assertEqual(operator.length_hint(repeat(None, -2)), 0)
2003 self.assertEqual(operator.length_hint(repeat(None, times=-1)), 0)
2004 self.assertEqual(operator.length_hint(repeat(None, times=-2)), 0)
2005
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002006class RegressionTests(unittest.TestCase):
2007
2008 def test_sf_793826(self):
2009 # Fix Armin Rigo's successful efforts to wreak havoc
2010
2011 def mutatingtuple(tuple1, f, tuple2):
2012 # this builds a tuple t which is a copy of tuple1,
2013 # then calls f(t), then mutates t to be equal to tuple2
2014 # (needs len(tuple1) == len(tuple2)).
2015 def g(value, first=[1]):
2016 if first:
2017 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00002018 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002019 return value
2020 items = list(tuple2)
2021 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002022 gen = map(g, items)
2023 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00002024 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002025
2026 def f(t):
2027 global T
2028 T = t
2029 first[:] = list(T)
2030
2031 first = []
2032 mutatingtuple((1,2,3), f, (4,5,6))
2033 second = list(T)
2034 self.assertEqual(first, second)
2035
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002036
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00002037 def test_sf_950057(self):
2038 # Make sure that chain() and cycle() catch exceptions immediately
2039 # rather than when shifting between input sources
2040
2041 def gen1():
2042 hist.append(0)
2043 yield 1
2044 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00002045 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00002046 hist.append(2)
2047
2048 def gen2(x):
2049 hist.append(3)
2050 yield 2
2051 hist.append(4)
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00002052
2053 hist = []
2054 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
2055 self.assertEqual(hist, [0,1])
2056
2057 hist = []
2058 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
2059 self.assertEqual(hist, [0,1])
2060
2061 hist = []
2062 self.assertRaises(AssertionError, list, cycle(gen1()))
2063 self.assertEqual(hist, [0,1])
2064
T. Wouters5466d4a2017-03-30 09:58:35 -07002065 def test_long_chain_of_empty_iterables(self):
2066 # Make sure itertools.chain doesn't run into recursion limits when
2067 # dealing with long chains of empty iterables. Even with a high
2068 # number this would probably only fail in Py_DEBUG mode.
2069 it = chain.from_iterable(() for unused in range(10000000))
2070 with self.assertRaises(StopIteration):
2071 next(it)
2072
Serhiy Storchakac740e4f2017-09-26 21:47:56 +03002073 def test_issue30347_1(self):
2074 def f(n):
2075 if n == 5:
2076 list(b)
2077 return n != 6
2078 for (k, b) in groupby(range(10), f):
2079 list(b) # shouldn't crash
2080
2081 def test_issue30347_2(self):
2082 class K:
2083 def __init__(self, v):
2084 pass
2085 def __eq__(self, other):
2086 nonlocal i
2087 i += 1
2088 if i == 1:
2089 next(g, None)
2090 return True
2091 i = 0
2092 g = next(groupby(range(10), K))[1]
2093 for j in range(2):
2094 next(g, None) # shouldn't crash
2095
2096
Thomas Woutersb2137042007-02-01 18:02:27 +00002097class SubclassWithKwargsTest(unittest.TestCase):
2098 def test_keywords_in_subclass(self):
2099 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002100 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002101 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00002102 class Subclass(cls):
2103 def __init__(self, newarg=None, *args):
2104 cls.__init__(self, *args)
2105 try:
2106 Subclass(newarg=1)
2107 except TypeError as err:
2108 # we expect type errors because of wrong argument count
Serhiy Storchaka5eb788b2017-06-06 18:45:22 +03002109 self.assertNotIn("keyword arguments", err.args[0])
Thomas Woutersb2137042007-02-01 18:02:27 +00002110
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002111@support.cpython_only
2112class SizeofTest(unittest.TestCase):
2113 def setUp(self):
2114 self.ssize_t = struct.calcsize('n')
2115
2116 check_sizeof = support.check_sizeof
2117
2118 def test_product_sizeof(self):
2119 basesize = support.calcobjsize('3Pi')
2120 check = self.check_sizeof
2121 check(product('ab', '12'), basesize + 2 * self.ssize_t)
2122 check(product(*(('abc',) * 10)), basesize + 10 * self.ssize_t)
2123
2124 def test_combinations_sizeof(self):
2125 basesize = support.calcobjsize('3Pni')
2126 check = self.check_sizeof
2127 check(combinations('abcd', 3), basesize + 3 * self.ssize_t)
2128 check(combinations(range(10), 4), basesize + 4 * self.ssize_t)
2129
2130 def test_combinations_with_replacement_sizeof(self):
2131 cwr = combinations_with_replacement
2132 basesize = support.calcobjsize('3Pni')
2133 check = self.check_sizeof
2134 check(cwr('abcd', 3), basesize + 3 * self.ssize_t)
2135 check(cwr(range(10), 4), basesize + 4 * self.ssize_t)
2136
2137 def test_permutations_sizeof(self):
2138 basesize = support.calcobjsize('4Pni')
2139 check = self.check_sizeof
2140 check(permutations('abcd'),
2141 basesize + 4 * self.ssize_t + 4 * self.ssize_t)
2142 check(permutations('abcd', 3),
2143 basesize + 4 * self.ssize_t + 3 * self.ssize_t)
2144 check(permutations('abcde', 3),
2145 basesize + 5 * self.ssize_t + 3 * self.ssize_t)
2146 check(permutations(range(10), 4),
2147 basesize + 10 * self.ssize_t + 4 * self.ssize_t)
2148
Thomas Woutersb2137042007-02-01 18:02:27 +00002149
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002150libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002151
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002152
2153>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002154>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00002155... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00002156...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002157Check 1200 is for $120.15
2158Check 1201 is for $764.05
2159Check 1202 is for $823.14
2160
2161>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002162>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00002163... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00002164...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000021651
21668
216727
2168
2169>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00002170>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00002171... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00002172...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002173Alex
2174Laura
2175Martin
2176Walter
2177Samuele
2178
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002179>>> from operator import itemgetter
2180>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00002181>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002182>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00002183... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00002184...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000021851 ['a', 'c', 'e']
21862 ['b', 'd', 'f']
21873 ['g']
2188
Raymond Hettinger734fb572004-01-20 20:04:40 +00002189# Find runs of consecutive numbers using groupby. The key to the solution
2190# is differencing with a range so that consecutive numbers all appear in
2191# same group.
2192>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00002193>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00002194... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00002195...
Raymond Hettinger734fb572004-01-20 20:04:40 +00002196[1]
2197[4, 5, 6]
2198[10]
2199[15, 16, 17, 18]
2200[22]
2201[25, 26, 27, 28]
2202
Georg Brandl3dbca812008-07-23 16:10:53 +00002203>>> def take(n, iterable):
2204... "Return first n items of the iterable as a list"
2205... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00002206
Raymond Hettinger9265dd72018-04-08 08:44:20 -07002207>>> def prepend(value, iterator):
2208... "Prepend a single value in front of an iterator"
2209... # prepend(1, [2, 3, 4]) -> 1 2 3 4
2210... return chain([value], iterator)
2211
Georg Brandl3dbca812008-07-23 16:10:53 +00002212>>> def enumerate(iterable, start=0):
2213... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002214
Georg Brandl3dbca812008-07-23 16:10:53 +00002215>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002216... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00002217... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002218
Cheryl Sabellada1734c2018-03-26 21:29:33 -04002219>>> import collections
2220>>> def consume(iterator, n=None):
2221... "Advance the iterator n-steps ahead. If n is None, consume entirely."
2222... # Use functions that consume iterators at C speed.
2223... if n is None:
2224... # feed the entire iterator into a zero-length deque
2225... collections.deque(iterator, maxlen=0)
2226... else:
2227... # advance to the empty slice starting at position n
2228... next(islice(iterator, n, n), None)
2229
Raymond Hettingercdf8ba32009-02-19 04:45:07 +00002230>>> def nth(iterable, n, default=None):
2231... "Returns the nth item or a default value"
2232... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002233
Raymond Hettingere525ee32016-03-06 18:11:38 -08002234>>> def all_equal(iterable):
2235... "Returns True if all the elements are equal to each other"
2236... g = groupby(iterable)
2237... return next(g, True) and not next(g, False)
2238
Georg Brandl3dbca812008-07-23 16:10:53 +00002239>>> def quantify(iterable, pred=bool):
2240... "Count how many times the predicate is true"
2241... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00002242
Georg Brandl3dbca812008-07-23 16:10:53 +00002243>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002244... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00002245... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002246
Georg Brandl3dbca812008-07-23 16:10:53 +00002247>>> def ncycles(iterable, n):
Ezio Melotti13925002011-03-16 11:05:33 +02002248... "Returns the sequence elements n times"
Georg Brandl3dbca812008-07-23 16:10:53 +00002249... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002250
2251>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002252... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002253
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002254>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002255... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002256
2257>>> def repeatfunc(func, times=None, *args):
2258... "Repeat calls to func with specified arguments."
2259... " Example: repeatfunc(random.random)"
2260... if times is None:
2261... return starmap(func, repeat(args))
2262... else:
2263... return starmap(func, repeat(args, times))
2264
Raymond Hettingerd591f662003-10-26 15:34:50 +00002265>>> def pairwise(iterable):
2266... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
2267... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00002268... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00002269... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00002270... except StopIteration:
2271... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002272... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002273
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002274>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00002275... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002276... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00002277... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002278
2279>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00002280... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002281... # Recipe credited to George Sakkis
2282... pending = len(iterables)
2283... nexts = cycle(iter(it).__next__ for it in iterables)
2284... while pending:
2285... try:
2286... for next in nexts:
2287... yield next()
2288... except StopIteration:
2289... pending -= 1
2290... nexts = cycle(islice(nexts, pending))
2291
2292>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00002293... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
2294... s = list(iterable)
2295... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002296
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00002297>>> def unique_everseen(iterable, key=None):
2298... "List unique elements, preserving order. Remember all elements ever seen."
2299... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
2300... # unique_everseen('ABBCcAD', str.lower) --> A B C D
2301... seen = set()
2302... seen_add = seen.add
2303... if key is None:
2304... for element in iterable:
2305... if element not in seen:
2306... seen_add(element)
2307... yield element
2308... else:
2309... for element in iterable:
2310... k = key(element)
2311... if k not in seen:
2312... seen_add(k)
2313... yield element
2314
2315>>> def unique_justseen(iterable, key=None):
2316... "List unique elements, preserving order. Remember only the element just seen."
2317... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
2318... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
2319... return map(next, map(itemgetter(1), groupby(iterable, key)))
2320
Raymond Hettinger31b26f62014-04-02 03:16:42 -07002321>>> def first_true(iterable, default=False, pred=None):
2322... '''Returns the first true value in the iterable.
2323...
2324... If no true value is found, returns *default*
2325...
2326... If *pred* is not None, returns the first item
2327... for which pred(item) is true.
2328...
2329... '''
2330... # first_true([a,b,c], x) --> a or b or c or x
2331... # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x
2332... return next(filter(pred, iterable), default)
2333
Raymond Hettingerd37258d2018-01-13 10:35:40 -08002334>>> def nth_combination(iterable, r, index):
2335... 'Equivalent to list(combinations(iterable, r))[index]'
2336... pool = tuple(iterable)
2337... n = len(pool)
2338... if r < 0 or r > n:
2339... raise ValueError
2340... c = 1
2341... k = min(r, n-r)
2342... for i in range(1, k+1):
2343... c = c * (n - k + i) // i
2344... if index < 0:
2345... index += c
2346... if index < 0 or index >= c:
2347... raise IndexError
2348... result = []
2349... while r:
2350... c, n, r = c*r//n, n-1, r-1
2351... while index >= c:
2352... index -= c
2353... c, n = c*(n-r)//n, n-1
2354... result.append(pool[-1-n])
2355... return tuple(result)
2356
2357
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002358This is not part of the examples but it tests to make sure the definitions
2359perform as purported.
2360
Raymond Hettingera098b332003-09-08 23:58:40 +00002361>>> take(10, count())
2362[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
2363
Raymond Hettinger9265dd72018-04-08 08:44:20 -07002364>>> list(prepend(1, [2, 3, 4]))
2365[1, 2, 3, 4]
2366
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002367>>> list(enumerate('abc'))
2368[(0, 'a'), (1, 'b'), (2, 'c')]
2369
2370>>> list(islice(tabulate(lambda x: 2*x), 4))
2371[0, 2, 4, 6]
2372
Cheryl Sabellada1734c2018-03-26 21:29:33 -04002373>>> it = iter(range(10))
2374>>> consume(it, 3)
2375>>> next(it)
23763
2377>>> consume(it)
2378>>> next(it, 'Done')
2379'Done'
2380
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002381>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00002382'd'
2383
2384>>> nth('abcde', 9) is None
2385True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002386
Raymond Hettingere525ee32016-03-06 18:11:38 -08002387>>> [all_equal(s) for s in ('', 'A', 'AAAA', 'AAAB', 'AAABA')]
2388[True, True, True, False, False]
2389
Guido van Rossum805365e2007-05-07 22:24:25 +00002390>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000239150
2392
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002393>>> a = [[1, 2, 3], [4, 5, 6]]
2394>>> flatten(a)
2395[1, 2, 3, 4, 5, 6]
2396
2397>>> list(repeatfunc(pow, 5, 2, 3))
2398[8, 8, 8, 8, 8]
2399
2400>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002401>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002402[0, 0, 0, 0, 0]
2403
Raymond Hettingerd591f662003-10-26 15:34:50 +00002404>>> list(pairwise('abcd'))
2405[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002406
Raymond Hettingerd591f662003-10-26 15:34:50 +00002407>>> list(pairwise([]))
2408[]
2409
2410>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00002411[]
2412
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002413>>> list(islice(padnone('abc'), 0, 6))
2414['a', 'b', 'c', None, None, None]
2415
2416>>> list(ncycles('abc', 3))
2417['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
2418
2419>>> dotproduct([1,2,3], [4,5,6])
242032
2421
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002422>>> list(grouper(3, 'abcdefg', 'x'))
2423[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
2424
2425>>> list(roundrobin('abc', 'd', 'ef'))
2426['a', 'd', 'e', 'b', 'f', 'c']
2427
Raymond Hettingerace67332009-01-26 02:23:50 +00002428>>> list(powerset([1,2,3]))
2429[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002430
Raymond Hettinger191e8502009-01-27 13:29:43 +00002431>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
2432True
2433
2434>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
2435True
2436
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00002437>>> list(unique_everseen('AAAABBBCCDAABBB'))
2438['A', 'B', 'C', 'D']
2439
2440>>> list(unique_everseen('ABBCcAD', str.lower))
2441['A', 'B', 'C', 'D']
2442
2443>>> list(unique_justseen('AAAABBBCCDAABBB'))
2444['A', 'B', 'C', 'D', 'A', 'B']
2445
2446>>> list(unique_justseen('ABBCcAD', str.lower))
2447['A', 'B', 'C', 'A', 'D']
2448
Raymond Hettinger31b26f62014-04-02 03:16:42 -07002449>>> first_true('ABC0DEF1', '9', str.isdigit)
2450'0'
2451
Raymond Hettingerd37258d2018-01-13 10:35:40 -08002452>>> population = 'ABCDEFGH'
2453>>> for r in range(len(population) + 1):
2454... seq = list(combinations(population, r))
2455... for i in range(len(seq)):
2456... assert nth_combination(population, r, i) == seq[i]
2457... for i in range(-len(seq), 0):
2458... assert nth_combination(population, r, i) == seq[i]
2459
2460
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002461"""
2462
2463__test__ = {'libreftest' : libreftest}
2464
2465def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002466 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00002467 RegressionTests, LengthTransparency,
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002468 SubclassWithKwargsTest, TestExamples,
Cheryl Sabellada1734c2018-03-26 21:29:33 -04002469 TestPurePythonRoughEquivalents,
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002470 SizeofTest)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002471 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002472
2473 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002474 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00002475 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00002476 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00002477 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002478 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00002479 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00002480 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00002481 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002482
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002483 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002484 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002485
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002486if __name__ == "__main__":
2487 test_main(verbose=True)