blob: df2997e87d494cf10084f30664732f7056f13472 [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
Serhiy Storchaka526a0142019-09-09 11:47:14 +030014import threading
Benjamin Petersonee8712c2008-05-20 21:35:26 +000015maxsize = support.MAX_Py_ssize_t
Guido van Rossum360e4b82007-05-14 22:51:27 +000016minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000017
Guido van Rossum801f0d72006-08-24 19:48:10 +000018def lzip(*args):
19 return list(zip(*args))
20
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000021def onearg(x):
22 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000023 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000024
25def errfunc(*args):
26 'Test function that raises an error'
27 raise ValueError
28
29def gen3():
30 'Non-restartable source sequence'
31 for i in (0, 1, 2):
32 yield i
33
34def isEven(x):
35 'Test predicate'
36 return x%2==0
37
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000038def isOdd(x):
39 'Test predicate'
40 return x%2==1
41
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000042def tupleize(*args):
43 return args
44
45def irange(n):
46 for i in range(n):
47 yield i
48
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000049class StopNow:
50 'Class emulating an empty iterable.'
51 def __iter__(self):
52 return self
Georg Brandla18af4e2007-04-21 15:47:16 +000053 def __next__(self):
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000054 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000055
Raymond Hettinger02420702003-06-29 20:36:23 +000056def take(n, seq):
57 'Convenience function for partially consuming a long of infinite iterable'
58 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000059
Christian Heimes78644762008-03-04 23:39:23 +000060def prod(iterable):
61 return reduce(operator.mul, iterable, 1)
62
Christian Heimes380f7f22008-02-28 11:19:05 +000063def fact(n):
64 'Factorial'
Christian Heimes78644762008-03-04 23:39:23 +000065 return prod(range(1, n+1))
66
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000067# root level methods for pickling ability
68def testR(r):
69 return r[0]
70
71def testR2(r):
72 return r[2]
73
74def underten(x):
75 return x<10
76
Serhiy Storchakabad12572014-12-15 14:03:42 +020077picklecopiers = [lambda s, proto=proto: pickle.loads(pickle.dumps(s, proto))
78 for proto in range(pickle.HIGHEST_PROTOCOL + 1)]
79
Raymond Hettinger96ef8112003-02-01 00:10:11 +000080class TestBasicOps(unittest.TestCase):
Raymond Hettinger482ba772010-12-01 22:48:00 +000081
Serhiy Storchakabad12572014-12-15 14:03:42 +020082 def pickletest(self, protocol, it, stop=4, take=1, compare=None):
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000083 """Test that an iterator is the same after pickling, also when part-consumed"""
84 def expand(it, i=0):
85 # Recursively expand iterables, within sensible bounds
86 if i > 10:
87 raise RuntimeError("infinite recursion encountered")
88 if isinstance(it, str):
89 return it
90 try:
91 l = list(islice(it, stop))
92 except TypeError:
93 return it # can't expand it
94 return [expand(e, i+1) for e in l]
95
96 # Test the initial copy against the original
Serhiy Storchakabad12572014-12-15 14:03:42 +020097 dump = pickle.dumps(it, protocol)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000098 i2 = pickle.loads(dump)
99 self.assertEqual(type(it), type(i2))
100 a, b = expand(it), expand(i2)
101 self.assertEqual(a, b)
102 if compare:
103 c = expand(compare)
104 self.assertEqual(a, c)
105
106 # Take from the copy, and create another copy and compare them.
107 i3 = pickle.loads(dump)
108 took = 0
109 try:
110 for i in range(take):
111 next(i3)
112 took += 1
113 except StopIteration:
114 pass #in case there is less data than 'take'
Serhiy Storchakabad12572014-12-15 14:03:42 +0200115 dump = pickle.dumps(i3, protocol)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000116 i4 = pickle.loads(dump)
117 a, b = expand(i3), expand(i4)
118 self.assertEqual(a, b)
119 if compare:
120 c = expand(compare[took:])
121 self.assertEqual(a, c);
122
Raymond Hettinger482ba772010-12-01 22:48:00 +0000123 def test_accumulate(self):
124 self.assertEqual(list(accumulate(range(10))), # one positional arg
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000125 [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])
126 self.assertEqual(list(accumulate(iterable=range(10))), # kw arg
127 [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])
Raymond Hettinger482ba772010-12-01 22:48:00 +0000128 for typ in int, complex, Decimal, Fraction: # multiple types
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000129 self.assertEqual(
130 list(accumulate(map(typ, range(10)))),
Raymond Hettinger482ba772010-12-01 22:48:00 +0000131 list(map(typ, [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])))
Raymond Hettinger2d93e6e2010-12-03 02:33:53 +0000132 self.assertEqual(list(accumulate('abc')), ['a', 'ab', 'abc']) # works with non-numeric
Raymond Hettinger482ba772010-12-01 22:48:00 +0000133 self.assertEqual(list(accumulate([])), []) # empty iterable
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000134 self.assertEqual(list(accumulate([7])), [7]) # iterable of length one
Raymond Hettinger5d446132011-03-27 18:52:10 -0700135 self.assertRaises(TypeError, accumulate, range(10), 5, 6) # too many args
Raymond Hettinger482ba772010-12-01 22:48:00 +0000136 self.assertRaises(TypeError, accumulate) # too few args
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000137 self.assertRaises(TypeError, accumulate, x=range(10)) # unexpected kwd arg
Raymond Hettinger482ba772010-12-01 22:48:00 +0000138 self.assertRaises(TypeError, list, accumulate([1, []])) # args that don't add
139
Raymond Hettinger5d446132011-03-27 18:52:10 -0700140 s = [2, 8, 9, 5, 7, 0, 3, 4, 1, 6]
141 self.assertEqual(list(accumulate(s, min)),
142 [2, 2, 2, 2, 2, 0, 0, 0, 0, 0])
143 self.assertEqual(list(accumulate(s, max)),
144 [2, 8, 9, 9, 9, 9, 9, 9, 9, 9])
145 self.assertEqual(list(accumulate(s, operator.mul)),
146 [2, 16, 144, 720, 5040, 0, 0, 0, 0, 0])
147 with self.assertRaises(TypeError):
148 list(accumulate(s, chr)) # unary-operation
Serhiy Storchakabad12572014-12-15 14:03:42 +0200149 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
150 self.pickletest(proto, accumulate(range(10))) # test pickling
Lisa Roach9718b592018-09-23 17:34:59 -0700151 self.pickletest(proto, accumulate(range(10), initial=7))
152 self.assertEqual(list(accumulate([10, 5, 1], initial=None)), [10, 15, 16])
153 self.assertEqual(list(accumulate([10, 5, 1], initial=100)), [100, 110, 115, 116])
154 self.assertEqual(list(accumulate([], initial=100)), [100])
155 with self.assertRaises(TypeError):
156 list(accumulate([10, 20], 100))
Raymond Hettinger5d446132011-03-27 18:52:10 -0700157
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000158 def test_chain(self):
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000159
160 def chain2(*iterables):
161 'Pure python version in the docs'
162 for it in iterables:
163 for element in it:
164 yield element
165
166 for c in (chain, chain2):
167 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
168 self.assertEqual(list(c('abc')), list('abc'))
169 self.assertEqual(list(c('')), [])
170 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
171 self.assertRaises(TypeError, list,c(2, 3))
Christian Heimesf16baeb2008-02-29 14:57:44 +0000172
173 def test_chain_from_iterable(self):
174 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
175 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
176 self.assertEqual(list(chain.from_iterable([''])), [])
177 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
178 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000179
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000180 def test_chain_reducible(self):
Serhiy Storchakabad12572014-12-15 14:03:42 +0200181 for oper in [copy.deepcopy] + picklecopiers:
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000182 it = chain('abc', 'def')
183 self.assertEqual(list(oper(it)), list('abcdef'))
184 self.assertEqual(next(it), 'a')
185 self.assertEqual(list(oper(it)), list('bcdef'))
186
187 self.assertEqual(list(oper(chain(''))), [])
188 self.assertEqual(take(4, oper(chain('abc', 'def'))), list('abcd'))
189 self.assertRaises(TypeError, list, oper(chain(2, 3)))
Serhiy Storchakabad12572014-12-15 14:03:42 +0200190 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
191 self.pickletest(proto, chain('abc', 'def'), compare=list('abcdef'))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000192
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300193 def test_chain_setstate(self):
194 self.assertRaises(TypeError, chain().__setstate__, ())
195 self.assertRaises(TypeError, chain().__setstate__, [])
196 self.assertRaises(TypeError, chain().__setstate__, 0)
197 self.assertRaises(TypeError, chain().__setstate__, ([],))
198 self.assertRaises(TypeError, chain().__setstate__, (iter([]), []))
199 it = chain()
200 it.__setstate__((iter(['abc', 'def']),))
201 self.assertEqual(list(it), ['a', 'b', 'c', 'd', 'e', 'f'])
202 it = chain()
203 it.__setstate__((iter(['abc', 'def']), iter(['ghi'])))
204 self.assertEqual(list(it), ['ghi', 'a', 'b', 'c', 'd', 'e', 'f'])
205
Christian Heimes380f7f22008-02-28 11:19:05 +0000206 def test_combinations(self):
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000207 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Christian Heimes380f7f22008-02-28 11:19:05 +0000208 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Christian Heimes78644762008-03-04 23:39:23 +0000209 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Christian Heimes380f7f22008-02-28 11:19:05 +0000210 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000211
Serhiy Storchakabad12572014-12-15 14:03:42 +0200212 for op in [lambda a:a] + picklecopiers:
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000213 self.assertEqual(list(op(combinations('abc', 32))), []) # r > n
214
215 self.assertEqual(list(op(combinations('ABCD', 2))),
216 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
217 testIntermediate = combinations('ABCD', 2)
218 next(testIntermediate)
219 self.assertEqual(list(op(testIntermediate)),
220 [('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
221
222 self.assertEqual(list(op(combinations(range(4), 3))),
223 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
224 testIntermediate = combinations(range(4), 3)
225 next(testIntermediate)
226 self.assertEqual(list(op(testIntermediate)),
227 [(0,1,3), (0,2,3), (1,2,3)])
228
Christian Heimes78644762008-03-04 23:39:23 +0000229
230 def combinations1(iterable, r):
231 'Pure python version shown in the docs'
232 pool = tuple(iterable)
233 n = len(pool)
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000234 if r > n:
235 return
Christian Heimes78644762008-03-04 23:39:23 +0000236 indices = list(range(r))
237 yield tuple(pool[i] for i in indices)
238 while 1:
239 for i in reversed(range(r)):
240 if indices[i] != i + n - r:
241 break
242 else:
243 return
244 indices[i] += 1
245 for j in range(i+1, r):
246 indices[j] = indices[j-1] + 1
247 yield tuple(pool[i] for i in indices)
248
249 def combinations2(iterable, r):
250 'Pure python version shown in the docs'
251 pool = tuple(iterable)
252 n = len(pool)
253 for indices in permutations(range(n), r):
254 if sorted(indices) == list(indices):
255 yield tuple(pool[i] for i in indices)
256
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000257 def combinations3(iterable, r):
258 'Pure python version from cwr()'
259 pool = tuple(iterable)
260 n = len(pool)
261 for indices in combinations_with_replacement(range(n), r):
262 if len(set(indices)) == r:
263 yield tuple(pool[i] for i in indices)
264
Christian Heimes78644762008-03-04 23:39:23 +0000265 for n in range(7):
Christian Heimes380f7f22008-02-28 11:19:05 +0000266 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000267 for r in range(n+2):
Christian Heimes380f7f22008-02-28 11:19:05 +0000268 result = list(combinations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000269 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 +0000270 self.assertEqual(len(result), len(set(result))) # no repeats
271 self.assertEqual(result, sorted(result)) # lexicographic order
272 for c in result:
273 self.assertEqual(len(c), r) # r-length combinations
274 self.assertEqual(len(set(c)), r) # no duplicate elements
275 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000276 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000277 self.assertEqual(list(c),
278 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000279 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000280 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000281 self.assertEqual(result, list(combinations3(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000282
Serhiy Storchakabad12572014-12-15 14:03:42 +0200283 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
284 self.pickletest(proto, combinations(values, r)) # test pickling
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000285
Benjamin Peterson4b40eeb2015-02-01 20:59:00 -0500286 @support.bigaddrspacetest
287 def test_combinations_overflow(self):
Serhiy Storchakadee948b2015-02-03 01:34:09 +0200288 with self.assertRaises((OverflowError, MemoryError)):
Benjamin Peterson4b40eeb2015-02-01 20:59:00 -0500289 combinations("AA", 2**29)
290
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000291 # Test implementation detail: tuple re-use
Alex Gaynore151d212011-07-17 16:21:30 -0700292 @support.impl_detail("tuple reuse is specific to CPython")
293 def test_combinations_tuple_reuse(self):
Christian Heimes78644762008-03-04 23:39:23 +0000294 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
295 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
296
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000297 def test_combinations_with_replacement(self):
298 cwr = combinations_with_replacement
299 self.assertRaises(TypeError, cwr, 'abc') # missing r argument
300 self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
301 self.assertRaises(TypeError, cwr, None) # pool is not iterable
302 self.assertRaises(ValueError, cwr, 'abc', -2) # r is negative
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000303
Serhiy Storchakabad12572014-12-15 14:03:42 +0200304 for op in [lambda a:a] + picklecopiers:
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000305 self.assertEqual(list(op(cwr('ABC', 2))),
306 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
307 testIntermediate = cwr('ABC', 2)
308 next(testIntermediate)
309 self.assertEqual(list(op(testIntermediate)),
310 [('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
311
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000312
313 def cwr1(iterable, r):
314 'Pure python version shown in the docs'
315 # number items returned: (n+r-1)! / r! / (n-1)! when n>0
316 pool = tuple(iterable)
317 n = len(pool)
318 if not n and r:
319 return
320 indices = [0] * r
321 yield tuple(pool[i] for i in indices)
322 while 1:
323 for i in reversed(range(r)):
324 if indices[i] != n - 1:
325 break
326 else:
327 return
328 indices[i:] = [indices[i] + 1] * (r - i)
329 yield tuple(pool[i] for i in indices)
330
331 def cwr2(iterable, r):
332 'Pure python version shown in the docs'
333 pool = tuple(iterable)
334 n = len(pool)
335 for indices in product(range(n), repeat=r):
336 if sorted(indices) == list(indices):
337 yield tuple(pool[i] for i in indices)
338
339 def numcombs(n, r):
340 if not n:
341 return 0 if r else 1
342 return fact(n+r-1) / fact(r)/ fact(n-1)
343
344 for n in range(7):
345 values = [5*x-12 for x in range(n)]
346 for r in range(n+2):
347 result = list(cwr(values, r))
348
349 self.assertEqual(len(result), numcombs(n, r)) # right number of combs
350 self.assertEqual(len(result), len(set(result))) # no repeats
351 self.assertEqual(result, sorted(result)) # lexicographic order
352
353 regular_combs = list(combinations(values, r)) # compare to combs without replacement
354 if n == 0 or r <= 1:
Ezio Melottib3aedd42010-11-20 19:04:17 +0000355 self.assertEqual(result, regular_combs) # cases that should be identical
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000356 else:
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000357 self.assertTrue(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000358
359 for c in result:
360 self.assertEqual(len(c), r) # r-length combinations
361 noruns = [k for k,v in groupby(c)] # combo without consecutive repeats
362 self.assertEqual(len(noruns), len(set(noruns))) # no repeats other than consecutive
363 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000364 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000365 self.assertEqual(noruns,
366 [e for e in values if e in c]) # comb is a subsequence of the input iterable
367 self.assertEqual(result, list(cwr1(values, r))) # matches first pure python version
368 self.assertEqual(result, list(cwr2(values, r))) # matches second pure python version
369
Serhiy Storchakabad12572014-12-15 14:03:42 +0200370 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
371 self.pickletest(proto, cwr(values,r)) # test pickling
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000372
Benjamin Peterson6f082292015-02-01 21:10:47 -0500373 @support.bigaddrspacetest
374 def test_combinations_with_replacement_overflow(self):
Serhiy Storchakadee948b2015-02-03 01:34:09 +0200375 with self.assertRaises((OverflowError, MemoryError)):
Benjamin Peterson6f082292015-02-01 21:10:47 -0500376 combinations_with_replacement("AA", 2**30)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000377
Benjamin Peterson6f082292015-02-01 21:10:47 -0500378 # Test implementation detail: tuple re-use
Alex Gaynore151d212011-07-17 16:21:30 -0700379 @support.impl_detail("tuple reuse is specific to CPython")
380 def test_combinations_with_replacement_tuple_reuse(self):
381 cwr = combinations_with_replacement
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000382 self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
383 self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
384
Christian Heimes78644762008-03-04 23:39:23 +0000385 def test_permutations(self):
386 self.assertRaises(TypeError, permutations) # too few arguments
387 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000388 self.assertRaises(TypeError, permutations, None) # pool is not iterable
389 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000390 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000391 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Christian Heimes78644762008-03-04 23:39:23 +0000392 self.assertEqual(list(permutations(range(3), 2)),
393 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
394
395 def permutations1(iterable, r=None):
396 'Pure python version shown in the docs'
397 pool = tuple(iterable)
398 n = len(pool)
399 r = n if r is None else r
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000400 if r > n:
401 return
Christian Heimes78644762008-03-04 23:39:23 +0000402 indices = list(range(n))
403 cycles = list(range(n-r+1, n+1))[::-1]
404 yield tuple(pool[i] for i in indices[:r])
405 while n:
406 for i in reversed(range(r)):
407 cycles[i] -= 1
408 if cycles[i] == 0:
409 indices[i:] = indices[i+1:] + indices[i:i+1]
410 cycles[i] = n - i
411 else:
412 j = cycles[i]
413 indices[i], indices[-j] = indices[-j], indices[i]
414 yield tuple(pool[i] for i in indices[:r])
415 break
416 else:
417 return
418
419 def permutations2(iterable, r=None):
420 'Pure python version shown in the docs'
421 pool = tuple(iterable)
422 n = len(pool)
423 r = n if r is None else r
424 for indices in product(range(n), repeat=r):
425 if len(set(indices)) == r:
426 yield tuple(pool[i] for i in indices)
427
428 for n in range(7):
429 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000430 for r in range(n+2):
Christian Heimes78644762008-03-04 23:39:23 +0000431 result = list(permutations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000432 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 +0000433 self.assertEqual(len(result), len(set(result))) # no repeats
434 self.assertEqual(result, sorted(result)) # lexicographic order
435 for p in result:
436 self.assertEqual(len(p), r) # r-length permutations
437 self.assertEqual(len(set(p)), r) # no duplicate elements
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000438 self.assertTrue(all(e in values for e in p)) # elements taken from input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000439 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000440 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000441 if r == n:
442 self.assertEqual(result, list(permutations(values, None))) # test r as None
443 self.assertEqual(result, list(permutations(values))) # test default r
444
Serhiy Storchakabad12572014-12-15 14:03:42 +0200445 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
446 self.pickletest(proto, permutations(values, r)) # test pickling
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000447
Benjamin Peterson0eaabf12015-02-01 21:34:07 -0500448 @support.bigaddrspacetest
449 def test_permutations_overflow(self):
Serhiy Storchakadee948b2015-02-03 01:34:09 +0200450 with self.assertRaises((OverflowError, MemoryError)):
Benjamin Peterson0eaabf12015-02-01 21:34:07 -0500451 permutations("A", 2**30)
Benjamin Peterson0eaabf12015-02-01 21:34:07 -0500452
Zachary Waredca807b2014-04-24 13:22:05 -0500453 @support.impl_detail("tuple reuse is specific to CPython")
Alex Gaynore151d212011-07-17 16:21:30 -0700454 def test_permutations_tuple_reuse(self):
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000455 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Christian Heimes78644762008-03-04 23:39:23 +0000456 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Christian Heimes380f7f22008-02-28 11:19:05 +0000457
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000458 def test_combinatorics(self):
459 # Test relationships between product(), permutations(),
460 # combinations() and combinations_with_replacement().
461
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000462 for n in range(6):
463 s = 'ABCDEFG'[:n]
464 for r in range(8):
465 prod = list(product(s, repeat=r))
466 cwr = list(combinations_with_replacement(s, r))
467 perm = list(permutations(s, r))
468 comb = list(combinations(s, r))
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000469
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000470 # Check size
Ezio Melottib3aedd42010-11-20 19:04:17 +0000471 self.assertEqual(len(prod), n**r)
472 self.assertEqual(len(cwr), (fact(n+r-1) / fact(r)/ fact(n-1)) if n else (not r))
473 self.assertEqual(len(perm), 0 if r>n else fact(n) / fact(n-r))
474 self.assertEqual(len(comb), 0 if r>n else fact(n) / fact(r) / fact(n-r))
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000475
476 # Check lexicographic order without repeated tuples
Ezio Melottib3aedd42010-11-20 19:04:17 +0000477 self.assertEqual(prod, sorted(set(prod)))
478 self.assertEqual(cwr, sorted(set(cwr)))
479 self.assertEqual(perm, sorted(set(perm)))
480 self.assertEqual(comb, sorted(set(comb)))
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000481
482 # Check interrelationships
Ezio Melottib3aedd42010-11-20 19:04:17 +0000483 self.assertEqual(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
484 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 +0000485 self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
486 self.assertEqual(comb, [t for t in cwr if len(set(t))==r]) # comb: cwrs without dups
487 self.assertEqual(comb, list(filter(set(cwr).__contains__, perm))) # comb: perm that is a cwr
488 self.assertEqual(comb, list(filter(set(perm).__contains__, cwr))) # comb: cwr that is a perm
489 self.assertEqual(comb, sorted(set(cwr) & set(perm))) # comb: both a cwr and a perm
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000490
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000491 def test_compress(self):
Raymond Hettinger15a49502009-02-19 02:17:09 +0000492 self.assertEqual(list(compress(data='ABCDEF', selectors=[1,0,1,0,1,1])), list('ACEF'))
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000493 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
494 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
495 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
496 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
497 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
498 n = 10000
499 data = chain.from_iterable(repeat(range(6), n))
500 selectors = chain.from_iterable(repeat((0, 1)))
501 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
502 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
503 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
504 self.assertRaises(TypeError, compress, range(6)) # too few args
505 self.assertRaises(TypeError, compress, range(6), None) # too many args
506
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000507 # check copy, deepcopy, pickle
Serhiy Storchakabad12572014-12-15 14:03:42 +0200508 for op in [lambda a:copy.copy(a), lambda a:copy.deepcopy(a)] + picklecopiers:
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000509 for data, selectors, result1, result2 in [
510 ('ABCDEF', [1,0,1,0,1,1], 'ACEF', 'CEF'),
511 ('ABCDEF', [0,0,0,0,0,0], '', ''),
512 ('ABCDEF', [1,1,1,1,1,1], 'ABCDEF', 'BCDEF'),
513 ('ABCDEF', [1,0,1], 'AC', 'C'),
514 ('ABC', [0,1,1,1,1,1], 'BC', 'C'),
515 ]:
516
517 self.assertEqual(list(op(compress(data=data, selectors=selectors))), list(result1))
518 self.assertEqual(list(op(compress(data, selectors))), list(result1))
519 testIntermediate = compress(data, selectors)
520 if result1:
521 next(testIntermediate)
522 self.assertEqual(list(op(testIntermediate)), list(result2))
523
524
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000525 def test_count(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000526 self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
527 self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
528 self.assertEqual(take(2, lzip('abc',count(3))), [('a', 3), ('b', 4)])
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000529 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
530 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000531 self.assertRaises(TypeError, count, 2, 3, 4)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000532 self.assertRaises(TypeError, count, 'a')
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300533 self.assertEqual(take(10, count(maxsize-5)),
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000534 list(range(maxsize-5, maxsize+5)))
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300535 self.assertEqual(take(10, count(-maxsize-5)),
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000536 list(range(-maxsize-5, -maxsize+5)))
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300537 self.assertEqual(take(3, count(3.25)), [3.25, 4.25, 5.25])
538 self.assertEqual(take(3, count(3.25-4j)), [3.25-4j, 4.25-4j, 5.25-4j])
539 self.assertEqual(take(3, count(Decimal('1.1'))),
540 [Decimal('1.1'), Decimal('2.1'), Decimal('3.1')])
541 self.assertEqual(take(3, count(Fraction(2, 3))),
542 [Fraction(2, 3), Fraction(5, 3), Fraction(8, 3)])
543 BIGINT = 1<<1000
544 self.assertEqual(take(3, count(BIGINT)), [BIGINT, BIGINT+1, BIGINT+2])
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000545 c = count(3)
546 self.assertEqual(repr(c), 'count(3)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000547 next(c)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000548 self.assertEqual(repr(c), 'count(4)')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000549 c = count(-9)
550 self.assertEqual(repr(c), 'count(-9)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000551 next(c)
552 self.assertEqual(next(c), -8)
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300553 self.assertEqual(repr(count(10.25)), 'count(10.25)')
554 self.assertEqual(repr(count(10.0)), 'count(10.0)')
555 self.assertEqual(type(next(count(10.0))), float)
Christian Heimesa37d4c62007-12-04 23:02:19 +0000556 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 +0300557 # Test repr
558 r1 = repr(count(i))
559 r2 = 'count(%r)'.__mod__(i)
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000560 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000561
Raymond Hettingera3e1ad22009-11-30 22:02:31 +0000562 # check copy, deepcopy, pickle
563 for value in -3, 3, maxsize-5, maxsize+5:
564 c = count(value)
565 self.assertEqual(next(copy.copy(c)), value)
566 self.assertEqual(next(copy.deepcopy(c)), value)
Serhiy Storchakabad12572014-12-15 14:03:42 +0200567 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
568 self.pickletest(proto, count(value))
Raymond Hettingera3e1ad22009-11-30 22:02:31 +0000569
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +0000570 #check proper internal error handling for large "step' sizes
571 count(1, maxsize+5); sys.exc_info()
572
Raymond Hettinger30729212009-02-12 06:28:27 +0000573 def test_count_with_stride(self):
574 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingereb13fdd2009-02-21 22:30:12 +0000575 self.assertEqual(lzip('abc',count(start=2,step=3)),
576 [('a', 2), ('b', 5), ('c', 8)])
577 self.assertEqual(lzip('abc',count(step=-1)),
578 [('a', 0), ('b', -1), ('c', -2)])
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300579 self.assertRaises(TypeError, count, 'a', 'b')
Raymond Hettinger30729212009-02-12 06:28:27 +0000580 self.assertEqual(lzip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
581 self.assertEqual(lzip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000582 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000583 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
584 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300585 self.assertEqual(take(3, count(10, maxsize+5)),
586 list(range(10, 10+3*(maxsize+5), maxsize+5)))
587 self.assertEqual(take(3, count(2, 1.25)), [2, 3.25, 4.5])
Raymond Hettinger30729212009-02-12 06:28:27 +0000588 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettinger8a882a72009-02-12 12:08:18 +0000589 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
590 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerde09acf2009-02-12 12:53:53 +0000591 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
592 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300593 BIGINT = 1<<1000
594 self.assertEqual(take(3, count(step=BIGINT)), [0, BIGINT, 2*BIGINT])
Raymond Hettinger30729212009-02-12 06:28:27 +0000595 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
596 c = count(3, 5)
597 self.assertEqual(repr(c), 'count(3, 5)')
598 next(c)
599 self.assertEqual(repr(c), 'count(8, 5)')
600 c = count(-9, 0)
601 self.assertEqual(repr(c), 'count(-9, 0)')
602 next(c)
603 self.assertEqual(repr(c), 'count(-9, 0)')
604 c = count(-9, -3)
605 self.assertEqual(repr(c), 'count(-9, -3)')
606 next(c)
607 self.assertEqual(repr(c), 'count(-12, -3)')
608 self.assertEqual(repr(c), 'count(-12, -3)')
609 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
610 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
611 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 +0300612 self.assertEqual(repr(count(10, 1.00)), 'count(10, 1.0)')
613 c = count(10, 1.0)
614 self.assertEqual(type(next(c)), int)
615 self.assertEqual(type(next(c)), float)
Raymond Hettinger30729212009-02-12 06:28:27 +0000616 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
617 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 +0300618 # Test repr
619 r1 = repr(count(i, j))
Raymond Hettinger30729212009-02-12 06:28:27 +0000620 if j == 1:
Serhiy Storchaka95949422013-08-27 19:40:23 +0300621 r2 = ('count(%r)' % i)
Raymond Hettinger30729212009-02-12 06:28:27 +0000622 else:
Serhiy Storchaka95949422013-08-27 19:40:23 +0300623 r2 = ('count(%r, %r)' % (i, j))
Raymond Hettinger30729212009-02-12 06:28:27 +0000624 self.assertEqual(r1, r2)
Serhiy Storchakabad12572014-12-15 14:03:42 +0200625 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
626 self.pickletest(proto, count(i, j))
Raymond Hettinger30729212009-02-12 06:28:27 +0000627
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000628 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000629 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000630 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000631 self.assertRaises(TypeError, cycle)
632 self.assertRaises(TypeError, cycle, 5)
633 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000634
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000635 # check copy, deepcopy, pickle
636 c = cycle('abc')
637 self.assertEqual(next(c), 'a')
638 #simple copy currently not supported, because __reduce__ returns
639 #an internal iterator
640 #self.assertEqual(take(10, copy.copy(c)), list('bcabcabcab'))
641 self.assertEqual(take(10, copy.deepcopy(c)), list('bcabcabcab'))
Serhiy Storchakabad12572014-12-15 14:03:42 +0200642 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
643 self.assertEqual(take(10, pickle.loads(pickle.dumps(c, proto))),
644 list('bcabcabcab'))
645 next(c)
646 self.assertEqual(take(10, pickle.loads(pickle.dumps(c, proto))),
647 list('cabcabcabc'))
648 next(c)
649 next(c)
650 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
651 self.pickletest(proto, cycle('abc'))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000652
Raymond Hettingera166ce52015-08-15 14:45:49 -0700653 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
654 # test with partial consumed input iterable
655 it = iter('abcde')
656 c = cycle(it)
Raymond Hettinger1cadf762015-08-15 14:47:27 -0700657 _ = [next(c) for i in range(2)] # consume 2 of 5 inputs
Raymond Hettingera166ce52015-08-15 14:45:49 -0700658 p = pickle.dumps(c, proto)
659 d = pickle.loads(p) # rebuild the cycle object
660 self.assertEqual(take(20, d), list('cdeabcdeabcdeabcdeab'))
661
662 # test with completely consumed input iterable
663 it = iter('abcde')
664 c = cycle(it)
Raymond Hettinger1cadf762015-08-15 14:47:27 -0700665 _ = [next(c) for i in range(7)] # consume 7 of 5 inputs
Raymond Hettingera166ce52015-08-15 14:45:49 -0700666 p = pickle.dumps(c, proto)
667 d = pickle.loads(p) # rebuild the cycle object
668 self.assertEqual(take(20, d), list('cdeabcdeabcdeabcdeab'))
669
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700670 def test_cycle_setstate(self):
671 # Verify both modes for restoring state
672
673 # Mode 0 is efficient. It uses an incompletely consumed input
674 # iterator to build a cycle object and then passes in state with
675 # a list of previously consumed values. There is no data
Martin Panterf1579822016-05-26 06:03:33 +0000676 # overlap between the two.
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700677 c = cycle('defg')
678 c.__setstate__((list('abc'), 0))
679 self.assertEqual(take(20, c), list('defgabcdefgabcdefgab'))
680
681 # Mode 1 is inefficient. It starts with a cycle object built
682 # from an iterator over the remaining elements in a partial
683 # cycle and then passes in state with all of the previously
684 # seen values (this overlaps values included in the iterator).
685 c = cycle('defg')
686 c.__setstate__((list('abcdefg'), 1))
687 self.assertEqual(take(20, c), list('defgabcdefgabcdefgab'))
688
689 # The first argument to setstate needs to be a tuple
Serhiy Storchaka8f0f2052016-10-02 09:13:14 +0300690 with self.assertRaises(TypeError):
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700691 cycle('defg').__setstate__([list('abcdefg'), 0])
692
693 # The first argument in the setstate tuple must be a list
694 with self.assertRaises(TypeError):
695 c = cycle('defg')
Serhiy Storchaka8f0f2052016-10-02 09:13:14 +0300696 c.__setstate__((tuple('defg'), 0))
697 take(20, c)
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700698
Serhiy Storchaka8f0f2052016-10-02 09:13:14 +0300699 # The second argument in the setstate tuple must be an int
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700700 with self.assertRaises(TypeError):
701 cycle('defg').__setstate__((list('abcdefg'), 'x'))
702
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300703 self.assertRaises(TypeError, cycle('').__setstate__, ())
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300704 self.assertRaises(TypeError, cycle('').__setstate__, ([],))
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300705
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000706 def test_groupby(self):
707 # Check whether it accepts arguments correctly
708 self.assertEqual([], list(groupby([])))
709 self.assertEqual([], list(groupby([], key=id)))
710 self.assertRaises(TypeError, list, groupby('abc', []))
711 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000712 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000713
714 # Check normal input
715 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
716 (2,15,22), (3,16,23), (3,17,23)]
717 dup = []
718 for k, g in groupby(s, lambda r:r[0]):
719 for elem in g:
720 self.assertEqual(k, elem[0])
721 dup.append(elem)
722 self.assertEqual(s, dup)
723
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000724 # Check normal pickled
Serhiy Storchakabad12572014-12-15 14:03:42 +0200725 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
726 dup = []
727 for k, g in pickle.loads(pickle.dumps(groupby(s, testR), proto)):
728 for elem in g:
729 self.assertEqual(k, elem[0])
730 dup.append(elem)
731 self.assertEqual(s, dup)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000732
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000733 # Check nested case
734 dup = []
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000735 for k, g in groupby(s, testR):
736 for ik, ig in groupby(g, testR2):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000737 for elem in ig:
738 self.assertEqual(k, elem[0])
739 self.assertEqual(ik, elem[2])
740 dup.append(elem)
741 self.assertEqual(s, dup)
742
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000743 # Check nested and pickled
Serhiy Storchakabad12572014-12-15 14:03:42 +0200744 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
745 dup = []
746 for k, g in pickle.loads(pickle.dumps(groupby(s, testR), proto)):
747 for ik, ig in pickle.loads(pickle.dumps(groupby(g, testR2), proto)):
748 for elem in ig:
749 self.assertEqual(k, elem[0])
750 self.assertEqual(ik, elem[2])
751 dup.append(elem)
752 self.assertEqual(s, dup)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000753
754
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000755 # Check case where inner iterator is not used
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000756 keys = [k for k, g in groupby(s, testR)]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000757 expectedkeys = set([r[0] for r in s])
758 self.assertEqual(set(keys), expectedkeys)
759 self.assertEqual(len(keys), len(expectedkeys))
760
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300761 # Check case where inner iterator is used after advancing the groupby
762 # iterator
763 s = list(zip('AABBBAAAA', range(9)))
764 it = groupby(s, testR)
765 _, g1 = next(it)
766 _, g2 = next(it)
767 _, g3 = next(it)
768 self.assertEqual(list(g1), [])
769 self.assertEqual(list(g2), [])
770 self.assertEqual(next(g3), ('A', 5))
771 list(it) # exhaust the groupby iterator
772 self.assertEqual(list(g3), [])
773
774 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
775 it = groupby(s, testR)
776 _, g = next(it)
777 next(it)
778 next(it)
779 self.assertEqual(list(pickle.loads(pickle.dumps(g, proto))), [])
780
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000781 # Exercise pipes and filters style
782 s = 'abracadabra'
783 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000784 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000785 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
786 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000787 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000788 self.assertEqual(r, ['a', 'b', 'r'])
789 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000790 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000791 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
792 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000793 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000794 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
795
Georg Brandla18af4e2007-04-21 15:47:16 +0000796 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000797 class ExpectedError(Exception):
798 pass
799 def delayed_raise(n=0):
800 for i in range(n):
801 yield 'yo'
802 raise ExpectedError
803 def gulp(iterable, keyp=None, func=list):
804 return [func(g) for k, g in groupby(iterable, keyp)]
805
Georg Brandla18af4e2007-04-21 15:47:16 +0000806 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000807 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000808 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000809 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
810
Serhiy Storchakaa60c2fe2015-03-12 21:56:08 +0200811 # __eq__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000812 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000813 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000814 raise ExpectedError
815 s = [DummyCmp(), DummyCmp(), None]
816
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000817 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000818 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000819 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000820 self.assertRaises(ExpectedError, gulp, s)
821
822 # keyfunc failure
823 def keyfunc(obj):
824 if keyfunc.skip > 0:
825 keyfunc.skip -= 1
826 return obj
827 else:
828 raise ExpectedError
829
830 # keyfunc failure on outer object
831 keyfunc.skip = 0
832 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
833 keyfunc.skip = 1
834 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
835
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000836 def test_filter(self):
837 self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
838 self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])
839 self.assertEqual(list(filter(bool, [0,1,0,2,0])), [1,2])
840 self.assertEqual(take(4, filter(isEven, count())), [0,2,4,6])
841 self.assertRaises(TypeError, filter)
842 self.assertRaises(TypeError, filter, lambda x:x)
843 self.assertRaises(TypeError, filter, lambda x:x, range(6), 7)
844 self.assertRaises(TypeError, filter, isEven, 3)
845 self.assertRaises(TypeError, next, filter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000846
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000847 # check copy, deepcopy, pickle
848 ans = [0,2,4]
849
850 c = filter(isEven, range(6))
851 self.assertEqual(list(copy.copy(c)), ans)
852 c = filter(isEven, range(6))
853 self.assertEqual(list(copy.deepcopy(c)), ans)
Serhiy Storchakabad12572014-12-15 14:03:42 +0200854 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
855 c = filter(isEven, range(6))
856 self.assertEqual(list(pickle.loads(pickle.dumps(c, proto))), ans)
857 next(c)
858 self.assertEqual(list(pickle.loads(pickle.dumps(c, proto))), ans[1:])
859 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
860 c = filter(isEven, range(6))
861 self.pickletest(proto, c)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000862
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000863 def test_filterfalse(self):
864 self.assertEqual(list(filterfalse(isEven, range(6))), [1,3,5])
865 self.assertEqual(list(filterfalse(None, [0,1,0,2,0])), [0,0,0])
866 self.assertEqual(list(filterfalse(bool, [0,1,0,2,0])), [0,0,0])
867 self.assertEqual(take(4, filterfalse(isEven, count())), [1,3,5,7])
868 self.assertRaises(TypeError, filterfalse)
869 self.assertRaises(TypeError, filterfalse, lambda x:x)
870 self.assertRaises(TypeError, filterfalse, lambda x:x, range(6), 7)
871 self.assertRaises(TypeError, filterfalse, isEven, 3)
872 self.assertRaises(TypeError, next, filterfalse(range(6), range(6)))
Serhiy Storchakabad12572014-12-15 14:03:42 +0200873 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
874 self.pickletest(proto, filterfalse(isEven, range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000875
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000876 def test_zip(self):
877 # XXX This is rather silly now that builtin zip() calls zip()...
878 ans = [(x,y) for x, y in zip('abc',count())]
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000879 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000880 self.assertEqual(list(zip('abc', range(6))), lzip('abc', range(6)))
881 self.assertEqual(list(zip('abcdef', range(3))), lzip('abcdef', range(3)))
882 self.assertEqual(take(3,zip('abcdef', count())), lzip('abcdef', range(3)))
883 self.assertEqual(list(zip('abcdef')), lzip('abcdef'))
884 self.assertEqual(list(zip()), lzip())
885 self.assertRaises(TypeError, zip, 3)
886 self.assertRaises(TypeError, zip, range(3), 3)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000887 self.assertEqual([tuple(list(pair)) for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000888 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000889 self.assertEqual([pair for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000890 lzip('abc', 'def'))
Alex Gaynore151d212011-07-17 16:21:30 -0700891
892 @support.impl_detail("tuple reuse is specific to CPython")
893 def test_zip_tuple_reuse(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000894 ids = list(map(id, zip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000895 self.assertEqual(min(ids), max(ids))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000896 ids = list(map(id, list(zip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000897 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000898
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000899 # check copy, deepcopy, pickle
900 ans = [(x,y) for x, y in copy.copy(zip('abc',count()))]
901 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
902
903 ans = [(x,y) for x, y in copy.deepcopy(zip('abc',count()))]
904 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
905
Serhiy Storchakabad12572014-12-15 14:03:42 +0200906 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
907 ans = [(x,y) for x, y in pickle.loads(pickle.dumps(zip('abc',count()), proto))]
908 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000909
Serhiy Storchakabad12572014-12-15 14:03:42 +0200910 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
911 testIntermediate = zip('abc',count())
912 next(testIntermediate)
913 ans = [(x,y) for x, y in pickle.loads(pickle.dumps(testIntermediate, proto))]
914 self.assertEqual(ans, [('b', 1), ('c', 2)])
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000915
Serhiy Storchakabad12572014-12-15 14:03:42 +0200916 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
917 self.pickletest(proto, zip('abc', count()))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000918
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000919 def test_ziplongest(self):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000920 for args in [
921 ['abc', range(6)],
922 [range(6), 'abc'],
923 [range(1000), range(2000,2100), range(3000,3050)],
924 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
925 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
926 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000927 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
928 for i in range(max(map(len, args)))]
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000929 self.assertEqual(list(zip_longest(*args)), target)
930 self.assertEqual(list(zip_longest(*args, **{})), target)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000931 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 +0000932 self.assertEqual(list(zip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000933
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000934 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 +0000935
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000936 self.assertEqual(list(zip_longest()), list(zip()))
937 self.assertEqual(list(zip_longest([])), list(zip([])))
938 self.assertEqual(list(zip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000939
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000940 self.assertEqual(list(zip_longest('abc', 'defg', **{})),
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000941 list(zip(list('abc')+[None], 'defg'))) # empty keyword dict
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000942 self.assertRaises(TypeError, zip_longest, 3)
943 self.assertRaises(TypeError, zip_longest, range(3), 3)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000944
945 for stmt in [
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000946 "zip_longest('abc', fv=1)",
947 "zip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000948 ]:
949 try:
950 eval(stmt, globals(), locals())
951 except TypeError:
952 pass
953 else:
954 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000955
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000956 self.assertEqual([tuple(list(pair)) for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000957 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000958 self.assertEqual([pair for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000959 list(zip('abc', 'def')))
Alex Gaynore151d212011-07-17 16:21:30 -0700960
961 @support.impl_detail("tuple reuse is specific to CPython")
962 def test_zip_longest_tuple_reuse(self):
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000963 ids = list(map(id, zip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000964 self.assertEqual(min(ids), max(ids))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000965 ids = list(map(id, list(zip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000966 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
967
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000968 def test_zip_longest_pickling(self):
Serhiy Storchakabad12572014-12-15 14:03:42 +0200969 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
970 self.pickletest(proto, zip_longest("abc", "def"))
971 self.pickletest(proto, zip_longest("abc", "defgh"))
972 self.pickletest(proto, zip_longest("abc", "defgh", fillvalue=1))
973 self.pickletest(proto, zip_longest("", "defgh"))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000974
Sergey Fedoseev6a650aa2019-08-30 09:25:48 +0500975 def test_zip_longest_bad_iterable(self):
976 exception = TypeError()
977
978 class BadIterable:
979 def __iter__(self):
980 raise exception
981
982 with self.assertRaises(TypeError) as cm:
983 zip_longest(BadIterable())
984
985 self.assertIs(cm.exception, exception)
986
Raymond Hettingerfc438512009-11-01 20:55:33 +0000987 def test_bug_7244(self):
988
989 class Repeater:
990 # this class is similar to itertools.repeat
991 def __init__(self, o, t, e):
992 self.o = o
993 self.t = int(t)
994 self.e = e
995 def __iter__(self): # its iterator is itself
996 return self
997 def __next__(self):
998 if self.t > 0:
999 self.t -= 1
1000 return self.o
1001 else:
1002 raise self.e
1003
1004 # Formerly this code in would fail in debug mode
1005 # with Undetected Error and Stop Iteration
1006 r1 = Repeater(1, 3, StopIteration)
1007 r2 = Repeater(2, 4, StopIteration)
1008 def run(r1, r2):
1009 result = []
1010 for i, j in zip_longest(r1, r2, fillvalue=0):
1011 with support.captured_output('stdout'):
1012 print((i, j))
1013 result.append((i, j))
1014 return result
1015 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
1016
1017 # Formerly, the RuntimeError would be lost
1018 # and StopIteration would stop as expected
1019 r1 = Repeater(1, 3, RuntimeError)
1020 r2 = Repeater(2, 4, StopIteration)
1021 it = zip_longest(r1, r2, fillvalue=0)
1022 self.assertEqual(next(it), (1, 2))
1023 self.assertEqual(next(it), (1, 2))
1024 self.assertEqual(next(it), (1, 2))
1025 self.assertRaises(RuntimeError, next, it)
1026
Raymond Hettingercc061d02020-11-30 20:42:54 -08001027 def test_pairwise(self):
1028 self.assertEqual(list(pairwise('')), [])
1029 self.assertEqual(list(pairwise('a')), [])
1030 self.assertEqual(list(pairwise('ab')),
1031 [('a', 'b')]),
1032 self.assertEqual(list(pairwise('abcde')),
1033 [('a', 'b'), ('b', 'c'), ('c', 'd'), ('d', 'e')])
1034 self.assertEqual(list(pairwise(range(10_000))),
1035 list(zip(range(10_000), range(1, 10_000))))
1036
1037 with self.assertRaises(TypeError):
1038 pairwise() # too few arguments
1039 with self.assertRaises(TypeError):
1040 pairwise('abc', 10) # too many arguments
1041 with self.assertRaises(TypeError):
1042 pairwise(iterable='abc') # keyword arguments
1043 with self.assertRaises(TypeError):
1044 pairwise(None) # non-iterable argument
1045
Christian Heimesc3f30c42008-02-22 16:37:40 +00001046 def test_product(self):
1047 for args, result in [
Christian Heimes78644762008-03-04 23:39:23 +00001048 ([], [()]), # zero iterables
Christian Heimesc3f30c42008-02-22 16:37:40 +00001049 (['ab'], [('a',), ('b',)]), # one iterable
1050 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
1051 ([range(0), range(2), range(3)], []), # first iterable with zero length
1052 ([range(2), range(0), range(3)], []), # middle iterable with zero length
1053 ([range(2), range(3), range(0)], []), # last iterable with zero length
1054 ]:
1055 self.assertEqual(list(product(*args)), result)
Christian Heimesf16baeb2008-02-29 14:57:44 +00001056 for r in range(4):
1057 self.assertEqual(list(product(*(args*r))),
1058 list(product(*args, **dict(repeat=r))))
Christian Heimesc3f30c42008-02-22 16:37:40 +00001059 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
1060 self.assertRaises(TypeError, product, range(6), None)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001061
1062 def product1(*args, **kwds):
1063 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
1064 n = len(pools)
1065 if n == 0:
1066 yield ()
1067 return
1068 if any(len(pool) == 0 for pool in pools):
1069 return
1070 indices = [0] * n
1071 yield tuple(pool[i] for pool, i in zip(pools, indices))
1072 while 1:
1073 for i in reversed(range(n)): # right to left
1074 if indices[i] == len(pools[i]) - 1:
1075 continue
1076 indices[i] += 1
1077 for j in range(i+1, n):
1078 indices[j] = 0
1079 yield tuple(pool[i] for pool, i in zip(pools, indices))
1080 break
1081 else:
1082 return
1083
1084 def product2(*args, **kwds):
1085 'Pure python version used in docs'
1086 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
1087 result = [[]]
1088 for pool in pools:
1089 result = [x+[y] for x in result for y in pool]
1090 for prod in result:
1091 yield tuple(prod)
1092
Christian Heimesc3f30c42008-02-22 16:37:40 +00001093 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
1094 set('abcdefg'), range(11), tuple(range(13))]
1095 for i in range(100):
1096 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Christian Heimes78644762008-03-04 23:39:23 +00001097 expected_len = prod(map(len, args))
1098 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001099 self.assertEqual(list(product(*args)), list(product1(*args)))
1100 self.assertEqual(list(product(*args)), list(product2(*args)))
Christian Heimesc3f30c42008-02-22 16:37:40 +00001101 args = map(iter, args)
Christian Heimes78644762008-03-04 23:39:23 +00001102 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesc3f30c42008-02-22 16:37:40 +00001103
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05001104 @support.bigaddrspacetest
1105 def test_product_overflow(self):
Serhiy Storchakadee948b2015-02-03 01:34:09 +02001106 with self.assertRaises((OverflowError, MemoryError)):
1107 product(*(['ab']*2**5), repeat=2**25)
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05001108
Alex Gaynore151d212011-07-17 16:21:30 -07001109 @support.impl_detail("tuple reuse is specific to CPython")
1110 def test_product_tuple_reuse(self):
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001111 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
1112 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +00001113
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001114 def test_product_pickling(self):
1115 # check copy, deepcopy, pickle
1116 for args, result in [
1117 ([], [()]), # zero iterables
1118 (['ab'], [('a',), ('b',)]), # one iterable
1119 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
1120 ([range(0), range(2), range(3)], []), # first iterable with zero length
1121 ([range(2), range(0), range(3)], []), # middle iterable with zero length
1122 ([range(2), range(3), range(0)], []), # last iterable with zero length
1123 ]:
1124 self.assertEqual(list(copy.copy(product(*args))), result)
1125 self.assertEqual(list(copy.deepcopy(product(*args))), result)
Serhiy Storchakabad12572014-12-15 14:03:42 +02001126 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1127 self.pickletest(proto, product(*args))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001128
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00001129 def test_product_issue_25021(self):
1130 # test that indices are properly clamped to the length of the tuples
1131 p = product((1, 2),(3,))
1132 p.__setstate__((0, 0x1000)) # will access tuple element 1 if not clamped
1133 self.assertEqual(next(p), (2, 3))
1134 # test that empty tuple in the list will result in an immediate StopIteration
1135 p = product((1, 2), (), (3,))
1136 p.__setstate__((0, 0, 0x1000)) # will access tuple element 1 if not clamped
1137 self.assertRaises(StopIteration, next, p)
1138
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001139 def test_repeat(self):
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00001140 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
Guido van Rossum805365e2007-05-07 22:24:25 +00001141 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001142 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001143 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +00001144 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001145 self.assertEqual(list(repeat('a', 0)), [])
1146 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001147 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001148 self.assertRaises(TypeError, repeat, None, 3, 4)
1149 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001150 r = repeat(1+0j)
1151 self.assertEqual(repr(r), 'repeat((1+0j))')
1152 r = repeat(1+0j, 5)
1153 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
1154 list(r)
1155 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001156
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001157 # check copy, deepcopy, pickle
1158 c = repeat(object='a', times=10)
1159 self.assertEqual(next(c), 'a')
1160 self.assertEqual(take(2, copy.copy(c)), list('a' * 2))
1161 self.assertEqual(take(2, copy.deepcopy(c)), list('a' * 2))
Serhiy Storchakabad12572014-12-15 14:03:42 +02001162 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1163 self.pickletest(proto, repeat(object='a', times=10))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001164
Raymond Hettinger97d35552014-06-24 21:36:58 -07001165 def test_repeat_with_negative_times(self):
1166 self.assertEqual(repr(repeat('a', -1)), "repeat('a', 0)")
1167 self.assertEqual(repr(repeat('a', -2)), "repeat('a', 0)")
1168 self.assertEqual(repr(repeat('a', times=-1)), "repeat('a', 0)")
1169 self.assertEqual(repr(repeat('a', times=-2)), "repeat('a', 0)")
1170
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001171 def test_map(self):
1172 self.assertEqual(list(map(operator.pow, range(3), range(1,7))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001173 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001174 self.assertEqual(list(map(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001175 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001176 self.assertEqual(list(map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +00001177 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001178 self.assertEqual(take(2,map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +00001179 [('a',0),('b',1)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001180 self.assertEqual(list(map(operator.pow, [])), [])
1181 self.assertRaises(TypeError, map)
1182 self.assertRaises(TypeError, list, map(None, range(3), range(3)))
1183 self.assertRaises(TypeError, map, operator.neg)
1184 self.assertRaises(TypeError, next, map(10, range(5)))
1185 self.assertRaises(ValueError, next, map(errfunc, [4], [5]))
1186 self.assertRaises(TypeError, next, map(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001187
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001188 # check copy, deepcopy, pickle
1189 ans = [('a',0),('b',1),('c',2)]
1190
1191 c = map(tupleize, 'abc', count())
1192 self.assertEqual(list(copy.copy(c)), ans)
1193
1194 c = map(tupleize, 'abc', count())
1195 self.assertEqual(list(copy.deepcopy(c)), ans)
1196
Serhiy Storchakabad12572014-12-15 14:03:42 +02001197 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1198 c = map(tupleize, 'abc', count())
1199 self.pickletest(proto, c)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001200
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001201 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001202 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
1203 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001204 self.assertEqual(take(3, starmap(operator.pow, zip(count(), count(1)))),
Raymond Hettinger02420702003-06-29 20:36:23 +00001205 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001206 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +00001207 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
1208 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001209 self.assertRaises(TypeError, starmap)
1210 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001211 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
1212 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
1213 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001214
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001215 # check copy, deepcopy, pickle
1216 ans = [0**1, 1**2, 2**3]
1217
1218 c = starmap(operator.pow, zip(range(3), range(1,7)))
1219 self.assertEqual(list(copy.copy(c)), ans)
1220
1221 c = starmap(operator.pow, zip(range(3), range(1,7)))
1222 self.assertEqual(list(copy.deepcopy(c)), ans)
1223
Serhiy Storchakabad12572014-12-15 14:03:42 +02001224 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1225 c = starmap(operator.pow, zip(range(3), range(1,7)))
1226 self.pickletest(proto, c)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001227
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001228 def test_islice(self):
1229 for args in [ # islice(args) should agree with range(args)
1230 (10, 20, 3),
1231 (10, 3, 20),
1232 (10, 20),
Cheryl Sabellada1734c2018-03-26 21:29:33 -04001233 (10, 10),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001234 (10, 3),
1235 (20,)
1236 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +00001237 self.assertEqual(list(islice(range(100), *args)),
1238 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001239
1240 for args, tgtargs in [ # Stop when seqn is exhausted
1241 ((10, 110, 3), ((10, 100, 3))),
1242 ((10, 110), ((10, 100))),
1243 ((110,), (100,))
1244 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +00001245 self.assertEqual(list(islice(range(100), *args)),
1246 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001247
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001248 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +00001249 self.assertEqual(list(islice(range(10), None)), list(range(10)))
1250 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
1251 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
1252 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
1253 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001254
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +00001255 # Test number of items consumed SF #1171417
1256 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +00001257 self.assertEqual(list(islice(it, 3)), list(range(3)))
1258 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +00001259
Cheryl Sabellada1734c2018-03-26 21:29:33 -04001260 it = iter(range(10))
1261 self.assertEqual(list(islice(it, 3, 3)), [])
1262 self.assertEqual(list(it), list(range(3, 10)))
1263
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001264 # Test invalid arguments
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001265 ra = range(10)
1266 self.assertRaises(TypeError, islice, ra)
1267 self.assertRaises(TypeError, islice, ra, 1, 2, 3, 4)
1268 self.assertRaises(ValueError, islice, ra, -5, 10, 1)
1269 self.assertRaises(ValueError, islice, ra, 1, -5, -1)
1270 self.assertRaises(ValueError, islice, ra, 1, 10, -1)
1271 self.assertRaises(ValueError, islice, ra, 1, 10, 0)
1272 self.assertRaises(ValueError, islice, ra, 'a')
1273 self.assertRaises(ValueError, islice, ra, 'a', 1)
1274 self.assertRaises(ValueError, islice, ra, 1, 'a')
1275 self.assertRaises(ValueError, islice, ra, 'a', 1, 1)
1276 self.assertRaises(ValueError, islice, ra, 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +00001277 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001278
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001279 # Issue #10323: Less islice in a predictable state
1280 c = count()
1281 self.assertEqual(list(islice(c, 1, 3, 50)), [1])
1282 self.assertEqual(next(c), 3)
1283
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001284 # check copy, deepcopy, pickle
1285 for args in [ # islice(args) should agree with range(args)
1286 (10, 20, 3),
1287 (10, 3, 20),
1288 (10, 20),
1289 (10, 3),
1290 (20,)
1291 ]:
1292 self.assertEqual(list(copy.copy(islice(range(100), *args))),
1293 list(range(*args)))
1294 self.assertEqual(list(copy.deepcopy(islice(range(100), *args))),
1295 list(range(*args)))
Serhiy Storchakabad12572014-12-15 14:03:42 +02001296 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1297 self.pickletest(proto, islice(range(100), *args))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001298
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001299 # Issue #21321: check source iterator is not referenced
1300 # from islice() after the latter has been exhausted
1301 it = (x for x in (1, 2))
1302 wr = weakref.ref(it)
1303 it = islice(it, 1)
1304 self.assertIsNotNone(wr())
1305 list(it) # exhaust the iterator
Benjamin Peterson8e163512014-08-24 18:07:28 -05001306 support.gc_collect()
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001307 self.assertIsNone(wr())
1308
Will Roberts0ecdc522017-06-08 08:03:04 +02001309 # Issue #30537: islice can accept integer-like objects as
1310 # arguments
1311 class IntLike(object):
1312 def __init__(self, val):
1313 self.val = val
1314 def __index__(self):
1315 return self.val
1316 self.assertEqual(list(islice(range(100), IntLike(10))), list(range(10)))
1317 self.assertEqual(list(islice(range(100), IntLike(10), IntLike(50))),
1318 list(range(10, 50)))
1319 self.assertEqual(list(islice(range(100), IntLike(10), IntLike(50), IntLike(5))),
1320 list(range(10,50,5)))
1321
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001322 def test_takewhile(self):
1323 data = [1, 3, 5, 20, 2, 4, 6, 8]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001324 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001325 self.assertEqual(list(takewhile(underten, [])), [])
1326 self.assertRaises(TypeError, takewhile)
1327 self.assertRaises(TypeError, takewhile, operator.pow)
1328 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001329 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
1330 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001331 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
1332 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +00001333 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001334
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001335 # check copy, deepcopy, pickle
1336 self.assertEqual(list(copy.copy(takewhile(underten, data))), [1, 3, 5])
1337 self.assertEqual(list(copy.deepcopy(takewhile(underten, data))),
1338 [1, 3, 5])
Serhiy Storchakabad12572014-12-15 14:03:42 +02001339 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1340 self.pickletest(proto, takewhile(underten, data))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001341
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001342 def test_dropwhile(self):
1343 data = [1, 3, 5, 20, 2, 4, 6, 8]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001344 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001345 self.assertEqual(list(dropwhile(underten, [])), [])
1346 self.assertRaises(TypeError, dropwhile)
1347 self.assertRaises(TypeError, dropwhile, operator.pow)
1348 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001349 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
1350 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001351
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001352 # check copy, deepcopy, pickle
1353 self.assertEqual(list(copy.copy(dropwhile(underten, data))), [20, 2, 4, 6, 8])
1354 self.assertEqual(list(copy.deepcopy(dropwhile(underten, data))),
1355 [20, 2, 4, 6, 8])
Serhiy Storchakabad12572014-12-15 14:03:42 +02001356 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1357 self.pickletest(proto, dropwhile(underten, data))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001358
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001359 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +00001360 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001361
1362 a, b = tee([]) # test empty iterator
1363 self.assertEqual(list(a), [])
1364 self.assertEqual(list(b), [])
1365
1366 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +00001367 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001368
1369 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +00001370 self.assertEqual(list(a), list(range(n)))
1371 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001372
1373 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +00001374 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001375 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001376 del a
Guido van Rossum805365e2007-05-07 22:24:25 +00001377 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001378
1379 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +00001380 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001381 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001382 del b
Guido van Rossum805365e2007-05-07 22:24:25 +00001383 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001384
Guido van Rossum805365e2007-05-07 22:24:25 +00001385 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001386 order = [0]*n + [1]*n
1387 random.shuffle(order)
1388 lists = ([], [])
1389 its = tee(irange(n))
1390 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +00001391 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001392 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +00001393 self.assertEqual(lists[0], list(range(n)))
1394 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001395
Raymond Hettingerad983e72003-11-12 14:32:26 +00001396 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001397 self.assertRaises(TypeError, tee)
1398 self.assertRaises(TypeError, tee, 3)
1399 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +00001400 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001401
Raymond Hettingerad983e72003-11-12 14:32:26 +00001402 # tee object should be instantiable
1403 a, b = tee('abc')
1404 c = type(a)('def')
1405 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +00001406
Raymond Hettingerad983e72003-11-12 14:32:26 +00001407 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +00001408 a, b, c = tee(range(2000), 3)
1409 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001410 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +00001411 self.assertEqual(list(b), list(range(2000)))
1412 self.assertEqual([next(c), next(c)], list(range(2)))
1413 self.assertEqual(list(a), list(range(100,2000)))
1414 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +00001415
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001416 # test values of n
1417 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +00001418 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +00001419 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001420 result = tee('abc', n)
1421 self.assertEqual(type(result), tuple)
1422 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001423 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001424
Raymond Hettingerad983e72003-11-12 14:32:26 +00001425 # tee pass-through to copyable iterator
1426 a, b = tee('abc')
1427 c, d = tee(a)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +00001428 self.assertTrue(a is c)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001429
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001430 # test tee_new
1431 t1, t2 = tee('abc')
1432 tnew = type(t1)
1433 self.assertRaises(TypeError, tnew)
1434 self.assertRaises(TypeError, tnew, 10)
1435 t3 = tnew(t1)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +00001436 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +00001437
Raymond Hettingera9f60922004-10-17 16:40:14 +00001438 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +00001439 a, b = tee(range(10))
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001440 p = weakref.proxy(a)
Raymond Hettingera9f60922004-10-17 16:40:14 +00001441 self.assertEqual(getattr(p, '__class__'), type(b))
1442 del a
1443 self.assertRaises(ReferenceError, getattr, p, '__class__')
1444
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001445 ans = list('abc')
1446 long_ans = list(range(10000))
1447
1448 # check copy
1449 a, b = tee('abc')
1450 self.assertEqual(list(copy.copy(a)), ans)
1451 self.assertEqual(list(copy.copy(b)), ans)
1452 a, b = tee(list(range(10000)))
1453 self.assertEqual(list(copy.copy(a)), long_ans)
1454 self.assertEqual(list(copy.copy(b)), long_ans)
1455
1456 # check partially consumed copy
1457 a, b = tee('abc')
1458 take(2, a)
1459 take(1, b)
1460 self.assertEqual(list(copy.copy(a)), ans[2:])
1461 self.assertEqual(list(copy.copy(b)), ans[1:])
1462 self.assertEqual(list(a), ans[2:])
1463 self.assertEqual(list(b), ans[1:])
1464 a, b = tee(range(10000))
1465 take(100, a)
1466 take(60, b)
1467 self.assertEqual(list(copy.copy(a)), long_ans[100:])
1468 self.assertEqual(list(copy.copy(b)), long_ans[60:])
1469 self.assertEqual(list(a), long_ans[100:])
1470 self.assertEqual(list(b), long_ans[60:])
1471
1472 # check deepcopy
1473 a, b = tee('abc')
1474 self.assertEqual(list(copy.deepcopy(a)), ans)
1475 self.assertEqual(list(copy.deepcopy(b)), ans)
1476 self.assertEqual(list(a), ans)
1477 self.assertEqual(list(b), ans)
1478 a, b = tee(range(10000))
1479 self.assertEqual(list(copy.deepcopy(a)), long_ans)
1480 self.assertEqual(list(copy.deepcopy(b)), long_ans)
1481 self.assertEqual(list(a), long_ans)
1482 self.assertEqual(list(b), long_ans)
1483
1484 # check partially consumed deepcopy
1485 a, b = tee('abc')
1486 take(2, a)
1487 take(1, b)
1488 self.assertEqual(list(copy.deepcopy(a)), ans[2:])
1489 self.assertEqual(list(copy.deepcopy(b)), ans[1:])
1490 self.assertEqual(list(a), ans[2:])
1491 self.assertEqual(list(b), ans[1:])
1492 a, b = tee(range(10000))
1493 take(100, a)
1494 take(60, b)
1495 self.assertEqual(list(copy.deepcopy(a)), long_ans[100:])
1496 self.assertEqual(list(copy.deepcopy(b)), long_ans[60:])
1497 self.assertEqual(list(a), long_ans[100:])
1498 self.assertEqual(list(b), long_ans[60:])
1499
1500 # check pickle
Serhiy Storchakabad12572014-12-15 14:03:42 +02001501 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1502 self.pickletest(proto, iter(tee('abc')))
1503 a, b = tee('abc')
1504 self.pickletest(proto, a, compare=ans)
1505 self.pickletest(proto, b, compare=ans)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001506
Serhiy Storchakaa3e91282013-01-25 13:19:31 +02001507 # Issue 13454: Crash when deleting backward iterator from tee()
1508 def test_tee_del_backward(self):
Serhiy Storchaka5bb893c2013-01-26 11:52:06 +02001509 forward, backward = tee(repeat(None, 20000000))
Serhiy Storchaka9db55002015-03-28 20:38:37 +02001510 try:
1511 any(forward) # exhaust the iterator
1512 del backward
1513 except:
1514 del forward, backward
1515 raise
Serhiy Storchakaa3e91282013-01-25 13:19:31 +02001516
Serhiy Storchaka526a0142019-09-09 11:47:14 +03001517 def test_tee_reenter(self):
1518 class I:
1519 first = True
1520 def __iter__(self):
1521 return self
1522 def __next__(self):
1523 first = self.first
1524 self.first = False
1525 if first:
1526 return next(b)
1527
1528 a, b = tee(I())
1529 with self.assertRaisesRegex(RuntimeError, "tee"):
1530 next(a)
1531
1532 def test_tee_concurrent(self):
1533 start = threading.Event()
1534 finish = threading.Event()
1535 class I:
1536 def __iter__(self):
1537 return self
1538 def __next__(self):
1539 start.set()
1540 finish.wait()
1541
1542 a, b = tee(I())
1543 thread = threading.Thread(target=next, args=[a])
1544 thread.start()
1545 try:
1546 start.wait()
1547 with self.assertRaisesRegex(RuntimeError, "tee"):
1548 next(b)
1549 finally:
1550 finish.set()
1551 thread.join()
1552
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001553 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001554 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +00001555
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001556 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +00001557 self.assertRaises(StopIteration, next, f([]))
1558 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001559
Georg Brandla18af4e2007-04-21 15:47:16 +00001560 self.assertRaises(StopIteration, next, islice([], None))
1561 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001562
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001563 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +00001564 self.assertRaises(StopIteration, next, p)
1565 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001566 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +00001567 self.assertRaises(StopIteration, next, p)
1568 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001569
Georg Brandla18af4e2007-04-21 15:47:16 +00001570 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001571
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001572 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +00001573 self.assertRaises(StopIteration, next, f(lambda x:x, []))
1574 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001575
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001576class TestExamples(unittest.TestCase):
1577
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001578 def test_accumulate(self):
Raymond Hettinger482ba772010-12-01 22:48:00 +00001579 self.assertEqual(list(accumulate([1,2,3,4,5])), [1, 3, 6, 10, 15])
1580
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001581 def test_accumulate_reducible(self):
1582 # check copy, deepcopy, pickle
1583 data = [1, 2, 3, 4, 5]
1584 accumulated = [1, 3, 6, 10, 15]
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001585
Serhiy Storchakabad12572014-12-15 14:03:42 +02001586 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1587 it = accumulate(data)
1588 self.assertEqual(list(pickle.loads(pickle.dumps(it, proto))), accumulated[:])
1589 self.assertEqual(next(it), 1)
1590 self.assertEqual(list(pickle.loads(pickle.dumps(it, proto))), accumulated[1:])
1591 it = accumulate(data)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001592 self.assertEqual(next(it), 1)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001593 self.assertEqual(list(copy.deepcopy(it)), accumulated[1:])
1594 self.assertEqual(list(copy.copy(it)), accumulated[1:])
1595
Serhiy Storchakad5516252016-03-06 14:00:45 +02001596 def test_accumulate_reducible_none(self):
1597 # Issue #25718: total is None
1598 it = accumulate([None, None, None], operator.is_)
1599 self.assertEqual(next(it), None)
1600 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1601 it_copy = pickle.loads(pickle.dumps(it, proto))
1602 self.assertEqual(list(it_copy), [True, False])
1603 self.assertEqual(list(copy.deepcopy(it)), [True, False])
1604 self.assertEqual(list(copy.copy(it)), [True, False])
1605
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001606 def test_chain(self):
1607 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
1608
1609 def test_chain_from_iterable(self):
1610 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
1611
1612 def test_combinations(self):
1613 self.assertEqual(list(combinations('ABCD', 2)),
1614 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
1615 self.assertEqual(list(combinations(range(4), 3)),
1616 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
1617
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001618 def test_combinations_with_replacement(self):
1619 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
1620 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
1621
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001622 def test_compress(self):
1623 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
1624
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001625 def test_count(self):
1626 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
1627
1628 def test_cycle(self):
1629 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
1630
1631 def test_dropwhile(self):
1632 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
1633
1634 def test_groupby(self):
1635 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
1636 list('ABCDAB'))
1637 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
1638 [list('AAAA'), list('BBB'), list('CC'), list('D')])
1639
1640 def test_filter(self):
1641 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
1642
1643 def test_filterfalse(self):
1644 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
1645
1646 def test_map(self):
1647 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
1648
1649 def test_islice(self):
1650 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
1651 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
1652 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
1653 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1654
1655 def test_zip(self):
1656 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
1657
1658 def test_zip_longest(self):
1659 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
1660 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
1661
1662 def test_permutations(self):
1663 self.assertEqual(list(permutations('ABCD', 2)),
1664 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
1665 self.assertEqual(list(permutations(range(3))),
1666 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
1667
1668 def test_product(self):
1669 self.assertEqual(list(product('ABCD', 'xy')),
1670 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
1671 self.assertEqual(list(product(range(2), repeat=3)),
1672 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
1673 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
1674
1675 def test_repeat(self):
1676 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
1677
1678 def test_stapmap(self):
1679 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
1680 [32, 9, 1000])
1681
1682 def test_takewhile(self):
1683 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
1684
1685
Cheryl Sabellada1734c2018-03-26 21:29:33 -04001686class TestPurePythonRoughEquivalents(unittest.TestCase):
1687
1688 @staticmethod
1689 def islice(iterable, *args):
1690 s = slice(*args)
1691 start, stop, step = s.start or 0, s.stop or sys.maxsize, s.step or 1
1692 it = iter(range(start, stop, step))
1693 try:
1694 nexti = next(it)
1695 except StopIteration:
1696 # Consume *iterable* up to the *start* position.
1697 for i, element in zip(range(start), iterable):
1698 pass
1699 return
1700 try:
1701 for i, element in enumerate(iterable):
1702 if i == nexti:
1703 yield element
1704 nexti = next(it)
1705 except StopIteration:
1706 # Consume to *stop*.
1707 for i, element in zip(range(i + 1, stop), iterable):
1708 pass
1709
1710 def test_islice_recipe(self):
1711 self.assertEqual(list(self.islice('ABCDEFG', 2)), list('AB'))
1712 self.assertEqual(list(self.islice('ABCDEFG', 2, 4)), list('CD'))
1713 self.assertEqual(list(self.islice('ABCDEFG', 2, None)), list('CDEFG'))
1714 self.assertEqual(list(self.islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1715 # Test items consumed.
1716 it = iter(range(10))
1717 self.assertEqual(list(self.islice(it, 3)), list(range(3)))
1718 self.assertEqual(list(it), list(range(3, 10)))
1719 it = iter(range(10))
1720 self.assertEqual(list(self.islice(it, 3, 3)), [])
1721 self.assertEqual(list(it), list(range(3, 10)))
1722 # Test that slice finishes in predictable state.
1723 c = count()
1724 self.assertEqual(list(self.islice(c, 1, 3, 50)), [1])
1725 self.assertEqual(next(c), 3)
1726
1727
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001728class TestGC(unittest.TestCase):
1729
1730 def makecycle(self, iterator, container):
1731 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +00001732 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001733 del container, iterator
1734
Raymond Hettinger482ba772010-12-01 22:48:00 +00001735 def test_accumulate(self):
1736 a = []
1737 self.makecycle(accumulate([1,2,a,3]), a)
1738
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001739 def test_chain(self):
1740 a = []
1741 self.makecycle(chain(a), a)
1742
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001743 def test_chain_from_iterable(self):
1744 a = []
1745 self.makecycle(chain.from_iterable([a]), a)
1746
1747 def test_combinations(self):
1748 a = []
1749 self.makecycle(combinations([1,2,a,3], 3), a)
1750
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001751 def test_combinations_with_replacement(self):
1752 a = []
1753 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1754
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001755 def test_compress(self):
1756 a = []
1757 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1758
Raymond Hettingerd6280f42009-02-16 20:50:56 +00001759 def test_count(self):
1760 a = []
1761 Int = type('Int', (int,), dict(x=a))
1762 self.makecycle(count(Int(0), Int(1)), a)
1763
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001764 def test_cycle(self):
1765 a = []
1766 self.makecycle(cycle([a]*2), a)
1767
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001768 def test_dropwhile(self):
1769 a = []
1770 self.makecycle(dropwhile(bool, [0, a, a]), a)
1771
1772 def test_groupby(self):
1773 a = []
1774 self.makecycle(groupby([a]*2, lambda x:x), a)
1775
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001776 def test_issue2246(self):
1777 # Issue 2246 -- the _grouper iterator was not included in GC
1778 n = 10
1779 keyfunc = lambda x: x
1780 for i, j in groupby(range(n), key=keyfunc):
1781 keyfunc.__dict__.setdefault('x',[]).append(j)
1782
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001783 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001784 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001785 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001786
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001787 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001788 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001789 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001790
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001791 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001792 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001793 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001794
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001795 def test_zip_longest(self):
1796 a = []
1797 self.makecycle(zip_longest([a]*2, [a]*3), a)
1798 b = [a, None]
1799 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1800
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001801 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001802 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001803 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001804
1805 def test_islice(self):
1806 a = []
1807 self.makecycle(islice([a]*2, None), a)
1808
Raymond Hettingercc061d02020-11-30 20:42:54 -08001809 def test_pairwise(self):
1810 a = []
1811 self.makecycle(pairwise([a]*5), a)
1812
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001813 def test_permutations(self):
1814 a = []
1815 self.makecycle(permutations([1,2,a,3], 3), a)
1816
1817 def test_product(self):
1818 a = []
1819 self.makecycle(product([1,2,a,3], repeat=3), a)
1820
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001821 def test_repeat(self):
1822 a = []
1823 self.makecycle(repeat(a), a)
1824
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001825 def test_starmap(self):
1826 a = []
1827 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1828
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001829 def test_takewhile(self):
1830 a = []
1831 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1832
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001833def R(seqn):
1834 'Regular generator'
1835 for i in seqn:
1836 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001837
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001838class G:
1839 'Sequence using __getitem__'
1840 def __init__(self, seqn):
1841 self.seqn = seqn
1842 def __getitem__(self, i):
1843 return self.seqn[i]
1844
1845class I:
1846 'Sequence using iterator protocol'
1847 def __init__(self, seqn):
1848 self.seqn = seqn
1849 self.i = 0
1850 def __iter__(self):
1851 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001852 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001853 if self.i >= len(self.seqn): raise StopIteration
1854 v = self.seqn[self.i]
1855 self.i += 1
1856 return v
1857
1858class Ig:
1859 'Sequence using iterator protocol defined with a generator'
1860 def __init__(self, seqn):
1861 self.seqn = seqn
1862 self.i = 0
1863 def __iter__(self):
1864 for val in self.seqn:
1865 yield val
1866
1867class X:
1868 'Missing __getitem__ and __iter__'
1869 def __init__(self, seqn):
1870 self.seqn = seqn
1871 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001872 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001873 if self.i >= len(self.seqn): raise StopIteration
1874 v = self.seqn[self.i]
1875 self.i += 1
1876 return v
1877
1878class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001879 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001880 def __init__(self, seqn):
1881 self.seqn = seqn
1882 self.i = 0
1883 def __iter__(self):
1884 return self
1885
1886class E:
1887 'Test propagation of exceptions'
1888 def __init__(self, seqn):
1889 self.seqn = seqn
1890 self.i = 0
1891 def __iter__(self):
1892 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001893 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001894 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001895
1896class S:
1897 'Test immediate stop'
1898 def __init__(self, seqn):
1899 pass
1900 def __iter__(self):
1901 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001902 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001903 raise StopIteration
1904
1905def L(seqn):
1906 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001907 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001908
1909
1910class TestVariousIteratorArgs(unittest.TestCase):
1911
Raymond Hettinger482ba772010-12-01 22:48:00 +00001912 def test_accumulate(self):
1913 s = [1,2,3,4,5]
1914 r = [1,3,6,10,15]
1915 n = len(s)
1916 for g in (G, I, Ig, L, R):
1917 self.assertEqual(list(accumulate(g(s))), r)
1918 self.assertEqual(list(accumulate(S(s))), [])
1919 self.assertRaises(TypeError, accumulate, X(s))
1920 self.assertRaises(TypeError, accumulate, N(s))
1921 self.assertRaises(ZeroDivisionError, list, accumulate(E(s)))
1922
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001923 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001924 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001925 for g in (G, I, Ig, S, L, R):
1926 self.assertEqual(list(chain(g(s))), list(g(s)))
1927 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001928 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001929 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001930 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1931
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001932 def test_compress(self):
1933 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1934 n = len(s)
1935 for g in (G, I, Ig, S, L, R):
1936 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1937 self.assertRaises(TypeError, compress, X(s), repeat(1))
1938 self.assertRaises(TypeError, compress, N(s), repeat(1))
1939 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1940
Christian Heimesc3f30c42008-02-22 16:37:40 +00001941 def test_product(self):
1942 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1943 self.assertRaises(TypeError, product, X(s))
1944 self.assertRaises(TypeError, product, N(s))
1945 self.assertRaises(ZeroDivisionError, product, E(s))
1946
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001947 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001948 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001949 for g in (G, I, Ig, S, L, R):
1950 tgtlen = len(s) * 3
1951 expected = list(g(s))*3
1952 actual = list(islice(cycle(g(s)), tgtlen))
1953 self.assertEqual(actual, expected)
1954 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001955 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001956 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1957
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001958 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001959 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001960 for g in (G, I, Ig, S, L, R):
1961 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1962 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001963 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001964 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1965
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001966 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001967 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001968 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001969 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001970 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001971 self.assertRaises(TypeError, filter, isEven, X(s))
1972 self.assertRaises(TypeError, filter, isEven, N(s))
1973 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001974
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001975 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001976 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001977 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001978 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001979 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001980 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1981 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1982 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001983
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001984 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001985 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001986 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001987 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1988 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1989 self.assertRaises(TypeError, zip, X(s))
1990 self.assertRaises(TypeError, zip, N(s))
1991 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001992
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001993 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001994 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001995 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001996 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1997 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1998 self.assertRaises(TypeError, zip_longest, X(s))
1999 self.assertRaises(TypeError, zip_longest, N(s))
2000 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00002001
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002002 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00002003 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002004 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002005 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00002006 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002007 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00002008 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002009 self.assertRaises(TypeError, map, onearg, X(s))
2010 self.assertRaises(TypeError, map, onearg, N(s))
2011 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002012
2013 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00002014 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002015 for g in (G, I, Ig, S, L, R):
2016 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
2017 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00002018 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002019 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
2020
Raymond Hettingercc061d02020-11-30 20:42:54 -08002021 def test_pairwise(self):
2022 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
2023 for g in (G, I, Ig, S, L, R):
2024 seq = list(g(s))
2025 expected = list(zip(seq, seq[1:]))
2026 actual = list(pairwise(g(s)))
2027 self.assertEqual(actual, expected)
2028 self.assertRaises(TypeError, pairwise, X(s))
2029 self.assertRaises(TypeError, pairwise, N(s))
2030 self.assertRaises(ZeroDivisionError, list, pairwise(E(s)))
2031
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002032 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00002033 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002034 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00002035 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00002036 self.assertEqual(list(starmap(operator.pow, g(ss))),
2037 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002038 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00002039 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002040 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
2041
2042 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00002043 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002044 for g in (G, I, Ig, S, L, R):
2045 tgt = []
2046 for elem in g(s):
2047 if not isEven(elem): break
2048 tgt.append(elem)
2049 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
2050 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00002051 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002052 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
2053
2054 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00002055 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002056 for g in (G, I, Ig, S, L, R):
2057 tgt = []
2058 for elem in g(s):
2059 if not tgt and isOdd(elem): continue
2060 tgt.append(elem)
2061 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
2062 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00002063 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002064 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
2065
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002066 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00002067 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002068 for g in (G, I, Ig, S, L, R):
2069 it1, it2 = tee(g(s))
2070 self.assertEqual(list(it1), list(g(s)))
2071 self.assertEqual(list(it2), list(g(s)))
2072 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00002073 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002074 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
2075
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002076class LengthTransparency(unittest.TestCase):
2077
2078 def test_repeat(self):
Armin Ronacheraa9a79d2012-10-06 14:03:24 +02002079 self.assertEqual(operator.length_hint(repeat(None, 50)), 50)
Raymond Hettinger97d35552014-06-24 21:36:58 -07002080 self.assertEqual(operator.length_hint(repeat(None, 0)), 0)
Armin Ronacheraa9a79d2012-10-06 14:03:24 +02002081 self.assertEqual(operator.length_hint(repeat(None), 12), 12)
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002082
Raymond Hettinger97d35552014-06-24 21:36:58 -07002083 def test_repeat_with_negative_times(self):
2084 self.assertEqual(operator.length_hint(repeat(None, -1)), 0)
2085 self.assertEqual(operator.length_hint(repeat(None, -2)), 0)
2086 self.assertEqual(operator.length_hint(repeat(None, times=-1)), 0)
2087 self.assertEqual(operator.length_hint(repeat(None, times=-2)), 0)
2088
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002089class RegressionTests(unittest.TestCase):
2090
2091 def test_sf_793826(self):
2092 # Fix Armin Rigo's successful efforts to wreak havoc
2093
2094 def mutatingtuple(tuple1, f, tuple2):
2095 # this builds a tuple t which is a copy of tuple1,
2096 # then calls f(t), then mutates t to be equal to tuple2
2097 # (needs len(tuple1) == len(tuple2)).
2098 def g(value, first=[1]):
2099 if first:
2100 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00002101 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002102 return value
2103 items = list(tuple2)
2104 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002105 gen = map(g, items)
2106 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00002107 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002108
2109 def f(t):
2110 global T
2111 T = t
2112 first[:] = list(T)
2113
2114 first = []
2115 mutatingtuple((1,2,3), f, (4,5,6))
2116 second = list(T)
2117 self.assertEqual(first, second)
2118
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002119
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00002120 def test_sf_950057(self):
2121 # Make sure that chain() and cycle() catch exceptions immediately
2122 # rather than when shifting between input sources
2123
2124 def gen1():
2125 hist.append(0)
2126 yield 1
2127 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00002128 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00002129 hist.append(2)
2130
2131 def gen2(x):
2132 hist.append(3)
2133 yield 2
2134 hist.append(4)
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00002135
2136 hist = []
2137 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
2138 self.assertEqual(hist, [0,1])
2139
2140 hist = []
2141 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
2142 self.assertEqual(hist, [0,1])
2143
2144 hist = []
2145 self.assertRaises(AssertionError, list, cycle(gen1()))
2146 self.assertEqual(hist, [0,1])
2147
Neil Schemenauer52a48e62019-07-30 11:08:18 -07002148 @support.skip_if_pgo_task
T. Wouters5466d4a2017-03-30 09:58:35 -07002149 def test_long_chain_of_empty_iterables(self):
2150 # Make sure itertools.chain doesn't run into recursion limits when
2151 # dealing with long chains of empty iterables. Even with a high
2152 # number this would probably only fail in Py_DEBUG mode.
2153 it = chain.from_iterable(() for unused in range(10000000))
2154 with self.assertRaises(StopIteration):
2155 next(it)
2156
Serhiy Storchakac740e4f2017-09-26 21:47:56 +03002157 def test_issue30347_1(self):
2158 def f(n):
2159 if n == 5:
2160 list(b)
2161 return n != 6
2162 for (k, b) in groupby(range(10), f):
2163 list(b) # shouldn't crash
2164
2165 def test_issue30347_2(self):
2166 class K:
2167 def __init__(self, v):
2168 pass
2169 def __eq__(self, other):
2170 nonlocal i
2171 i += 1
2172 if i == 1:
2173 next(g, None)
2174 return True
2175 i = 0
2176 g = next(groupby(range(10), K))[1]
2177 for j in range(2):
2178 next(g, None) # shouldn't crash
2179
2180
Thomas Woutersb2137042007-02-01 18:02:27 +00002181class SubclassWithKwargsTest(unittest.TestCase):
2182 def test_keywords_in_subclass(self):
2183 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002184 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002185 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00002186 class Subclass(cls):
2187 def __init__(self, newarg=None, *args):
2188 cls.__init__(self, *args)
2189 try:
2190 Subclass(newarg=1)
2191 except TypeError as err:
2192 # we expect type errors because of wrong argument count
Serhiy Storchaka5eb788b2017-06-06 18:45:22 +03002193 self.assertNotIn("keyword arguments", err.args[0])
Thomas Woutersb2137042007-02-01 18:02:27 +00002194
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002195@support.cpython_only
2196class SizeofTest(unittest.TestCase):
2197 def setUp(self):
2198 self.ssize_t = struct.calcsize('n')
2199
2200 check_sizeof = support.check_sizeof
2201
2202 def test_product_sizeof(self):
2203 basesize = support.calcobjsize('3Pi')
2204 check = self.check_sizeof
2205 check(product('ab', '12'), basesize + 2 * self.ssize_t)
2206 check(product(*(('abc',) * 10)), basesize + 10 * self.ssize_t)
2207
2208 def test_combinations_sizeof(self):
2209 basesize = support.calcobjsize('3Pni')
2210 check = self.check_sizeof
2211 check(combinations('abcd', 3), basesize + 3 * self.ssize_t)
2212 check(combinations(range(10), 4), basesize + 4 * self.ssize_t)
2213
2214 def test_combinations_with_replacement_sizeof(self):
2215 cwr = combinations_with_replacement
2216 basesize = support.calcobjsize('3Pni')
2217 check = self.check_sizeof
2218 check(cwr('abcd', 3), basesize + 3 * self.ssize_t)
2219 check(cwr(range(10), 4), basesize + 4 * self.ssize_t)
2220
2221 def test_permutations_sizeof(self):
2222 basesize = support.calcobjsize('4Pni')
2223 check = self.check_sizeof
2224 check(permutations('abcd'),
2225 basesize + 4 * self.ssize_t + 4 * self.ssize_t)
2226 check(permutations('abcd', 3),
2227 basesize + 4 * self.ssize_t + 3 * self.ssize_t)
2228 check(permutations('abcde', 3),
2229 basesize + 5 * self.ssize_t + 3 * self.ssize_t)
2230 check(permutations(range(10), 4),
2231 basesize + 10 * self.ssize_t + 4 * self.ssize_t)
2232
Thomas Woutersb2137042007-02-01 18:02:27 +00002233
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002234libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002235
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002236
2237>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002238>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00002239... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00002240...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002241Check 1200 is for $120.15
2242Check 1201 is for $764.05
2243Check 1202 is for $823.14
2244
2245>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002246>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00002247... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00002248...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000022491
22508
225127
2252
2253>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00002254>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00002255... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00002256...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002257Alex
2258Laura
2259Martin
2260Walter
2261Samuele
2262
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002263>>> from operator import itemgetter
2264>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00002265>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002266>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00002267... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00002268...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000022691 ['a', 'c', 'e']
22702 ['b', 'd', 'f']
22713 ['g']
2272
Raymond Hettinger734fb572004-01-20 20:04:40 +00002273# Find runs of consecutive numbers using groupby. The key to the solution
2274# is differencing with a range so that consecutive numbers all appear in
2275# same group.
2276>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00002277>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00002278... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00002279...
Raymond Hettinger734fb572004-01-20 20:04:40 +00002280[1]
2281[4, 5, 6]
2282[10]
2283[15, 16, 17, 18]
2284[22]
2285[25, 26, 27, 28]
2286
Georg Brandl3dbca812008-07-23 16:10:53 +00002287>>> def take(n, iterable):
2288... "Return first n items of the iterable as a list"
2289... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00002290
Raymond Hettinger9265dd72018-04-08 08:44:20 -07002291>>> def prepend(value, iterator):
2292... "Prepend a single value in front of an iterator"
2293... # prepend(1, [2, 3, 4]) -> 1 2 3 4
2294... return chain([value], iterator)
2295
Georg Brandl3dbca812008-07-23 16:10:53 +00002296>>> def enumerate(iterable, start=0):
2297... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002298
Georg Brandl3dbca812008-07-23 16:10:53 +00002299>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002300... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00002301... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002302
Cheryl Sabellada1734c2018-03-26 21:29:33 -04002303>>> import collections
2304>>> def consume(iterator, n=None):
2305... "Advance the iterator n-steps ahead. If n is None, consume entirely."
2306... # Use functions that consume iterators at C speed.
2307... if n is None:
2308... # feed the entire iterator into a zero-length deque
2309... collections.deque(iterator, maxlen=0)
2310... else:
2311... # advance to the empty slice starting at position n
2312... next(islice(iterator, n, n), None)
2313
Raymond Hettingercdf8ba32009-02-19 04:45:07 +00002314>>> def nth(iterable, n, default=None):
2315... "Returns the nth item or a default value"
2316... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002317
Raymond Hettingere525ee32016-03-06 18:11:38 -08002318>>> def all_equal(iterable):
2319... "Returns True if all the elements are equal to each other"
2320... g = groupby(iterable)
2321... return next(g, True) and not next(g, False)
2322
Georg Brandl3dbca812008-07-23 16:10:53 +00002323>>> def quantify(iterable, pred=bool):
2324... "Count how many times the predicate is true"
2325... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00002326
Raymond Hettingerfc40b302020-11-29 10:47:22 -08002327>>> def pad_none(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002328... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00002329... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002330
Georg Brandl3dbca812008-07-23 16:10:53 +00002331>>> def ncycles(iterable, n):
Ezio Melotti13925002011-03-16 11:05:33 +02002332... "Returns the sequence elements n times"
Georg Brandl3dbca812008-07-23 16:10:53 +00002333... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002334
2335>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002336... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002337
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002338>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002339... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002340
2341>>> def repeatfunc(func, times=None, *args):
2342... "Repeat calls to func with specified arguments."
2343... " Example: repeatfunc(random.random)"
2344... if times is None:
2345... return starmap(func, repeat(args))
2346... else:
2347... return starmap(func, repeat(args, times))
2348
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002349>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00002350... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002351... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00002352... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002353
2354>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00002355... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002356... # Recipe credited to George Sakkis
2357... pending = len(iterables)
2358... nexts = cycle(iter(it).__next__ for it in iterables)
2359... while pending:
2360... try:
2361... for next in nexts:
2362... yield next()
2363... except StopIteration:
2364... pending -= 1
2365... nexts = cycle(islice(nexts, pending))
2366
2367>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00002368... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
2369... s = list(iterable)
2370... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002371
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00002372>>> def unique_everseen(iterable, key=None):
2373... "List unique elements, preserving order. Remember all elements ever seen."
2374... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
2375... # unique_everseen('ABBCcAD', str.lower) --> A B C D
2376... seen = set()
2377... seen_add = seen.add
2378... if key is None:
2379... for element in iterable:
2380... if element not in seen:
2381... seen_add(element)
2382... yield element
2383... else:
2384... for element in iterable:
2385... k = key(element)
2386... if k not in seen:
2387... seen_add(k)
2388... yield element
2389
2390>>> def unique_justseen(iterable, key=None):
2391... "List unique elements, preserving order. Remember only the element just seen."
2392... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
2393... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
2394... return map(next, map(itemgetter(1), groupby(iterable, key)))
2395
Raymond Hettinger31b26f62014-04-02 03:16:42 -07002396>>> def first_true(iterable, default=False, pred=None):
2397... '''Returns the first true value in the iterable.
2398...
2399... If no true value is found, returns *default*
2400...
2401... If *pred* is not None, returns the first item
2402... for which pred(item) is true.
2403...
2404... '''
2405... # first_true([a,b,c], x) --> a or b or c or x
2406... # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x
2407... return next(filter(pred, iterable), default)
2408
Raymond Hettingerd37258d2018-01-13 10:35:40 -08002409>>> def nth_combination(iterable, r, index):
2410... 'Equivalent to list(combinations(iterable, r))[index]'
2411... pool = tuple(iterable)
2412... n = len(pool)
2413... if r < 0 or r > n:
2414... raise ValueError
2415... c = 1
2416... k = min(r, n-r)
2417... for i in range(1, k+1):
2418... c = c * (n - k + i) // i
2419... if index < 0:
2420... index += c
2421... if index < 0 or index >= c:
2422... raise IndexError
2423... result = []
2424... while r:
2425... c, n, r = c*r//n, n-1, r-1
2426... while index >= c:
2427... index -= c
2428... c, n = c*(n-r)//n, n-1
2429... result.append(pool[-1-n])
2430... return tuple(result)
2431
2432
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002433This is not part of the examples but it tests to make sure the definitions
2434perform as purported.
2435
Raymond Hettingera098b332003-09-08 23:58:40 +00002436>>> take(10, count())
2437[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
2438
Raymond Hettinger9265dd72018-04-08 08:44:20 -07002439>>> list(prepend(1, [2, 3, 4]))
2440[1, 2, 3, 4]
2441
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002442>>> list(enumerate('abc'))
2443[(0, 'a'), (1, 'b'), (2, 'c')]
2444
2445>>> list(islice(tabulate(lambda x: 2*x), 4))
2446[0, 2, 4, 6]
2447
Cheryl Sabellada1734c2018-03-26 21:29:33 -04002448>>> it = iter(range(10))
2449>>> consume(it, 3)
2450>>> next(it)
24513
2452>>> consume(it)
2453>>> next(it, 'Done')
2454'Done'
2455
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002456>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00002457'd'
2458
2459>>> nth('abcde', 9) is None
2460True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002461
Raymond Hettingere525ee32016-03-06 18:11:38 -08002462>>> [all_equal(s) for s in ('', 'A', 'AAAA', 'AAAB', 'AAABA')]
2463[True, True, True, False, False]
2464
Guido van Rossum805365e2007-05-07 22:24:25 +00002465>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000246650
2467
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002468>>> a = [[1, 2, 3], [4, 5, 6]]
2469>>> flatten(a)
2470[1, 2, 3, 4, 5, 6]
2471
2472>>> list(repeatfunc(pow, 5, 2, 3))
2473[8, 8, 8, 8, 8]
2474
2475>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002476>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002477[0, 0, 0, 0, 0]
2478
Raymond Hettingerfc40b302020-11-29 10:47:22 -08002479>>> list(islice(pad_none('abc'), 0, 6))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002480['a', 'b', 'c', None, None, None]
2481
2482>>> list(ncycles('abc', 3))
2483['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
2484
2485>>> dotproduct([1,2,3], [4,5,6])
248632
2487
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002488>>> list(grouper(3, 'abcdefg', 'x'))
2489[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
2490
2491>>> list(roundrobin('abc', 'd', 'ef'))
2492['a', 'd', 'e', 'b', 'f', 'c']
2493
Raymond Hettingerace67332009-01-26 02:23:50 +00002494>>> list(powerset([1,2,3]))
2495[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002496
Raymond Hettinger191e8502009-01-27 13:29:43 +00002497>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
2498True
2499
2500>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
2501True
2502
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00002503>>> list(unique_everseen('AAAABBBCCDAABBB'))
2504['A', 'B', 'C', 'D']
2505
2506>>> list(unique_everseen('ABBCcAD', str.lower))
2507['A', 'B', 'C', 'D']
2508
2509>>> list(unique_justseen('AAAABBBCCDAABBB'))
2510['A', 'B', 'C', 'D', 'A', 'B']
2511
2512>>> list(unique_justseen('ABBCcAD', str.lower))
2513['A', 'B', 'C', 'A', 'D']
2514
Raymond Hettinger31b26f62014-04-02 03:16:42 -07002515>>> first_true('ABC0DEF1', '9', str.isdigit)
2516'0'
2517
Raymond Hettingerd37258d2018-01-13 10:35:40 -08002518>>> population = 'ABCDEFGH'
2519>>> for r in range(len(population) + 1):
2520... seq = list(combinations(population, r))
2521... for i in range(len(seq)):
2522... assert nth_combination(population, r, i) == seq[i]
2523... for i in range(-len(seq), 0):
2524... assert nth_combination(population, r, i) == seq[i]
2525
2526
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002527"""
2528
2529__test__ = {'libreftest' : libreftest}
2530
2531def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002532 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00002533 RegressionTests, LengthTransparency,
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002534 SubclassWithKwargsTest, TestExamples,
Cheryl Sabellada1734c2018-03-26 21:29:33 -04002535 TestPurePythonRoughEquivalents,
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002536 SizeofTest)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002537 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002538
2539 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002540 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00002541 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00002542 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00002543 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002544 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00002545 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00002546 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00002547 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002548
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002549 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002550 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002551
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002552if __name__ == "__main__":
2553 test_main(verbose=True)