blob: 6f8f87684e29cf1f9af529bf16a98e3ba8d9250c [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
Brandt Bucher226a0122020-12-04 19:45:57 -080015import gc
16
Benjamin Petersonee8712c2008-05-20 21:35:26 +000017maxsize = support.MAX_Py_ssize_t
Guido van Rossum360e4b82007-05-14 22:51:27 +000018minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000019
Guido van Rossum801f0d72006-08-24 19:48:10 +000020def lzip(*args):
21 return list(zip(*args))
22
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000023def onearg(x):
24 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000025 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000026
27def errfunc(*args):
28 'Test function that raises an error'
29 raise ValueError
30
31def gen3():
32 'Non-restartable source sequence'
33 for i in (0, 1, 2):
34 yield i
35
36def isEven(x):
37 'Test predicate'
38 return x%2==0
39
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000040def isOdd(x):
41 'Test predicate'
42 return x%2==1
43
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000044def tupleize(*args):
45 return args
46
47def irange(n):
48 for i in range(n):
49 yield i
50
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000051class StopNow:
52 'Class emulating an empty iterable.'
53 def __iter__(self):
54 return self
Georg Brandla18af4e2007-04-21 15:47:16 +000055 def __next__(self):
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000056 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000057
Raymond Hettinger02420702003-06-29 20:36:23 +000058def take(n, seq):
59 'Convenience function for partially consuming a long of infinite iterable'
60 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000061
Christian Heimes78644762008-03-04 23:39:23 +000062def prod(iterable):
63 return reduce(operator.mul, iterable, 1)
64
Christian Heimes380f7f22008-02-28 11:19:05 +000065def fact(n):
66 'Factorial'
Christian Heimes78644762008-03-04 23:39:23 +000067 return prod(range(1, n+1))
68
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000069# root level methods for pickling ability
70def testR(r):
71 return r[0]
72
73def testR2(r):
74 return r[2]
75
76def underten(x):
77 return x<10
78
Serhiy Storchakabad12572014-12-15 14:03:42 +020079picklecopiers = [lambda s, proto=proto: pickle.loads(pickle.dumps(s, proto))
80 for proto in range(pickle.HIGHEST_PROTOCOL + 1)]
81
Raymond Hettinger96ef8112003-02-01 00:10:11 +000082class TestBasicOps(unittest.TestCase):
Raymond Hettinger482ba772010-12-01 22:48:00 +000083
Serhiy Storchakabad12572014-12-15 14:03:42 +020084 def pickletest(self, protocol, it, stop=4, take=1, compare=None):
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000085 """Test that an iterator is the same after pickling, also when part-consumed"""
86 def expand(it, i=0):
87 # Recursively expand iterables, within sensible bounds
88 if i > 10:
89 raise RuntimeError("infinite recursion encountered")
90 if isinstance(it, str):
91 return it
92 try:
93 l = list(islice(it, stop))
94 except TypeError:
95 return it # can't expand it
96 return [expand(e, i+1) for e in l]
97
98 # Test the initial copy against the original
Serhiy Storchakabad12572014-12-15 14:03:42 +020099 dump = pickle.dumps(it, protocol)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000100 i2 = pickle.loads(dump)
101 self.assertEqual(type(it), type(i2))
102 a, b = expand(it), expand(i2)
103 self.assertEqual(a, b)
104 if compare:
105 c = expand(compare)
106 self.assertEqual(a, c)
107
108 # Take from the copy, and create another copy and compare them.
109 i3 = pickle.loads(dump)
110 took = 0
111 try:
112 for i in range(take):
113 next(i3)
114 took += 1
115 except StopIteration:
116 pass #in case there is less data than 'take'
Serhiy Storchakabad12572014-12-15 14:03:42 +0200117 dump = pickle.dumps(i3, protocol)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000118 i4 = pickle.loads(dump)
119 a, b = expand(i3), expand(i4)
120 self.assertEqual(a, b)
121 if compare:
122 c = expand(compare[took:])
123 self.assertEqual(a, c);
124
Raymond Hettinger482ba772010-12-01 22:48:00 +0000125 def test_accumulate(self):
126 self.assertEqual(list(accumulate(range(10))), # one positional arg
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000127 [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])
128 self.assertEqual(list(accumulate(iterable=range(10))), # kw arg
129 [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])
Raymond Hettinger482ba772010-12-01 22:48:00 +0000130 for typ in int, complex, Decimal, Fraction: # multiple types
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000131 self.assertEqual(
132 list(accumulate(map(typ, range(10)))),
Raymond Hettinger482ba772010-12-01 22:48:00 +0000133 list(map(typ, [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])))
Raymond Hettinger2d93e6e2010-12-03 02:33:53 +0000134 self.assertEqual(list(accumulate('abc')), ['a', 'ab', 'abc']) # works with non-numeric
Raymond Hettinger482ba772010-12-01 22:48:00 +0000135 self.assertEqual(list(accumulate([])), []) # empty iterable
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000136 self.assertEqual(list(accumulate([7])), [7]) # iterable of length one
Raymond Hettinger5d446132011-03-27 18:52:10 -0700137 self.assertRaises(TypeError, accumulate, range(10), 5, 6) # too many args
Raymond Hettinger482ba772010-12-01 22:48:00 +0000138 self.assertRaises(TypeError, accumulate) # too few args
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000139 self.assertRaises(TypeError, accumulate, x=range(10)) # unexpected kwd arg
Raymond Hettinger482ba772010-12-01 22:48:00 +0000140 self.assertRaises(TypeError, list, accumulate([1, []])) # args that don't add
141
Raymond Hettinger5d446132011-03-27 18:52:10 -0700142 s = [2, 8, 9, 5, 7, 0, 3, 4, 1, 6]
143 self.assertEqual(list(accumulate(s, min)),
144 [2, 2, 2, 2, 2, 0, 0, 0, 0, 0])
145 self.assertEqual(list(accumulate(s, max)),
146 [2, 8, 9, 9, 9, 9, 9, 9, 9, 9])
147 self.assertEqual(list(accumulate(s, operator.mul)),
148 [2, 16, 144, 720, 5040, 0, 0, 0, 0, 0])
149 with self.assertRaises(TypeError):
150 list(accumulate(s, chr)) # unary-operation
Serhiy Storchakabad12572014-12-15 14:03:42 +0200151 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
152 self.pickletest(proto, accumulate(range(10))) # test pickling
Lisa Roach9718b592018-09-23 17:34:59 -0700153 self.pickletest(proto, accumulate(range(10), initial=7))
154 self.assertEqual(list(accumulate([10, 5, 1], initial=None)), [10, 15, 16])
155 self.assertEqual(list(accumulate([10, 5, 1], initial=100)), [100, 110, 115, 116])
156 self.assertEqual(list(accumulate([], initial=100)), [100])
157 with self.assertRaises(TypeError):
158 list(accumulate([10, 20], 100))
Raymond Hettinger5d446132011-03-27 18:52:10 -0700159
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000160 def test_chain(self):
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000161
162 def chain2(*iterables):
163 'Pure python version in the docs'
164 for it in iterables:
165 for element in it:
166 yield element
167
168 for c in (chain, chain2):
169 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
170 self.assertEqual(list(c('abc')), list('abc'))
171 self.assertEqual(list(c('')), [])
172 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
173 self.assertRaises(TypeError, list,c(2, 3))
Christian Heimesf16baeb2008-02-29 14:57:44 +0000174
175 def test_chain_from_iterable(self):
176 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
177 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
178 self.assertEqual(list(chain.from_iterable([''])), [])
179 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
180 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000181
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000182 def test_chain_reducible(self):
Serhiy Storchakabad12572014-12-15 14:03:42 +0200183 for oper in [copy.deepcopy] + picklecopiers:
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000184 it = chain('abc', 'def')
185 self.assertEqual(list(oper(it)), list('abcdef'))
186 self.assertEqual(next(it), 'a')
187 self.assertEqual(list(oper(it)), list('bcdef'))
188
189 self.assertEqual(list(oper(chain(''))), [])
190 self.assertEqual(take(4, oper(chain('abc', 'def'))), list('abcd'))
191 self.assertRaises(TypeError, list, oper(chain(2, 3)))
Serhiy Storchakabad12572014-12-15 14:03:42 +0200192 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
193 self.pickletest(proto, chain('abc', 'def'), compare=list('abcdef'))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000194
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300195 def test_chain_setstate(self):
196 self.assertRaises(TypeError, chain().__setstate__, ())
197 self.assertRaises(TypeError, chain().__setstate__, [])
198 self.assertRaises(TypeError, chain().__setstate__, 0)
199 self.assertRaises(TypeError, chain().__setstate__, ([],))
200 self.assertRaises(TypeError, chain().__setstate__, (iter([]), []))
201 it = chain()
202 it.__setstate__((iter(['abc', 'def']),))
203 self.assertEqual(list(it), ['a', 'b', 'c', 'd', 'e', 'f'])
204 it = chain()
205 it.__setstate__((iter(['abc', 'def']), iter(['ghi'])))
206 self.assertEqual(list(it), ['ghi', 'a', 'b', 'c', 'd', 'e', 'f'])
207
Christian Heimes380f7f22008-02-28 11:19:05 +0000208 def test_combinations(self):
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000209 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Christian Heimes380f7f22008-02-28 11:19:05 +0000210 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Christian Heimes78644762008-03-04 23:39:23 +0000211 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Christian Heimes380f7f22008-02-28 11:19:05 +0000212 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000213
Serhiy Storchakabad12572014-12-15 14:03:42 +0200214 for op in [lambda a:a] + picklecopiers:
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000215 self.assertEqual(list(op(combinations('abc', 32))), []) # r > n
216
217 self.assertEqual(list(op(combinations('ABCD', 2))),
218 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
219 testIntermediate = combinations('ABCD', 2)
220 next(testIntermediate)
221 self.assertEqual(list(op(testIntermediate)),
222 [('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
223
224 self.assertEqual(list(op(combinations(range(4), 3))),
225 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
226 testIntermediate = combinations(range(4), 3)
227 next(testIntermediate)
228 self.assertEqual(list(op(testIntermediate)),
229 [(0,1,3), (0,2,3), (1,2,3)])
230
Christian Heimes78644762008-03-04 23:39:23 +0000231
232 def combinations1(iterable, r):
233 'Pure python version shown in the docs'
234 pool = tuple(iterable)
235 n = len(pool)
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000236 if r > n:
237 return
Christian Heimes78644762008-03-04 23:39:23 +0000238 indices = list(range(r))
239 yield tuple(pool[i] for i in indices)
240 while 1:
241 for i in reversed(range(r)):
242 if indices[i] != i + n - r:
243 break
244 else:
245 return
246 indices[i] += 1
247 for j in range(i+1, r):
248 indices[j] = indices[j-1] + 1
249 yield tuple(pool[i] for i in indices)
250
251 def combinations2(iterable, r):
252 'Pure python version shown in the docs'
253 pool = tuple(iterable)
254 n = len(pool)
255 for indices in permutations(range(n), r):
256 if sorted(indices) == list(indices):
257 yield tuple(pool[i] for i in indices)
258
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000259 def combinations3(iterable, r):
260 'Pure python version from cwr()'
261 pool = tuple(iterable)
262 n = len(pool)
263 for indices in combinations_with_replacement(range(n), r):
264 if len(set(indices)) == r:
265 yield tuple(pool[i] for i in indices)
266
Christian Heimes78644762008-03-04 23:39:23 +0000267 for n in range(7):
Christian Heimes380f7f22008-02-28 11:19:05 +0000268 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000269 for r in range(n+2):
Christian Heimes380f7f22008-02-28 11:19:05 +0000270 result = list(combinations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000271 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 +0000272 self.assertEqual(len(result), len(set(result))) # no repeats
273 self.assertEqual(result, sorted(result)) # lexicographic order
274 for c in result:
275 self.assertEqual(len(c), r) # r-length combinations
276 self.assertEqual(len(set(c)), r) # no duplicate elements
277 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000278 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000279 self.assertEqual(list(c),
280 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000281 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000282 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000283 self.assertEqual(result, list(combinations3(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000284
Serhiy Storchakabad12572014-12-15 14:03:42 +0200285 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
286 self.pickletest(proto, combinations(values, r)) # test pickling
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000287
Benjamin Peterson4b40eeb2015-02-01 20:59:00 -0500288 @support.bigaddrspacetest
289 def test_combinations_overflow(self):
Serhiy Storchakadee948b2015-02-03 01:34:09 +0200290 with self.assertRaises((OverflowError, MemoryError)):
Benjamin Peterson4b40eeb2015-02-01 20:59:00 -0500291 combinations("AA", 2**29)
292
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000293 # Test implementation detail: tuple re-use
Alex Gaynore151d212011-07-17 16:21:30 -0700294 @support.impl_detail("tuple reuse is specific to CPython")
295 def test_combinations_tuple_reuse(self):
Christian Heimes78644762008-03-04 23:39:23 +0000296 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
297 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
298
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000299 def test_combinations_with_replacement(self):
300 cwr = combinations_with_replacement
301 self.assertRaises(TypeError, cwr, 'abc') # missing r argument
302 self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
303 self.assertRaises(TypeError, cwr, None) # pool is not iterable
304 self.assertRaises(ValueError, cwr, 'abc', -2) # r is negative
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000305
Serhiy Storchakabad12572014-12-15 14:03:42 +0200306 for op in [lambda a:a] + picklecopiers:
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000307 self.assertEqual(list(op(cwr('ABC', 2))),
308 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
309 testIntermediate = cwr('ABC', 2)
310 next(testIntermediate)
311 self.assertEqual(list(op(testIntermediate)),
312 [('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
313
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000314
315 def cwr1(iterable, r):
316 'Pure python version shown in the docs'
317 # number items returned: (n+r-1)! / r! / (n-1)! when n>0
318 pool = tuple(iterable)
319 n = len(pool)
320 if not n and r:
321 return
322 indices = [0] * r
323 yield tuple(pool[i] for i in indices)
324 while 1:
325 for i in reversed(range(r)):
326 if indices[i] != n - 1:
327 break
328 else:
329 return
330 indices[i:] = [indices[i] + 1] * (r - i)
331 yield tuple(pool[i] for i in indices)
332
333 def cwr2(iterable, r):
334 'Pure python version shown in the docs'
335 pool = tuple(iterable)
336 n = len(pool)
337 for indices in product(range(n), repeat=r):
338 if sorted(indices) == list(indices):
339 yield tuple(pool[i] for i in indices)
340
341 def numcombs(n, r):
342 if not n:
343 return 0 if r else 1
344 return fact(n+r-1) / fact(r)/ fact(n-1)
345
346 for n in range(7):
347 values = [5*x-12 for x in range(n)]
348 for r in range(n+2):
349 result = list(cwr(values, r))
350
351 self.assertEqual(len(result), numcombs(n, r)) # right number of combs
352 self.assertEqual(len(result), len(set(result))) # no repeats
353 self.assertEqual(result, sorted(result)) # lexicographic order
354
355 regular_combs = list(combinations(values, r)) # compare to combs without replacement
356 if n == 0 or r <= 1:
Ezio Melottib3aedd42010-11-20 19:04:17 +0000357 self.assertEqual(result, regular_combs) # cases that should be identical
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000358 else:
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000359 self.assertTrue(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000360
361 for c in result:
362 self.assertEqual(len(c), r) # r-length combinations
363 noruns = [k for k,v in groupby(c)] # combo without consecutive repeats
364 self.assertEqual(len(noruns), len(set(noruns))) # no repeats other than consecutive
365 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000366 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000367 self.assertEqual(noruns,
368 [e for e in values if e in c]) # comb is a subsequence of the input iterable
369 self.assertEqual(result, list(cwr1(values, r))) # matches first pure python version
370 self.assertEqual(result, list(cwr2(values, r))) # matches second pure python version
371
Serhiy Storchakabad12572014-12-15 14:03:42 +0200372 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
373 self.pickletest(proto, cwr(values,r)) # test pickling
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000374
Benjamin Peterson6f082292015-02-01 21:10:47 -0500375 @support.bigaddrspacetest
376 def test_combinations_with_replacement_overflow(self):
Serhiy Storchakadee948b2015-02-03 01:34:09 +0200377 with self.assertRaises((OverflowError, MemoryError)):
Benjamin Peterson6f082292015-02-01 21:10:47 -0500378 combinations_with_replacement("AA", 2**30)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000379
Benjamin Peterson6f082292015-02-01 21:10:47 -0500380 # Test implementation detail: tuple re-use
Alex Gaynore151d212011-07-17 16:21:30 -0700381 @support.impl_detail("tuple reuse is specific to CPython")
382 def test_combinations_with_replacement_tuple_reuse(self):
383 cwr = combinations_with_replacement
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000384 self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
385 self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
386
Christian Heimes78644762008-03-04 23:39:23 +0000387 def test_permutations(self):
388 self.assertRaises(TypeError, permutations) # too few arguments
389 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000390 self.assertRaises(TypeError, permutations, None) # pool is not iterable
391 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000392 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000393 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Christian Heimes78644762008-03-04 23:39:23 +0000394 self.assertEqual(list(permutations(range(3), 2)),
395 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
396
397 def permutations1(iterable, r=None):
398 'Pure python version shown in the docs'
399 pool = tuple(iterable)
400 n = len(pool)
401 r = n if r is None else r
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000402 if r > n:
403 return
Christian Heimes78644762008-03-04 23:39:23 +0000404 indices = list(range(n))
405 cycles = list(range(n-r+1, n+1))[::-1]
406 yield tuple(pool[i] for i in indices[:r])
407 while n:
408 for i in reversed(range(r)):
409 cycles[i] -= 1
410 if cycles[i] == 0:
411 indices[i:] = indices[i+1:] + indices[i:i+1]
412 cycles[i] = n - i
413 else:
414 j = cycles[i]
415 indices[i], indices[-j] = indices[-j], indices[i]
416 yield tuple(pool[i] for i in indices[:r])
417 break
418 else:
419 return
420
421 def permutations2(iterable, r=None):
422 'Pure python version shown in the docs'
423 pool = tuple(iterable)
424 n = len(pool)
425 r = n if r is None else r
426 for indices in product(range(n), repeat=r):
427 if len(set(indices)) == r:
428 yield tuple(pool[i] for i in indices)
429
430 for n in range(7):
431 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000432 for r in range(n+2):
Christian Heimes78644762008-03-04 23:39:23 +0000433 result = list(permutations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000434 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 +0000435 self.assertEqual(len(result), len(set(result))) # no repeats
436 self.assertEqual(result, sorted(result)) # lexicographic order
437 for p in result:
438 self.assertEqual(len(p), r) # r-length permutations
439 self.assertEqual(len(set(p)), r) # no duplicate elements
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000440 self.assertTrue(all(e in values for e in p)) # elements taken from input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000441 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000442 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000443 if r == n:
444 self.assertEqual(result, list(permutations(values, None))) # test r as None
445 self.assertEqual(result, list(permutations(values))) # test default r
446
Serhiy Storchakabad12572014-12-15 14:03:42 +0200447 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
448 self.pickletest(proto, permutations(values, r)) # test pickling
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000449
Benjamin Peterson0eaabf12015-02-01 21:34:07 -0500450 @support.bigaddrspacetest
451 def test_permutations_overflow(self):
Serhiy Storchakadee948b2015-02-03 01:34:09 +0200452 with self.assertRaises((OverflowError, MemoryError)):
Benjamin Peterson0eaabf12015-02-01 21:34:07 -0500453 permutations("A", 2**30)
Benjamin Peterson0eaabf12015-02-01 21:34:07 -0500454
Zachary Waredca807b2014-04-24 13:22:05 -0500455 @support.impl_detail("tuple reuse is specific to CPython")
Alex Gaynore151d212011-07-17 16:21:30 -0700456 def test_permutations_tuple_reuse(self):
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000457 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Christian Heimes78644762008-03-04 23:39:23 +0000458 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Christian Heimes380f7f22008-02-28 11:19:05 +0000459
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000460 def test_combinatorics(self):
461 # Test relationships between product(), permutations(),
462 # combinations() and combinations_with_replacement().
463
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000464 for n in range(6):
465 s = 'ABCDEFG'[:n]
466 for r in range(8):
467 prod = list(product(s, repeat=r))
468 cwr = list(combinations_with_replacement(s, r))
469 perm = list(permutations(s, r))
470 comb = list(combinations(s, r))
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000471
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000472 # Check size
Ezio Melottib3aedd42010-11-20 19:04:17 +0000473 self.assertEqual(len(prod), n**r)
474 self.assertEqual(len(cwr), (fact(n+r-1) / fact(r)/ fact(n-1)) if n else (not r))
475 self.assertEqual(len(perm), 0 if r>n else fact(n) / fact(n-r))
476 self.assertEqual(len(comb), 0 if r>n else fact(n) / fact(r) / fact(n-r))
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000477
478 # Check lexicographic order without repeated tuples
Ezio Melottib3aedd42010-11-20 19:04:17 +0000479 self.assertEqual(prod, sorted(set(prod)))
480 self.assertEqual(cwr, sorted(set(cwr)))
481 self.assertEqual(perm, sorted(set(perm)))
482 self.assertEqual(comb, sorted(set(comb)))
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000483
484 # Check interrelationships
Ezio Melottib3aedd42010-11-20 19:04:17 +0000485 self.assertEqual(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
486 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 +0000487 self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
488 self.assertEqual(comb, [t for t in cwr if len(set(t))==r]) # comb: cwrs without dups
489 self.assertEqual(comb, list(filter(set(cwr).__contains__, perm))) # comb: perm that is a cwr
490 self.assertEqual(comb, list(filter(set(perm).__contains__, cwr))) # comb: cwr that is a perm
491 self.assertEqual(comb, sorted(set(cwr) & set(perm))) # comb: both a cwr and a perm
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000492
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000493 def test_compress(self):
Raymond Hettinger15a49502009-02-19 02:17:09 +0000494 self.assertEqual(list(compress(data='ABCDEF', selectors=[1,0,1,0,1,1])), list('ACEF'))
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000495 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
496 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
497 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
498 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
499 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
500 n = 10000
501 data = chain.from_iterable(repeat(range(6), n))
502 selectors = chain.from_iterable(repeat((0, 1)))
503 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
504 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
505 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
506 self.assertRaises(TypeError, compress, range(6)) # too few args
507 self.assertRaises(TypeError, compress, range(6), None) # too many args
508
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000509 # check copy, deepcopy, pickle
Serhiy Storchakabad12572014-12-15 14:03:42 +0200510 for op in [lambda a:copy.copy(a), lambda a:copy.deepcopy(a)] + picklecopiers:
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000511 for data, selectors, result1, result2 in [
512 ('ABCDEF', [1,0,1,0,1,1], 'ACEF', 'CEF'),
513 ('ABCDEF', [0,0,0,0,0,0], '', ''),
514 ('ABCDEF', [1,1,1,1,1,1], 'ABCDEF', 'BCDEF'),
515 ('ABCDEF', [1,0,1], 'AC', 'C'),
516 ('ABC', [0,1,1,1,1,1], 'BC', 'C'),
517 ]:
518
519 self.assertEqual(list(op(compress(data=data, selectors=selectors))), list(result1))
520 self.assertEqual(list(op(compress(data, selectors))), list(result1))
521 testIntermediate = compress(data, selectors)
522 if result1:
523 next(testIntermediate)
524 self.assertEqual(list(op(testIntermediate)), list(result2))
525
526
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000527 def test_count(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000528 self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
529 self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
530 self.assertEqual(take(2, lzip('abc',count(3))), [('a', 3), ('b', 4)])
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000531 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
532 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000533 self.assertRaises(TypeError, count, 2, 3, 4)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000534 self.assertRaises(TypeError, count, 'a')
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(10, count(-maxsize-5)),
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000538 list(range(-maxsize-5, -maxsize+5)))
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300539 self.assertEqual(take(3, count(3.25)), [3.25, 4.25, 5.25])
540 self.assertEqual(take(3, count(3.25-4j)), [3.25-4j, 4.25-4j, 5.25-4j])
541 self.assertEqual(take(3, count(Decimal('1.1'))),
542 [Decimal('1.1'), Decimal('2.1'), Decimal('3.1')])
543 self.assertEqual(take(3, count(Fraction(2, 3))),
544 [Fraction(2, 3), Fraction(5, 3), Fraction(8, 3)])
545 BIGINT = 1<<1000
546 self.assertEqual(take(3, count(BIGINT)), [BIGINT, BIGINT+1, BIGINT+2])
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000547 c = count(3)
548 self.assertEqual(repr(c), 'count(3)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000549 next(c)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000550 self.assertEqual(repr(c), 'count(4)')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000551 c = count(-9)
552 self.assertEqual(repr(c), 'count(-9)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000553 next(c)
554 self.assertEqual(next(c), -8)
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300555 self.assertEqual(repr(count(10.25)), 'count(10.25)')
556 self.assertEqual(repr(count(10.0)), 'count(10.0)')
557 self.assertEqual(type(next(count(10.0))), float)
Christian Heimesa37d4c62007-12-04 23:02:19 +0000558 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 +0300559 # Test repr
560 r1 = repr(count(i))
561 r2 = 'count(%r)'.__mod__(i)
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000562 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000563
Raymond Hettingera3e1ad22009-11-30 22:02:31 +0000564 # check copy, deepcopy, pickle
565 for value in -3, 3, maxsize-5, maxsize+5:
566 c = count(value)
567 self.assertEqual(next(copy.copy(c)), value)
568 self.assertEqual(next(copy.deepcopy(c)), value)
Serhiy Storchakabad12572014-12-15 14:03:42 +0200569 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
570 self.pickletest(proto, count(value))
Raymond Hettingera3e1ad22009-11-30 22:02:31 +0000571
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +0000572 #check proper internal error handling for large "step' sizes
573 count(1, maxsize+5); sys.exc_info()
574
Raymond Hettinger30729212009-02-12 06:28:27 +0000575 def test_count_with_stride(self):
576 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingereb13fdd2009-02-21 22:30:12 +0000577 self.assertEqual(lzip('abc',count(start=2,step=3)),
578 [('a', 2), ('b', 5), ('c', 8)])
579 self.assertEqual(lzip('abc',count(step=-1)),
580 [('a', 0), ('b', -1), ('c', -2)])
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300581 self.assertRaises(TypeError, count, 'a', 'b')
Raymond Hettinger30729212009-02-12 06:28:27 +0000582 self.assertEqual(lzip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
583 self.assertEqual(lzip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000584 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000585 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
586 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300587 self.assertEqual(take(3, count(10, maxsize+5)),
588 list(range(10, 10+3*(maxsize+5), maxsize+5)))
589 self.assertEqual(take(3, count(2, 1.25)), [2, 3.25, 4.5])
Raymond Hettinger30729212009-02-12 06:28:27 +0000590 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettinger8a882a72009-02-12 12:08:18 +0000591 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
592 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerde09acf2009-02-12 12:53:53 +0000593 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
594 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300595 BIGINT = 1<<1000
596 self.assertEqual(take(3, count(step=BIGINT)), [0, BIGINT, 2*BIGINT])
Raymond Hettinger30729212009-02-12 06:28:27 +0000597 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
598 c = count(3, 5)
599 self.assertEqual(repr(c), 'count(3, 5)')
600 next(c)
601 self.assertEqual(repr(c), 'count(8, 5)')
602 c = count(-9, 0)
603 self.assertEqual(repr(c), 'count(-9, 0)')
604 next(c)
605 self.assertEqual(repr(c), 'count(-9, 0)')
606 c = count(-9, -3)
607 self.assertEqual(repr(c), 'count(-9, -3)')
608 next(c)
609 self.assertEqual(repr(c), 'count(-12, -3)')
610 self.assertEqual(repr(c), 'count(-12, -3)')
611 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
612 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
613 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 +0300614 self.assertEqual(repr(count(10, 1.00)), 'count(10, 1.0)')
615 c = count(10, 1.0)
616 self.assertEqual(type(next(c)), int)
617 self.assertEqual(type(next(c)), float)
Raymond Hettinger30729212009-02-12 06:28:27 +0000618 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
619 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 +0300620 # Test repr
621 r1 = repr(count(i, j))
Raymond Hettinger30729212009-02-12 06:28:27 +0000622 if j == 1:
Serhiy Storchaka95949422013-08-27 19:40:23 +0300623 r2 = ('count(%r)' % i)
Raymond Hettinger30729212009-02-12 06:28:27 +0000624 else:
Serhiy Storchaka95949422013-08-27 19:40:23 +0300625 r2 = ('count(%r, %r)' % (i, j))
Raymond Hettinger30729212009-02-12 06:28:27 +0000626 self.assertEqual(r1, r2)
Serhiy Storchakabad12572014-12-15 14:03:42 +0200627 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
628 self.pickletest(proto, count(i, j))
Raymond Hettinger30729212009-02-12 06:28:27 +0000629
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000630 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000631 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000632 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000633 self.assertRaises(TypeError, cycle)
634 self.assertRaises(TypeError, cycle, 5)
635 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000636
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000637 # check copy, deepcopy, pickle
638 c = cycle('abc')
639 self.assertEqual(next(c), 'a')
640 #simple copy currently not supported, because __reduce__ returns
641 #an internal iterator
642 #self.assertEqual(take(10, copy.copy(c)), list('bcabcabcab'))
643 self.assertEqual(take(10, copy.deepcopy(c)), list('bcabcabcab'))
Serhiy Storchakabad12572014-12-15 14:03:42 +0200644 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
645 self.assertEqual(take(10, pickle.loads(pickle.dumps(c, proto))),
646 list('bcabcabcab'))
647 next(c)
648 self.assertEqual(take(10, pickle.loads(pickle.dumps(c, proto))),
649 list('cabcabcabc'))
650 next(c)
651 next(c)
652 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
653 self.pickletest(proto, cycle('abc'))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000654
Raymond Hettingera166ce52015-08-15 14:45:49 -0700655 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
656 # test with partial consumed input iterable
657 it = iter('abcde')
658 c = cycle(it)
Raymond Hettinger1cadf762015-08-15 14:47:27 -0700659 _ = [next(c) for i in range(2)] # consume 2 of 5 inputs
Raymond Hettingera166ce52015-08-15 14:45:49 -0700660 p = pickle.dumps(c, proto)
661 d = pickle.loads(p) # rebuild the cycle object
662 self.assertEqual(take(20, d), list('cdeabcdeabcdeabcdeab'))
663
664 # test with completely consumed input iterable
665 it = iter('abcde')
666 c = cycle(it)
Raymond Hettinger1cadf762015-08-15 14:47:27 -0700667 _ = [next(c) for i in range(7)] # consume 7 of 5 inputs
Raymond Hettingera166ce52015-08-15 14:45:49 -0700668 p = pickle.dumps(c, proto)
669 d = pickle.loads(p) # rebuild the cycle object
670 self.assertEqual(take(20, d), list('cdeabcdeabcdeabcdeab'))
671
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700672 def test_cycle_setstate(self):
673 # Verify both modes for restoring state
674
675 # Mode 0 is efficient. It uses an incompletely consumed input
676 # iterator to build a cycle object and then passes in state with
677 # a list of previously consumed values. There is no data
Martin Panterf1579822016-05-26 06:03:33 +0000678 # overlap between the two.
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700679 c = cycle('defg')
680 c.__setstate__((list('abc'), 0))
681 self.assertEqual(take(20, c), list('defgabcdefgabcdefgab'))
682
683 # Mode 1 is inefficient. It starts with a cycle object built
684 # from an iterator over the remaining elements in a partial
685 # cycle and then passes in state with all of the previously
686 # seen values (this overlaps values included in the iterator).
687 c = cycle('defg')
688 c.__setstate__((list('abcdefg'), 1))
689 self.assertEqual(take(20, c), list('defgabcdefgabcdefgab'))
690
691 # The first argument to setstate needs to be a tuple
Serhiy Storchaka8f0f2052016-10-02 09:13:14 +0300692 with self.assertRaises(TypeError):
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700693 cycle('defg').__setstate__([list('abcdefg'), 0])
694
695 # The first argument in the setstate tuple must be a list
696 with self.assertRaises(TypeError):
697 c = cycle('defg')
Serhiy Storchaka8f0f2052016-10-02 09:13:14 +0300698 c.__setstate__((tuple('defg'), 0))
699 take(20, c)
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700700
Serhiy Storchaka8f0f2052016-10-02 09:13:14 +0300701 # The second argument in the setstate tuple must be an int
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700702 with self.assertRaises(TypeError):
703 cycle('defg').__setstate__((list('abcdefg'), 'x'))
704
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300705 self.assertRaises(TypeError, cycle('').__setstate__, ())
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300706 self.assertRaises(TypeError, cycle('').__setstate__, ([],))
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300707
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000708 def test_groupby(self):
709 # Check whether it accepts arguments correctly
710 self.assertEqual([], list(groupby([])))
711 self.assertEqual([], list(groupby([], key=id)))
712 self.assertRaises(TypeError, list, groupby('abc', []))
713 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000714 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000715
716 # Check normal input
717 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
718 (2,15,22), (3,16,23), (3,17,23)]
719 dup = []
720 for k, g in groupby(s, lambda r:r[0]):
721 for elem in g:
722 self.assertEqual(k, elem[0])
723 dup.append(elem)
724 self.assertEqual(s, dup)
725
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000726 # Check normal pickled
Serhiy Storchakabad12572014-12-15 14:03:42 +0200727 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
728 dup = []
729 for k, g in pickle.loads(pickle.dumps(groupby(s, testR), proto)):
730 for elem in g:
731 self.assertEqual(k, elem[0])
732 dup.append(elem)
733 self.assertEqual(s, dup)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000734
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000735 # Check nested case
736 dup = []
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000737 for k, g in groupby(s, testR):
738 for ik, ig in groupby(g, testR2):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000739 for elem in ig:
740 self.assertEqual(k, elem[0])
741 self.assertEqual(ik, elem[2])
742 dup.append(elem)
743 self.assertEqual(s, dup)
744
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000745 # Check nested and pickled
Serhiy Storchakabad12572014-12-15 14:03:42 +0200746 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
747 dup = []
748 for k, g in pickle.loads(pickle.dumps(groupby(s, testR), proto)):
749 for ik, ig in pickle.loads(pickle.dumps(groupby(g, testR2), proto)):
750 for elem in ig:
751 self.assertEqual(k, elem[0])
752 self.assertEqual(ik, elem[2])
753 dup.append(elem)
754 self.assertEqual(s, dup)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000755
756
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000757 # Check case where inner iterator is not used
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000758 keys = [k for k, g in groupby(s, testR)]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000759 expectedkeys = set([r[0] for r in s])
760 self.assertEqual(set(keys), expectedkeys)
761 self.assertEqual(len(keys), len(expectedkeys))
762
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300763 # Check case where inner iterator is used after advancing the groupby
764 # iterator
765 s = list(zip('AABBBAAAA', range(9)))
766 it = groupby(s, testR)
767 _, g1 = next(it)
768 _, g2 = next(it)
769 _, g3 = next(it)
770 self.assertEqual(list(g1), [])
771 self.assertEqual(list(g2), [])
772 self.assertEqual(next(g3), ('A', 5))
773 list(it) # exhaust the groupby iterator
774 self.assertEqual(list(g3), [])
775
776 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
777 it = groupby(s, testR)
778 _, g = next(it)
779 next(it)
780 next(it)
781 self.assertEqual(list(pickle.loads(pickle.dumps(g, proto))), [])
782
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000783 # Exercise pipes and filters style
784 s = 'abracadabra'
785 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000786 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000787 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
788 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000789 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000790 self.assertEqual(r, ['a', 'b', 'r'])
791 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000792 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000793 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
794 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000795 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000796 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
797
Georg Brandla18af4e2007-04-21 15:47:16 +0000798 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000799 class ExpectedError(Exception):
800 pass
801 def delayed_raise(n=0):
802 for i in range(n):
803 yield 'yo'
804 raise ExpectedError
805 def gulp(iterable, keyp=None, func=list):
806 return [func(g) for k, g in groupby(iterable, keyp)]
807
Georg Brandla18af4e2007-04-21 15:47:16 +0000808 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000809 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000810 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000811 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
812
Serhiy Storchakaa60c2fe2015-03-12 21:56:08 +0200813 # __eq__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000814 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000815 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000816 raise ExpectedError
817 s = [DummyCmp(), DummyCmp(), None]
818
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000819 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000820 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000821 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000822 self.assertRaises(ExpectedError, gulp, s)
823
824 # keyfunc failure
825 def keyfunc(obj):
826 if keyfunc.skip > 0:
827 keyfunc.skip -= 1
828 return obj
829 else:
830 raise ExpectedError
831
832 # keyfunc failure on outer object
833 keyfunc.skip = 0
834 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
835 keyfunc.skip = 1
836 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
837
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000838 def test_filter(self):
839 self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
840 self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])
841 self.assertEqual(list(filter(bool, [0,1,0,2,0])), [1,2])
842 self.assertEqual(take(4, filter(isEven, count())), [0,2,4,6])
843 self.assertRaises(TypeError, filter)
844 self.assertRaises(TypeError, filter, lambda x:x)
845 self.assertRaises(TypeError, filter, lambda x:x, range(6), 7)
846 self.assertRaises(TypeError, filter, isEven, 3)
847 self.assertRaises(TypeError, next, filter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000848
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000849 # check copy, deepcopy, pickle
850 ans = [0,2,4]
851
852 c = filter(isEven, range(6))
853 self.assertEqual(list(copy.copy(c)), ans)
854 c = filter(isEven, range(6))
855 self.assertEqual(list(copy.deepcopy(c)), ans)
Serhiy Storchakabad12572014-12-15 14:03:42 +0200856 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
857 c = filter(isEven, range(6))
858 self.assertEqual(list(pickle.loads(pickle.dumps(c, proto))), ans)
859 next(c)
860 self.assertEqual(list(pickle.loads(pickle.dumps(c, proto))), ans[1:])
861 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
862 c = filter(isEven, range(6))
863 self.pickletest(proto, c)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000864
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000865 def test_filterfalse(self):
866 self.assertEqual(list(filterfalse(isEven, range(6))), [1,3,5])
867 self.assertEqual(list(filterfalse(None, [0,1,0,2,0])), [0,0,0])
868 self.assertEqual(list(filterfalse(bool, [0,1,0,2,0])), [0,0,0])
869 self.assertEqual(take(4, filterfalse(isEven, count())), [1,3,5,7])
870 self.assertRaises(TypeError, filterfalse)
871 self.assertRaises(TypeError, filterfalse, lambda x:x)
872 self.assertRaises(TypeError, filterfalse, lambda x:x, range(6), 7)
873 self.assertRaises(TypeError, filterfalse, isEven, 3)
874 self.assertRaises(TypeError, next, filterfalse(range(6), range(6)))
Serhiy Storchakabad12572014-12-15 14:03:42 +0200875 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
876 self.pickletest(proto, filterfalse(isEven, range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000877
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000878 def test_zip(self):
879 # XXX This is rather silly now that builtin zip() calls zip()...
880 ans = [(x,y) for x, y in zip('abc',count())]
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000881 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000882 self.assertEqual(list(zip('abc', range(6))), lzip('abc', range(6)))
883 self.assertEqual(list(zip('abcdef', range(3))), lzip('abcdef', range(3)))
884 self.assertEqual(take(3,zip('abcdef', count())), lzip('abcdef', range(3)))
885 self.assertEqual(list(zip('abcdef')), lzip('abcdef'))
886 self.assertEqual(list(zip()), lzip())
887 self.assertRaises(TypeError, zip, 3)
888 self.assertRaises(TypeError, zip, range(3), 3)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000889 self.assertEqual([tuple(list(pair)) for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000890 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000891 self.assertEqual([pair for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000892 lzip('abc', 'def'))
Alex Gaynore151d212011-07-17 16:21:30 -0700893
894 @support.impl_detail("tuple reuse is specific to CPython")
895 def test_zip_tuple_reuse(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000896 ids = list(map(id, zip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000897 self.assertEqual(min(ids), max(ids))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000898 ids = list(map(id, list(zip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000899 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000900
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000901 # check copy, deepcopy, pickle
902 ans = [(x,y) for x, y in copy.copy(zip('abc',count()))]
903 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
904
905 ans = [(x,y) for x, y in copy.deepcopy(zip('abc',count()))]
906 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
907
Serhiy Storchakabad12572014-12-15 14:03:42 +0200908 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
909 ans = [(x,y) for x, y in pickle.loads(pickle.dumps(zip('abc',count()), proto))]
910 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000911
Serhiy Storchakabad12572014-12-15 14:03:42 +0200912 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
913 testIntermediate = zip('abc',count())
914 next(testIntermediate)
915 ans = [(x,y) for x, y in pickle.loads(pickle.dumps(testIntermediate, proto))]
916 self.assertEqual(ans, [('b', 1), ('c', 2)])
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000917
Serhiy Storchakabad12572014-12-15 14:03:42 +0200918 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
919 self.pickletest(proto, zip('abc', count()))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000920
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000921 def test_ziplongest(self):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000922 for args in [
923 ['abc', range(6)],
924 [range(6), 'abc'],
925 [range(1000), range(2000,2100), range(3000,3050)],
926 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
927 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
928 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000929 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
930 for i in range(max(map(len, args)))]
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000931 self.assertEqual(list(zip_longest(*args)), target)
932 self.assertEqual(list(zip_longest(*args, **{})), target)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000933 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 +0000934 self.assertEqual(list(zip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000935
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000936 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 +0000937
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000938 self.assertEqual(list(zip_longest()), list(zip()))
939 self.assertEqual(list(zip_longest([])), list(zip([])))
940 self.assertEqual(list(zip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000941
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000942 self.assertEqual(list(zip_longest('abc', 'defg', **{})),
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000943 list(zip(list('abc')+[None], 'defg'))) # empty keyword dict
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000944 self.assertRaises(TypeError, zip_longest, 3)
945 self.assertRaises(TypeError, zip_longest, range(3), 3)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000946
947 for stmt in [
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000948 "zip_longest('abc', fv=1)",
949 "zip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000950 ]:
951 try:
952 eval(stmt, globals(), locals())
953 except TypeError:
954 pass
955 else:
956 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000957
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000958 self.assertEqual([tuple(list(pair)) for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000959 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000960 self.assertEqual([pair for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000961 list(zip('abc', 'def')))
Alex Gaynore151d212011-07-17 16:21:30 -0700962
963 @support.impl_detail("tuple reuse is specific to CPython")
964 def test_zip_longest_tuple_reuse(self):
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000965 ids = list(map(id, zip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000966 self.assertEqual(min(ids), max(ids))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000967 ids = list(map(id, list(zip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000968 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
969
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000970 def test_zip_longest_pickling(self):
Serhiy Storchakabad12572014-12-15 14:03:42 +0200971 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
972 self.pickletest(proto, zip_longest("abc", "def"))
973 self.pickletest(proto, zip_longest("abc", "defgh"))
974 self.pickletest(proto, zip_longest("abc", "defgh", fillvalue=1))
975 self.pickletest(proto, zip_longest("", "defgh"))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000976
Sergey Fedoseev6a650aa2019-08-30 09:25:48 +0500977 def test_zip_longest_bad_iterable(self):
978 exception = TypeError()
979
980 class BadIterable:
981 def __iter__(self):
982 raise exception
983
984 with self.assertRaises(TypeError) as cm:
985 zip_longest(BadIterable())
986
987 self.assertIs(cm.exception, exception)
988
Raymond Hettingerfc438512009-11-01 20:55:33 +0000989 def test_bug_7244(self):
990
991 class Repeater:
992 # this class is similar to itertools.repeat
993 def __init__(self, o, t, e):
994 self.o = o
995 self.t = int(t)
996 self.e = e
997 def __iter__(self): # its iterator is itself
998 return self
999 def __next__(self):
1000 if self.t > 0:
1001 self.t -= 1
1002 return self.o
1003 else:
1004 raise self.e
1005
1006 # Formerly this code in would fail in debug mode
1007 # with Undetected Error and Stop Iteration
1008 r1 = Repeater(1, 3, StopIteration)
1009 r2 = Repeater(2, 4, StopIteration)
1010 def run(r1, r2):
1011 result = []
1012 for i, j in zip_longest(r1, r2, fillvalue=0):
1013 with support.captured_output('stdout'):
1014 print((i, j))
1015 result.append((i, j))
1016 return result
1017 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
1018
1019 # Formerly, the RuntimeError would be lost
1020 # and StopIteration would stop as expected
1021 r1 = Repeater(1, 3, RuntimeError)
1022 r2 = Repeater(2, 4, StopIteration)
1023 it = zip_longest(r1, r2, fillvalue=0)
1024 self.assertEqual(next(it), (1, 2))
1025 self.assertEqual(next(it), (1, 2))
1026 self.assertEqual(next(it), (1, 2))
1027 self.assertRaises(RuntimeError, next, it)
1028
Raymond Hettingercc061d02020-11-30 20:42:54 -08001029 def test_pairwise(self):
1030 self.assertEqual(list(pairwise('')), [])
1031 self.assertEqual(list(pairwise('a')), [])
1032 self.assertEqual(list(pairwise('ab')),
1033 [('a', 'b')]),
1034 self.assertEqual(list(pairwise('abcde')),
1035 [('a', 'b'), ('b', 'c'), ('c', 'd'), ('d', 'e')])
1036 self.assertEqual(list(pairwise(range(10_000))),
1037 list(zip(range(10_000), range(1, 10_000))))
1038
1039 with self.assertRaises(TypeError):
1040 pairwise() # too few arguments
1041 with self.assertRaises(TypeError):
1042 pairwise('abc', 10) # too many arguments
1043 with self.assertRaises(TypeError):
1044 pairwise(iterable='abc') # keyword arguments
1045 with self.assertRaises(TypeError):
1046 pairwise(None) # non-iterable argument
1047
Christian Heimesc3f30c42008-02-22 16:37:40 +00001048 def test_product(self):
1049 for args, result in [
Christian Heimes78644762008-03-04 23:39:23 +00001050 ([], [()]), # zero iterables
Christian Heimesc3f30c42008-02-22 16:37:40 +00001051 (['ab'], [('a',), ('b',)]), # one iterable
1052 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
1053 ([range(0), range(2), range(3)], []), # first iterable with zero length
1054 ([range(2), range(0), range(3)], []), # middle iterable with zero length
1055 ([range(2), range(3), range(0)], []), # last iterable with zero length
1056 ]:
1057 self.assertEqual(list(product(*args)), result)
Christian Heimesf16baeb2008-02-29 14:57:44 +00001058 for r in range(4):
1059 self.assertEqual(list(product(*(args*r))),
1060 list(product(*args, **dict(repeat=r))))
Christian Heimesc3f30c42008-02-22 16:37:40 +00001061 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
1062 self.assertRaises(TypeError, product, range(6), None)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001063
1064 def product1(*args, **kwds):
1065 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
1066 n = len(pools)
1067 if n == 0:
1068 yield ()
1069 return
1070 if any(len(pool) == 0 for pool in pools):
1071 return
1072 indices = [0] * n
1073 yield tuple(pool[i] for pool, i in zip(pools, indices))
1074 while 1:
1075 for i in reversed(range(n)): # right to left
1076 if indices[i] == len(pools[i]) - 1:
1077 continue
1078 indices[i] += 1
1079 for j in range(i+1, n):
1080 indices[j] = 0
1081 yield tuple(pool[i] for pool, i in zip(pools, indices))
1082 break
1083 else:
1084 return
1085
1086 def product2(*args, **kwds):
1087 'Pure python version used in docs'
1088 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
1089 result = [[]]
1090 for pool in pools:
1091 result = [x+[y] for x in result for y in pool]
1092 for prod in result:
1093 yield tuple(prod)
1094
Christian Heimesc3f30c42008-02-22 16:37:40 +00001095 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
1096 set('abcdefg'), range(11), tuple(range(13))]
1097 for i in range(100):
1098 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Christian Heimes78644762008-03-04 23:39:23 +00001099 expected_len = prod(map(len, args))
1100 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001101 self.assertEqual(list(product(*args)), list(product1(*args)))
1102 self.assertEqual(list(product(*args)), list(product2(*args)))
Christian Heimesc3f30c42008-02-22 16:37:40 +00001103 args = map(iter, args)
Christian Heimes78644762008-03-04 23:39:23 +00001104 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesc3f30c42008-02-22 16:37:40 +00001105
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05001106 @support.bigaddrspacetest
1107 def test_product_overflow(self):
Serhiy Storchakadee948b2015-02-03 01:34:09 +02001108 with self.assertRaises((OverflowError, MemoryError)):
1109 product(*(['ab']*2**5), repeat=2**25)
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05001110
Alex Gaynore151d212011-07-17 16:21:30 -07001111 @support.impl_detail("tuple reuse is specific to CPython")
1112 def test_product_tuple_reuse(self):
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001113 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
1114 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +00001115
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001116 def test_product_pickling(self):
1117 # check copy, deepcopy, pickle
1118 for args, result in [
1119 ([], [()]), # zero iterables
1120 (['ab'], [('a',), ('b',)]), # one iterable
1121 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
1122 ([range(0), range(2), range(3)], []), # first iterable with zero length
1123 ([range(2), range(0), range(3)], []), # middle iterable with zero length
1124 ([range(2), range(3), range(0)], []), # last iterable with zero length
1125 ]:
1126 self.assertEqual(list(copy.copy(product(*args))), result)
1127 self.assertEqual(list(copy.deepcopy(product(*args))), result)
Serhiy Storchakabad12572014-12-15 14:03:42 +02001128 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1129 self.pickletest(proto, product(*args))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001130
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00001131 def test_product_issue_25021(self):
1132 # test that indices are properly clamped to the length of the tuples
1133 p = product((1, 2),(3,))
1134 p.__setstate__((0, 0x1000)) # will access tuple element 1 if not clamped
1135 self.assertEqual(next(p), (2, 3))
1136 # test that empty tuple in the list will result in an immediate StopIteration
1137 p = product((1, 2), (), (3,))
1138 p.__setstate__((0, 0, 0x1000)) # will access tuple element 1 if not clamped
1139 self.assertRaises(StopIteration, next, p)
1140
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001141 def test_repeat(self):
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00001142 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
Guido van Rossum805365e2007-05-07 22:24:25 +00001143 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001144 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001145 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +00001146 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001147 self.assertEqual(list(repeat('a', 0)), [])
1148 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001149 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001150 self.assertRaises(TypeError, repeat, None, 3, 4)
1151 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001152 r = repeat(1+0j)
1153 self.assertEqual(repr(r), 'repeat((1+0j))')
1154 r = repeat(1+0j, 5)
1155 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
1156 list(r)
1157 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001158
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001159 # check copy, deepcopy, pickle
1160 c = repeat(object='a', times=10)
1161 self.assertEqual(next(c), 'a')
1162 self.assertEqual(take(2, copy.copy(c)), list('a' * 2))
1163 self.assertEqual(take(2, copy.deepcopy(c)), list('a' * 2))
Serhiy Storchakabad12572014-12-15 14:03:42 +02001164 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1165 self.pickletest(proto, repeat(object='a', times=10))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001166
Raymond Hettinger97d35552014-06-24 21:36:58 -07001167 def test_repeat_with_negative_times(self):
1168 self.assertEqual(repr(repeat('a', -1)), "repeat('a', 0)")
1169 self.assertEqual(repr(repeat('a', -2)), "repeat('a', 0)")
1170 self.assertEqual(repr(repeat('a', times=-1)), "repeat('a', 0)")
1171 self.assertEqual(repr(repeat('a', times=-2)), "repeat('a', 0)")
1172
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001173 def test_map(self):
1174 self.assertEqual(list(map(operator.pow, range(3), range(1,7))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001175 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001176 self.assertEqual(list(map(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001177 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001178 self.assertEqual(list(map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +00001179 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001180 self.assertEqual(take(2,map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +00001181 [('a',0),('b',1)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001182 self.assertEqual(list(map(operator.pow, [])), [])
1183 self.assertRaises(TypeError, map)
1184 self.assertRaises(TypeError, list, map(None, range(3), range(3)))
1185 self.assertRaises(TypeError, map, operator.neg)
1186 self.assertRaises(TypeError, next, map(10, range(5)))
1187 self.assertRaises(ValueError, next, map(errfunc, [4], [5]))
1188 self.assertRaises(TypeError, next, map(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001189
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001190 # check copy, deepcopy, pickle
1191 ans = [('a',0),('b',1),('c',2)]
1192
1193 c = map(tupleize, 'abc', count())
1194 self.assertEqual(list(copy.copy(c)), ans)
1195
1196 c = map(tupleize, 'abc', count())
1197 self.assertEqual(list(copy.deepcopy(c)), ans)
1198
Serhiy Storchakabad12572014-12-15 14:03:42 +02001199 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1200 c = map(tupleize, 'abc', count())
1201 self.pickletest(proto, c)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001202
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001203 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001204 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
1205 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001206 self.assertEqual(take(3, starmap(operator.pow, zip(count(), count(1)))),
Raymond Hettinger02420702003-06-29 20:36:23 +00001207 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001208 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +00001209 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
1210 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001211 self.assertRaises(TypeError, starmap)
1212 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001213 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
1214 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
1215 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001216
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001217 # check copy, deepcopy, pickle
1218 ans = [0**1, 1**2, 2**3]
1219
1220 c = starmap(operator.pow, zip(range(3), range(1,7)))
1221 self.assertEqual(list(copy.copy(c)), ans)
1222
1223 c = starmap(operator.pow, zip(range(3), range(1,7)))
1224 self.assertEqual(list(copy.deepcopy(c)), ans)
1225
Serhiy Storchakabad12572014-12-15 14:03:42 +02001226 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1227 c = starmap(operator.pow, zip(range(3), range(1,7)))
1228 self.pickletest(proto, c)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001229
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001230 def test_islice(self):
1231 for args in [ # islice(args) should agree with range(args)
1232 (10, 20, 3),
1233 (10, 3, 20),
1234 (10, 20),
Cheryl Sabellada1734c2018-03-26 21:29:33 -04001235 (10, 10),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001236 (10, 3),
1237 (20,)
1238 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +00001239 self.assertEqual(list(islice(range(100), *args)),
1240 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001241
1242 for args, tgtargs in [ # Stop when seqn is exhausted
1243 ((10, 110, 3), ((10, 100, 3))),
1244 ((10, 110), ((10, 100))),
1245 ((110,), (100,))
1246 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +00001247 self.assertEqual(list(islice(range(100), *args)),
1248 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001249
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001250 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +00001251 self.assertEqual(list(islice(range(10), None)), list(range(10)))
1252 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
1253 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
1254 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
1255 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001256
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +00001257 # Test number of items consumed SF #1171417
1258 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +00001259 self.assertEqual(list(islice(it, 3)), list(range(3)))
1260 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +00001261
Cheryl Sabellada1734c2018-03-26 21:29:33 -04001262 it = iter(range(10))
1263 self.assertEqual(list(islice(it, 3, 3)), [])
1264 self.assertEqual(list(it), list(range(3, 10)))
1265
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001266 # Test invalid arguments
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001267 ra = range(10)
1268 self.assertRaises(TypeError, islice, ra)
1269 self.assertRaises(TypeError, islice, ra, 1, 2, 3, 4)
1270 self.assertRaises(ValueError, islice, ra, -5, 10, 1)
1271 self.assertRaises(ValueError, islice, ra, 1, -5, -1)
1272 self.assertRaises(ValueError, islice, ra, 1, 10, -1)
1273 self.assertRaises(ValueError, islice, ra, 1, 10, 0)
1274 self.assertRaises(ValueError, islice, ra, 'a')
1275 self.assertRaises(ValueError, islice, ra, 'a', 1)
1276 self.assertRaises(ValueError, islice, ra, 1, 'a')
1277 self.assertRaises(ValueError, islice, ra, 'a', 1, 1)
1278 self.assertRaises(ValueError, islice, ra, 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +00001279 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001280
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001281 # Issue #10323: Less islice in a predictable state
1282 c = count()
1283 self.assertEqual(list(islice(c, 1, 3, 50)), [1])
1284 self.assertEqual(next(c), 3)
1285
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001286 # check copy, deepcopy, pickle
1287 for args in [ # islice(args) should agree with range(args)
1288 (10, 20, 3),
1289 (10, 3, 20),
1290 (10, 20),
1291 (10, 3),
1292 (20,)
1293 ]:
1294 self.assertEqual(list(copy.copy(islice(range(100), *args))),
1295 list(range(*args)))
1296 self.assertEqual(list(copy.deepcopy(islice(range(100), *args))),
1297 list(range(*args)))
Serhiy Storchakabad12572014-12-15 14:03:42 +02001298 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1299 self.pickletest(proto, islice(range(100), *args))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001300
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001301 # Issue #21321: check source iterator is not referenced
1302 # from islice() after the latter has been exhausted
1303 it = (x for x in (1, 2))
1304 wr = weakref.ref(it)
1305 it = islice(it, 1)
1306 self.assertIsNotNone(wr())
1307 list(it) # exhaust the iterator
Benjamin Peterson8e163512014-08-24 18:07:28 -05001308 support.gc_collect()
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001309 self.assertIsNone(wr())
1310
Will Roberts0ecdc522017-06-08 08:03:04 +02001311 # Issue #30537: islice can accept integer-like objects as
1312 # arguments
1313 class IntLike(object):
1314 def __init__(self, val):
1315 self.val = val
1316 def __index__(self):
1317 return self.val
1318 self.assertEqual(list(islice(range(100), IntLike(10))), list(range(10)))
1319 self.assertEqual(list(islice(range(100), IntLike(10), IntLike(50))),
1320 list(range(10, 50)))
1321 self.assertEqual(list(islice(range(100), IntLike(10), IntLike(50), IntLike(5))),
1322 list(range(10,50,5)))
1323
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001324 def test_takewhile(self):
1325 data = [1, 3, 5, 20, 2, 4, 6, 8]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001326 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001327 self.assertEqual(list(takewhile(underten, [])), [])
1328 self.assertRaises(TypeError, takewhile)
1329 self.assertRaises(TypeError, takewhile, operator.pow)
1330 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001331 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
1332 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001333 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
1334 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +00001335 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001336
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001337 # check copy, deepcopy, pickle
1338 self.assertEqual(list(copy.copy(takewhile(underten, data))), [1, 3, 5])
1339 self.assertEqual(list(copy.deepcopy(takewhile(underten, data))),
1340 [1, 3, 5])
Serhiy Storchakabad12572014-12-15 14:03:42 +02001341 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1342 self.pickletest(proto, takewhile(underten, data))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001343
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001344 def test_dropwhile(self):
1345 data = [1, 3, 5, 20, 2, 4, 6, 8]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001346 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001347 self.assertEqual(list(dropwhile(underten, [])), [])
1348 self.assertRaises(TypeError, dropwhile)
1349 self.assertRaises(TypeError, dropwhile, operator.pow)
1350 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001351 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
1352 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001353
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001354 # check copy, deepcopy, pickle
1355 self.assertEqual(list(copy.copy(dropwhile(underten, data))), [20, 2, 4, 6, 8])
1356 self.assertEqual(list(copy.deepcopy(dropwhile(underten, data))),
1357 [20, 2, 4, 6, 8])
Serhiy Storchakabad12572014-12-15 14:03:42 +02001358 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1359 self.pickletest(proto, dropwhile(underten, data))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001360
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001361 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +00001362 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001363
1364 a, b = tee([]) # test empty iterator
1365 self.assertEqual(list(a), [])
1366 self.assertEqual(list(b), [])
1367
1368 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +00001369 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001370
1371 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +00001372 self.assertEqual(list(a), list(range(n)))
1373 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001374
1375 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +00001376 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001377 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001378 del a
Guido van Rossum805365e2007-05-07 22:24:25 +00001379 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001380
1381 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +00001382 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001383 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001384 del b
Guido van Rossum805365e2007-05-07 22:24:25 +00001385 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001386
Guido van Rossum805365e2007-05-07 22:24:25 +00001387 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001388 order = [0]*n + [1]*n
1389 random.shuffle(order)
1390 lists = ([], [])
1391 its = tee(irange(n))
1392 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +00001393 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001394 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +00001395 self.assertEqual(lists[0], list(range(n)))
1396 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001397
Raymond Hettingerad983e72003-11-12 14:32:26 +00001398 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001399 self.assertRaises(TypeError, tee)
1400 self.assertRaises(TypeError, tee, 3)
1401 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +00001402 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001403
Raymond Hettingerad983e72003-11-12 14:32:26 +00001404 # tee object should be instantiable
1405 a, b = tee('abc')
1406 c = type(a)('def')
1407 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +00001408
Raymond Hettingerad983e72003-11-12 14:32:26 +00001409 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +00001410 a, b, c = tee(range(2000), 3)
1411 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001412 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +00001413 self.assertEqual(list(b), list(range(2000)))
1414 self.assertEqual([next(c), next(c)], list(range(2)))
1415 self.assertEqual(list(a), list(range(100,2000)))
1416 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +00001417
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001418 # test values of n
1419 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +00001420 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +00001421 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001422 result = tee('abc', n)
1423 self.assertEqual(type(result), tuple)
1424 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001425 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001426
Raymond Hettingerad983e72003-11-12 14:32:26 +00001427 # tee pass-through to copyable iterator
1428 a, b = tee('abc')
1429 c, d = tee(a)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +00001430 self.assertTrue(a is c)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001431
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001432 # test tee_new
1433 t1, t2 = tee('abc')
1434 tnew = type(t1)
1435 self.assertRaises(TypeError, tnew)
1436 self.assertRaises(TypeError, tnew, 10)
1437 t3 = tnew(t1)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +00001438 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +00001439
Raymond Hettingera9f60922004-10-17 16:40:14 +00001440 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +00001441 a, b = tee(range(10))
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001442 p = weakref.proxy(a)
Raymond Hettingera9f60922004-10-17 16:40:14 +00001443 self.assertEqual(getattr(p, '__class__'), type(b))
1444 del a
1445 self.assertRaises(ReferenceError, getattr, p, '__class__')
1446
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001447 ans = list('abc')
1448 long_ans = list(range(10000))
1449
1450 # check copy
1451 a, b = tee('abc')
1452 self.assertEqual(list(copy.copy(a)), ans)
1453 self.assertEqual(list(copy.copy(b)), ans)
1454 a, b = tee(list(range(10000)))
1455 self.assertEqual(list(copy.copy(a)), long_ans)
1456 self.assertEqual(list(copy.copy(b)), long_ans)
1457
1458 # check partially consumed copy
1459 a, b = tee('abc')
1460 take(2, a)
1461 take(1, b)
1462 self.assertEqual(list(copy.copy(a)), ans[2:])
1463 self.assertEqual(list(copy.copy(b)), ans[1:])
1464 self.assertEqual(list(a), ans[2:])
1465 self.assertEqual(list(b), ans[1:])
1466 a, b = tee(range(10000))
1467 take(100, a)
1468 take(60, b)
1469 self.assertEqual(list(copy.copy(a)), long_ans[100:])
1470 self.assertEqual(list(copy.copy(b)), long_ans[60:])
1471 self.assertEqual(list(a), long_ans[100:])
1472 self.assertEqual(list(b), long_ans[60:])
1473
1474 # check deepcopy
1475 a, b = tee('abc')
1476 self.assertEqual(list(copy.deepcopy(a)), ans)
1477 self.assertEqual(list(copy.deepcopy(b)), ans)
1478 self.assertEqual(list(a), ans)
1479 self.assertEqual(list(b), ans)
1480 a, b = tee(range(10000))
1481 self.assertEqual(list(copy.deepcopy(a)), long_ans)
1482 self.assertEqual(list(copy.deepcopy(b)), long_ans)
1483 self.assertEqual(list(a), long_ans)
1484 self.assertEqual(list(b), long_ans)
1485
1486 # check partially consumed deepcopy
1487 a, b = tee('abc')
1488 take(2, a)
1489 take(1, b)
1490 self.assertEqual(list(copy.deepcopy(a)), ans[2:])
1491 self.assertEqual(list(copy.deepcopy(b)), ans[1:])
1492 self.assertEqual(list(a), ans[2:])
1493 self.assertEqual(list(b), ans[1:])
1494 a, b = tee(range(10000))
1495 take(100, a)
1496 take(60, b)
1497 self.assertEqual(list(copy.deepcopy(a)), long_ans[100:])
1498 self.assertEqual(list(copy.deepcopy(b)), long_ans[60:])
1499 self.assertEqual(list(a), long_ans[100:])
1500 self.assertEqual(list(b), long_ans[60:])
1501
1502 # check pickle
Serhiy Storchakabad12572014-12-15 14:03:42 +02001503 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1504 self.pickletest(proto, iter(tee('abc')))
1505 a, b = tee('abc')
1506 self.pickletest(proto, a, compare=ans)
1507 self.pickletest(proto, b, compare=ans)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001508
Serhiy Storchakaa3e91282013-01-25 13:19:31 +02001509 # Issue 13454: Crash when deleting backward iterator from tee()
1510 def test_tee_del_backward(self):
Serhiy Storchaka5bb893c2013-01-26 11:52:06 +02001511 forward, backward = tee(repeat(None, 20000000))
Serhiy Storchaka9db55002015-03-28 20:38:37 +02001512 try:
1513 any(forward) # exhaust the iterator
1514 del backward
1515 except:
1516 del forward, backward
1517 raise
Serhiy Storchakaa3e91282013-01-25 13:19:31 +02001518
Serhiy Storchaka526a0142019-09-09 11:47:14 +03001519 def test_tee_reenter(self):
1520 class I:
1521 first = True
1522 def __iter__(self):
1523 return self
1524 def __next__(self):
1525 first = self.first
1526 self.first = False
1527 if first:
1528 return next(b)
1529
1530 a, b = tee(I())
1531 with self.assertRaisesRegex(RuntimeError, "tee"):
1532 next(a)
1533
1534 def test_tee_concurrent(self):
1535 start = threading.Event()
1536 finish = threading.Event()
1537 class I:
1538 def __iter__(self):
1539 return self
1540 def __next__(self):
1541 start.set()
1542 finish.wait()
1543
1544 a, b = tee(I())
1545 thread = threading.Thread(target=next, args=[a])
1546 thread.start()
1547 try:
1548 start.wait()
1549 with self.assertRaisesRegex(RuntimeError, "tee"):
1550 next(b)
1551 finally:
1552 finish.set()
1553 thread.join()
1554
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001555 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001556 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +00001557
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001558 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +00001559 self.assertRaises(StopIteration, next, f([]))
1560 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001561
Georg Brandla18af4e2007-04-21 15:47:16 +00001562 self.assertRaises(StopIteration, next, islice([], None))
1563 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001564
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001565 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +00001566 self.assertRaises(StopIteration, next, p)
1567 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001568 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +00001569 self.assertRaises(StopIteration, next, p)
1570 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001571
Georg Brandla18af4e2007-04-21 15:47:16 +00001572 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001573
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001574 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +00001575 self.assertRaises(StopIteration, next, f(lambda x:x, []))
1576 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001577
Brandt Bucher226a0122020-12-04 19:45:57 -08001578 @support.cpython_only
1579 def test_combinations_result_gc(self):
1580 # bpo-42536: combinations's tuple-reuse speed trick breaks the GC's
1581 # assumptions about what can be untracked. Make sure we re-track result
1582 # tuples whenever we reuse them.
1583 it = combinations([None, []], 1)
1584 next(it)
1585 gc.collect()
1586 # That GC collection probably untracked the recycled internal result
1587 # tuple, which has the value (None,). Make sure it's re-tracked when
1588 # it's mutated and returned from __next__:
1589 self.assertTrue(gc.is_tracked(next(it)))
1590
1591 @support.cpython_only
1592 def test_combinations_with_replacement_result_gc(self):
1593 # Ditto for combinations_with_replacement.
1594 it = combinations_with_replacement([None, []], 1)
1595 next(it)
1596 gc.collect()
1597 self.assertTrue(gc.is_tracked(next(it)))
1598
1599 @support.cpython_only
1600 def test_permutations_result_gc(self):
1601 # Ditto for permutations.
1602 it = permutations([None, []], 1)
1603 next(it)
1604 gc.collect()
1605 self.assertTrue(gc.is_tracked(next(it)))
1606
1607 @support.cpython_only
1608 def test_product_result_gc(self):
1609 # Ditto for product.
1610 it = product([None, []])
1611 next(it)
1612 gc.collect()
1613 self.assertTrue(gc.is_tracked(next(it)))
1614
1615 @support.cpython_only
1616 def test_zip_longest_result_gc(self):
1617 # Ditto for zip_longest.
1618 it = zip_longest([[]])
1619 gc.collect()
1620 self.assertTrue(gc.is_tracked(next(it)))
1621
1622
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001623class TestExamples(unittest.TestCase):
1624
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001625 def test_accumulate(self):
Raymond Hettinger482ba772010-12-01 22:48:00 +00001626 self.assertEqual(list(accumulate([1,2,3,4,5])), [1, 3, 6, 10, 15])
1627
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001628 def test_accumulate_reducible(self):
1629 # check copy, deepcopy, pickle
1630 data = [1, 2, 3, 4, 5]
1631 accumulated = [1, 3, 6, 10, 15]
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001632
Serhiy Storchakabad12572014-12-15 14:03:42 +02001633 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1634 it = accumulate(data)
1635 self.assertEqual(list(pickle.loads(pickle.dumps(it, proto))), accumulated[:])
1636 self.assertEqual(next(it), 1)
1637 self.assertEqual(list(pickle.loads(pickle.dumps(it, proto))), accumulated[1:])
1638 it = accumulate(data)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001639 self.assertEqual(next(it), 1)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001640 self.assertEqual(list(copy.deepcopy(it)), accumulated[1:])
1641 self.assertEqual(list(copy.copy(it)), accumulated[1:])
1642
Serhiy Storchakad5516252016-03-06 14:00:45 +02001643 def test_accumulate_reducible_none(self):
1644 # Issue #25718: total is None
1645 it = accumulate([None, None, None], operator.is_)
1646 self.assertEqual(next(it), None)
1647 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1648 it_copy = pickle.loads(pickle.dumps(it, proto))
1649 self.assertEqual(list(it_copy), [True, False])
1650 self.assertEqual(list(copy.deepcopy(it)), [True, False])
1651 self.assertEqual(list(copy.copy(it)), [True, False])
1652
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001653 def test_chain(self):
1654 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
1655
1656 def test_chain_from_iterable(self):
1657 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
1658
1659 def test_combinations(self):
1660 self.assertEqual(list(combinations('ABCD', 2)),
1661 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
1662 self.assertEqual(list(combinations(range(4), 3)),
1663 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
1664
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001665 def test_combinations_with_replacement(self):
1666 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
1667 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
1668
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001669 def test_compress(self):
1670 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
1671
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001672 def test_count(self):
1673 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
1674
1675 def test_cycle(self):
1676 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
1677
1678 def test_dropwhile(self):
1679 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
1680
1681 def test_groupby(self):
1682 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
1683 list('ABCDAB'))
1684 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
1685 [list('AAAA'), list('BBB'), list('CC'), list('D')])
1686
1687 def test_filter(self):
1688 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
1689
1690 def test_filterfalse(self):
1691 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
1692
1693 def test_map(self):
1694 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
1695
1696 def test_islice(self):
1697 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
1698 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
1699 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
1700 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1701
1702 def test_zip(self):
1703 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
1704
1705 def test_zip_longest(self):
1706 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
1707 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
1708
1709 def test_permutations(self):
1710 self.assertEqual(list(permutations('ABCD', 2)),
1711 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
1712 self.assertEqual(list(permutations(range(3))),
1713 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
1714
1715 def test_product(self):
1716 self.assertEqual(list(product('ABCD', 'xy')),
1717 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
1718 self.assertEqual(list(product(range(2), repeat=3)),
1719 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
1720 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
1721
1722 def test_repeat(self):
1723 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
1724
1725 def test_stapmap(self):
1726 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
1727 [32, 9, 1000])
1728
1729 def test_takewhile(self):
1730 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
1731
1732
Cheryl Sabellada1734c2018-03-26 21:29:33 -04001733class TestPurePythonRoughEquivalents(unittest.TestCase):
1734
1735 @staticmethod
1736 def islice(iterable, *args):
1737 s = slice(*args)
1738 start, stop, step = s.start or 0, s.stop or sys.maxsize, s.step or 1
1739 it = iter(range(start, stop, step))
1740 try:
1741 nexti = next(it)
1742 except StopIteration:
1743 # Consume *iterable* up to the *start* position.
1744 for i, element in zip(range(start), iterable):
1745 pass
1746 return
1747 try:
1748 for i, element in enumerate(iterable):
1749 if i == nexti:
1750 yield element
1751 nexti = next(it)
1752 except StopIteration:
1753 # Consume to *stop*.
1754 for i, element in zip(range(i + 1, stop), iterable):
1755 pass
1756
1757 def test_islice_recipe(self):
1758 self.assertEqual(list(self.islice('ABCDEFG', 2)), list('AB'))
1759 self.assertEqual(list(self.islice('ABCDEFG', 2, 4)), list('CD'))
1760 self.assertEqual(list(self.islice('ABCDEFG', 2, None)), list('CDEFG'))
1761 self.assertEqual(list(self.islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1762 # Test items consumed.
1763 it = iter(range(10))
1764 self.assertEqual(list(self.islice(it, 3)), list(range(3)))
1765 self.assertEqual(list(it), list(range(3, 10)))
1766 it = iter(range(10))
1767 self.assertEqual(list(self.islice(it, 3, 3)), [])
1768 self.assertEqual(list(it), list(range(3, 10)))
1769 # Test that slice finishes in predictable state.
1770 c = count()
1771 self.assertEqual(list(self.islice(c, 1, 3, 50)), [1])
1772 self.assertEqual(next(c), 3)
1773
1774
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001775class TestGC(unittest.TestCase):
1776
1777 def makecycle(self, iterator, container):
1778 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +00001779 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001780 del container, iterator
1781
Raymond Hettinger482ba772010-12-01 22:48:00 +00001782 def test_accumulate(self):
1783 a = []
1784 self.makecycle(accumulate([1,2,a,3]), a)
1785
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001786 def test_chain(self):
1787 a = []
1788 self.makecycle(chain(a), a)
1789
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001790 def test_chain_from_iterable(self):
1791 a = []
1792 self.makecycle(chain.from_iterable([a]), a)
1793
1794 def test_combinations(self):
1795 a = []
1796 self.makecycle(combinations([1,2,a,3], 3), a)
1797
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001798 def test_combinations_with_replacement(self):
1799 a = []
1800 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1801
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001802 def test_compress(self):
1803 a = []
1804 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1805
Raymond Hettingerd6280f42009-02-16 20:50:56 +00001806 def test_count(self):
1807 a = []
1808 Int = type('Int', (int,), dict(x=a))
1809 self.makecycle(count(Int(0), Int(1)), a)
1810
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001811 def test_cycle(self):
1812 a = []
1813 self.makecycle(cycle([a]*2), a)
1814
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001815 def test_dropwhile(self):
1816 a = []
1817 self.makecycle(dropwhile(bool, [0, a, a]), a)
1818
1819 def test_groupby(self):
1820 a = []
1821 self.makecycle(groupby([a]*2, lambda x:x), a)
1822
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001823 def test_issue2246(self):
1824 # Issue 2246 -- the _grouper iterator was not included in GC
1825 n = 10
1826 keyfunc = lambda x: x
1827 for i, j in groupby(range(n), key=keyfunc):
1828 keyfunc.__dict__.setdefault('x',[]).append(j)
1829
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001830 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001831 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001832 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001833
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001834 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001835 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001836 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001837
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001838 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001839 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001840 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001841
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001842 def test_zip_longest(self):
1843 a = []
1844 self.makecycle(zip_longest([a]*2, [a]*3), a)
1845 b = [a, None]
1846 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1847
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001848 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001849 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001850 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001851
1852 def test_islice(self):
1853 a = []
1854 self.makecycle(islice([a]*2, None), a)
1855
Raymond Hettingercc061d02020-11-30 20:42:54 -08001856 def test_pairwise(self):
1857 a = []
1858 self.makecycle(pairwise([a]*5), a)
1859
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001860 def test_permutations(self):
1861 a = []
1862 self.makecycle(permutations([1,2,a,3], 3), a)
1863
1864 def test_product(self):
1865 a = []
1866 self.makecycle(product([1,2,a,3], repeat=3), a)
1867
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001868 def test_repeat(self):
1869 a = []
1870 self.makecycle(repeat(a), a)
1871
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001872 def test_starmap(self):
1873 a = []
1874 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1875
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001876 def test_takewhile(self):
1877 a = []
1878 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1879
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001880def R(seqn):
1881 'Regular generator'
1882 for i in seqn:
1883 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001884
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001885class G:
1886 'Sequence using __getitem__'
1887 def __init__(self, seqn):
1888 self.seqn = seqn
1889 def __getitem__(self, i):
1890 return self.seqn[i]
1891
1892class I:
1893 'Sequence using iterator protocol'
1894 def __init__(self, seqn):
1895 self.seqn = seqn
1896 self.i = 0
1897 def __iter__(self):
1898 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001899 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001900 if self.i >= len(self.seqn): raise StopIteration
1901 v = self.seqn[self.i]
1902 self.i += 1
1903 return v
1904
1905class Ig:
1906 'Sequence using iterator protocol defined with a generator'
1907 def __init__(self, seqn):
1908 self.seqn = seqn
1909 self.i = 0
1910 def __iter__(self):
1911 for val in self.seqn:
1912 yield val
1913
1914class X:
1915 'Missing __getitem__ and __iter__'
1916 def __init__(self, seqn):
1917 self.seqn = seqn
1918 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001919 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001920 if self.i >= len(self.seqn): raise StopIteration
1921 v = self.seqn[self.i]
1922 self.i += 1
1923 return v
1924
1925class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001926 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001927 def __init__(self, seqn):
1928 self.seqn = seqn
1929 self.i = 0
1930 def __iter__(self):
1931 return self
1932
1933class E:
1934 'Test propagation of exceptions'
1935 def __init__(self, seqn):
1936 self.seqn = seqn
1937 self.i = 0
1938 def __iter__(self):
1939 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001940 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001941 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001942
1943class S:
1944 'Test immediate stop'
1945 def __init__(self, seqn):
1946 pass
1947 def __iter__(self):
1948 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001949 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001950 raise StopIteration
1951
1952def L(seqn):
1953 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001954 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001955
1956
1957class TestVariousIteratorArgs(unittest.TestCase):
1958
Raymond Hettinger482ba772010-12-01 22:48:00 +00001959 def test_accumulate(self):
1960 s = [1,2,3,4,5]
1961 r = [1,3,6,10,15]
1962 n = len(s)
1963 for g in (G, I, Ig, L, R):
1964 self.assertEqual(list(accumulate(g(s))), r)
1965 self.assertEqual(list(accumulate(S(s))), [])
1966 self.assertRaises(TypeError, accumulate, X(s))
1967 self.assertRaises(TypeError, accumulate, N(s))
1968 self.assertRaises(ZeroDivisionError, list, accumulate(E(s)))
1969
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001970 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001971 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001972 for g in (G, I, Ig, S, L, R):
1973 self.assertEqual(list(chain(g(s))), list(g(s)))
1974 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001975 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001976 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001977 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1978
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001979 def test_compress(self):
1980 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1981 n = len(s)
1982 for g in (G, I, Ig, S, L, R):
1983 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1984 self.assertRaises(TypeError, compress, X(s), repeat(1))
1985 self.assertRaises(TypeError, compress, N(s), repeat(1))
1986 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1987
Christian Heimesc3f30c42008-02-22 16:37:40 +00001988 def test_product(self):
1989 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1990 self.assertRaises(TypeError, product, X(s))
1991 self.assertRaises(TypeError, product, N(s))
1992 self.assertRaises(ZeroDivisionError, product, E(s))
1993
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001994 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001995 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001996 for g in (G, I, Ig, S, L, R):
1997 tgtlen = len(s) * 3
1998 expected = list(g(s))*3
1999 actual = list(islice(cycle(g(s)), tgtlen))
2000 self.assertEqual(actual, expected)
2001 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00002002 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002003 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
2004
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002005 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00002006 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002007 for g in (G, I, Ig, S, L, R):
2008 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
2009 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00002010 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002011 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
2012
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002013 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00002014 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002015 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002016 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00002017 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002018 self.assertRaises(TypeError, filter, isEven, X(s))
2019 self.assertRaises(TypeError, filter, isEven, N(s))
2020 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002021
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002022 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00002023 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002024 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002025 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00002026 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002027 self.assertRaises(TypeError, filterfalse, isEven, X(s))
2028 self.assertRaises(TypeError, filterfalse, isEven, N(s))
2029 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002030
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002031 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00002032 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002033 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002034 self.assertEqual(list(zip(g(s))), lzip(g(s)))
2035 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
2036 self.assertRaises(TypeError, zip, X(s))
2037 self.assertRaises(TypeError, zip, N(s))
2038 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002039
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002040 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00002041 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00002042 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00002043 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
2044 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
2045 self.assertRaises(TypeError, zip_longest, X(s))
2046 self.assertRaises(TypeError, zip_longest, N(s))
2047 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00002048
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002049 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00002050 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002051 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002052 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00002053 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002054 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00002055 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002056 self.assertRaises(TypeError, map, onearg, X(s))
2057 self.assertRaises(TypeError, map, onearg, N(s))
2058 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002059
2060 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00002061 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002062 for g in (G, I, Ig, S, L, R):
2063 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
2064 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00002065 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002066 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
2067
Raymond Hettingercc061d02020-11-30 20:42:54 -08002068 def test_pairwise(self):
2069 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
2070 for g in (G, I, Ig, S, L, R):
2071 seq = list(g(s))
2072 expected = list(zip(seq, seq[1:]))
2073 actual = list(pairwise(g(s)))
2074 self.assertEqual(actual, expected)
2075 self.assertRaises(TypeError, pairwise, X(s))
2076 self.assertRaises(TypeError, pairwise, N(s))
2077 self.assertRaises(ZeroDivisionError, list, pairwise(E(s)))
2078
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002079 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00002080 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002081 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00002082 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00002083 self.assertEqual(list(starmap(operator.pow, g(ss))),
2084 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002085 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00002086 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002087 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
2088
2089 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00002090 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002091 for g in (G, I, Ig, S, L, R):
2092 tgt = []
2093 for elem in g(s):
2094 if not isEven(elem): break
2095 tgt.append(elem)
2096 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
2097 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00002098 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002099 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
2100
2101 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00002102 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002103 for g in (G, I, Ig, S, L, R):
2104 tgt = []
2105 for elem in g(s):
2106 if not tgt and isOdd(elem): continue
2107 tgt.append(elem)
2108 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
2109 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00002110 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002111 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
2112
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002113 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00002114 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002115 for g in (G, I, Ig, S, L, R):
2116 it1, it2 = tee(g(s))
2117 self.assertEqual(list(it1), list(g(s)))
2118 self.assertEqual(list(it2), list(g(s)))
2119 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00002120 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002121 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
2122
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002123class LengthTransparency(unittest.TestCase):
2124
2125 def test_repeat(self):
Armin Ronacheraa9a79d2012-10-06 14:03:24 +02002126 self.assertEqual(operator.length_hint(repeat(None, 50)), 50)
Raymond Hettinger97d35552014-06-24 21:36:58 -07002127 self.assertEqual(operator.length_hint(repeat(None, 0)), 0)
Armin Ronacheraa9a79d2012-10-06 14:03:24 +02002128 self.assertEqual(operator.length_hint(repeat(None), 12), 12)
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002129
Raymond Hettinger97d35552014-06-24 21:36:58 -07002130 def test_repeat_with_negative_times(self):
2131 self.assertEqual(operator.length_hint(repeat(None, -1)), 0)
2132 self.assertEqual(operator.length_hint(repeat(None, -2)), 0)
2133 self.assertEqual(operator.length_hint(repeat(None, times=-1)), 0)
2134 self.assertEqual(operator.length_hint(repeat(None, times=-2)), 0)
2135
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002136class RegressionTests(unittest.TestCase):
2137
2138 def test_sf_793826(self):
2139 # Fix Armin Rigo's successful efforts to wreak havoc
2140
2141 def mutatingtuple(tuple1, f, tuple2):
2142 # this builds a tuple t which is a copy of tuple1,
2143 # then calls f(t), then mutates t to be equal to tuple2
2144 # (needs len(tuple1) == len(tuple2)).
2145 def g(value, first=[1]):
2146 if first:
2147 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00002148 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002149 return value
2150 items = list(tuple2)
2151 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002152 gen = map(g, items)
2153 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00002154 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002155
2156 def f(t):
2157 global T
2158 T = t
2159 first[:] = list(T)
2160
2161 first = []
2162 mutatingtuple((1,2,3), f, (4,5,6))
2163 second = list(T)
2164 self.assertEqual(first, second)
2165
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002166
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00002167 def test_sf_950057(self):
2168 # Make sure that chain() and cycle() catch exceptions immediately
2169 # rather than when shifting between input sources
2170
2171 def gen1():
2172 hist.append(0)
2173 yield 1
2174 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00002175 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00002176 hist.append(2)
2177
2178 def gen2(x):
2179 hist.append(3)
2180 yield 2
2181 hist.append(4)
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00002182
2183 hist = []
2184 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
2185 self.assertEqual(hist, [0,1])
2186
2187 hist = []
2188 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
2189 self.assertEqual(hist, [0,1])
2190
2191 hist = []
2192 self.assertRaises(AssertionError, list, cycle(gen1()))
2193 self.assertEqual(hist, [0,1])
2194
Neil Schemenauer52a48e62019-07-30 11:08:18 -07002195 @support.skip_if_pgo_task
T. Wouters5466d4a2017-03-30 09:58:35 -07002196 def test_long_chain_of_empty_iterables(self):
2197 # Make sure itertools.chain doesn't run into recursion limits when
2198 # dealing with long chains of empty iterables. Even with a high
2199 # number this would probably only fail in Py_DEBUG mode.
2200 it = chain.from_iterable(() for unused in range(10000000))
2201 with self.assertRaises(StopIteration):
2202 next(it)
2203
Serhiy Storchakac740e4f2017-09-26 21:47:56 +03002204 def test_issue30347_1(self):
2205 def f(n):
2206 if n == 5:
2207 list(b)
2208 return n != 6
2209 for (k, b) in groupby(range(10), f):
2210 list(b) # shouldn't crash
2211
2212 def test_issue30347_2(self):
2213 class K:
2214 def __init__(self, v):
2215 pass
2216 def __eq__(self, other):
2217 nonlocal i
2218 i += 1
2219 if i == 1:
2220 next(g, None)
2221 return True
2222 i = 0
2223 g = next(groupby(range(10), K))[1]
2224 for j in range(2):
2225 next(g, None) # shouldn't crash
2226
2227
Thomas Woutersb2137042007-02-01 18:02:27 +00002228class SubclassWithKwargsTest(unittest.TestCase):
2229 def test_keywords_in_subclass(self):
2230 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002231 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002232 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00002233 class Subclass(cls):
2234 def __init__(self, newarg=None, *args):
2235 cls.__init__(self, *args)
2236 try:
2237 Subclass(newarg=1)
2238 except TypeError as err:
2239 # we expect type errors because of wrong argument count
Serhiy Storchaka5eb788b2017-06-06 18:45:22 +03002240 self.assertNotIn("keyword arguments", err.args[0])
Thomas Woutersb2137042007-02-01 18:02:27 +00002241
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002242@support.cpython_only
2243class SizeofTest(unittest.TestCase):
2244 def setUp(self):
2245 self.ssize_t = struct.calcsize('n')
2246
2247 check_sizeof = support.check_sizeof
2248
2249 def test_product_sizeof(self):
2250 basesize = support.calcobjsize('3Pi')
2251 check = self.check_sizeof
2252 check(product('ab', '12'), basesize + 2 * self.ssize_t)
2253 check(product(*(('abc',) * 10)), basesize + 10 * self.ssize_t)
2254
2255 def test_combinations_sizeof(self):
2256 basesize = support.calcobjsize('3Pni')
2257 check = self.check_sizeof
2258 check(combinations('abcd', 3), basesize + 3 * self.ssize_t)
2259 check(combinations(range(10), 4), basesize + 4 * self.ssize_t)
2260
2261 def test_combinations_with_replacement_sizeof(self):
2262 cwr = combinations_with_replacement
2263 basesize = support.calcobjsize('3Pni')
2264 check = self.check_sizeof
2265 check(cwr('abcd', 3), basesize + 3 * self.ssize_t)
2266 check(cwr(range(10), 4), basesize + 4 * self.ssize_t)
2267
2268 def test_permutations_sizeof(self):
2269 basesize = support.calcobjsize('4Pni')
2270 check = self.check_sizeof
2271 check(permutations('abcd'),
2272 basesize + 4 * self.ssize_t + 4 * self.ssize_t)
2273 check(permutations('abcd', 3),
2274 basesize + 4 * self.ssize_t + 3 * self.ssize_t)
2275 check(permutations('abcde', 3),
2276 basesize + 5 * self.ssize_t + 3 * self.ssize_t)
2277 check(permutations(range(10), 4),
2278 basesize + 10 * self.ssize_t + 4 * self.ssize_t)
2279
Thomas Woutersb2137042007-02-01 18:02:27 +00002280
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002281libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002282
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002283
2284>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002285>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00002286... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00002287...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002288Check 1200 is for $120.15
2289Check 1201 is for $764.05
2290Check 1202 is for $823.14
2291
2292>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002293>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00002294... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00002295...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000022961
22978
229827
2299
2300>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00002301>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00002302... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00002303...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002304Alex
2305Laura
2306Martin
2307Walter
2308Samuele
2309
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002310>>> from operator import itemgetter
2311>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00002312>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002313>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00002314... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00002315...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000023161 ['a', 'c', 'e']
23172 ['b', 'd', 'f']
23183 ['g']
2319
Raymond Hettinger734fb572004-01-20 20:04:40 +00002320# Find runs of consecutive numbers using groupby. The key to the solution
2321# is differencing with a range so that consecutive numbers all appear in
2322# same group.
2323>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00002324>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00002325... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00002326...
Raymond Hettinger734fb572004-01-20 20:04:40 +00002327[1]
2328[4, 5, 6]
2329[10]
2330[15, 16, 17, 18]
2331[22]
2332[25, 26, 27, 28]
2333
Georg Brandl3dbca812008-07-23 16:10:53 +00002334>>> def take(n, iterable):
2335... "Return first n items of the iterable as a list"
2336... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00002337
Raymond Hettinger9265dd72018-04-08 08:44:20 -07002338>>> def prepend(value, iterator):
2339... "Prepend a single value in front of an iterator"
2340... # prepend(1, [2, 3, 4]) -> 1 2 3 4
2341... return chain([value], iterator)
2342
Georg Brandl3dbca812008-07-23 16:10:53 +00002343>>> def enumerate(iterable, start=0):
2344... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002345
Georg Brandl3dbca812008-07-23 16:10:53 +00002346>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002347... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00002348... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002349
Cheryl Sabellada1734c2018-03-26 21:29:33 -04002350>>> import collections
2351>>> def consume(iterator, n=None):
2352... "Advance the iterator n-steps ahead. If n is None, consume entirely."
2353... # Use functions that consume iterators at C speed.
2354... if n is None:
2355... # feed the entire iterator into a zero-length deque
2356... collections.deque(iterator, maxlen=0)
2357... else:
2358... # advance to the empty slice starting at position n
2359... next(islice(iterator, n, n), None)
2360
Raymond Hettingercdf8ba32009-02-19 04:45:07 +00002361>>> def nth(iterable, n, default=None):
2362... "Returns the nth item or a default value"
2363... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002364
Raymond Hettingere525ee32016-03-06 18:11:38 -08002365>>> def all_equal(iterable):
2366... "Returns True if all the elements are equal to each other"
2367... g = groupby(iterable)
2368... return next(g, True) and not next(g, False)
2369
Georg Brandl3dbca812008-07-23 16:10:53 +00002370>>> def quantify(iterable, pred=bool):
2371... "Count how many times the predicate is true"
2372... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00002373
Raymond Hettingerfc40b302020-11-29 10:47:22 -08002374>>> def pad_none(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002375... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00002376... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002377
Georg Brandl3dbca812008-07-23 16:10:53 +00002378>>> def ncycles(iterable, n):
Ezio Melotti13925002011-03-16 11:05:33 +02002379... "Returns the sequence elements n times"
Georg Brandl3dbca812008-07-23 16:10:53 +00002380... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002381
2382>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002383... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002384
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002385>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002386... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002387
2388>>> def repeatfunc(func, times=None, *args):
2389... "Repeat calls to func with specified arguments."
2390... " Example: repeatfunc(random.random)"
2391... if times is None:
2392... return starmap(func, repeat(args))
2393... else:
2394... return starmap(func, repeat(args, times))
2395
Miss Islington (bot)be33e582021-09-07 10:52:26 -07002396>>> def triplewise(iterable):
2397... "Return overlapping triplets from an iterable"
2398... # pairwise('ABCDEFG') -> ABC BCD CDE DEF EFG
2399... for (a, _), (b, c) in pairwise(pairwise(iterable)):
2400... yield a, b, c
2401
2402>>> import collections
2403>>> def sliding_window(iterable, n):
2404... # sliding_window('ABCDEFG', 4) -> ABCD BCDE CDEF DEFG
2405... it = iter(iterable)
2406... window = collections.deque(islice(it, n), maxlen=n)
2407... if len(window) == n:
2408... yield tuple(window)
2409... for x in it:
2410... window.append(x)
2411... yield tuple(window)
2412
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002413>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00002414... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002415... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00002416... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002417
2418>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00002419... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002420... # Recipe credited to George Sakkis
2421... pending = len(iterables)
2422... nexts = cycle(iter(it).__next__ for it in iterables)
2423... while pending:
2424... try:
2425... for next in nexts:
2426... yield next()
2427... except StopIteration:
2428... pending -= 1
2429... nexts = cycle(islice(nexts, pending))
2430
Miss Islington (bot)656b0bd2021-09-04 22:30:37 -07002431>>> def partition(pred, iterable):
2432... "Use a predicate to partition entries into false entries and true entries"
2433... # partition(is_odd, range(10)) --> 0 2 4 6 8 and 1 3 5 7 9
2434... t1, t2 = tee(iterable)
2435... return filterfalse(pred, t1), filter(pred, t2)
2436
2437>>> def before_and_after(predicate, it):
2438... ''' Variant of takewhile() that allows complete
2439... access to the remainder of the iterator.
2440...
2441... >>> all_upper, remainder = before_and_after(str.isupper, 'ABCdEfGhI')
2442... >>> str.join('', all_upper)
2443... 'ABC'
2444... >>> str.join('', remainder)
2445... 'dEfGhI'
2446...
2447... Note that the first iterator must be fully
2448... consumed before the second iterator can
2449... generate valid results.
2450... '''
2451... it = iter(it)
2452... transition = []
2453... def true_iterator():
2454... for elem in it:
2455... if predicate(elem):
2456... yield elem
2457... else:
2458... transition.append(elem)
2459... return
2460... def remainder_iterator():
2461... yield from transition
2462... yield from it
2463... return true_iterator(), remainder_iterator()
2464
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002465>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00002466... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
2467... s = list(iterable)
2468... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002469
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00002470>>> def unique_everseen(iterable, key=None):
2471... "List unique elements, preserving order. Remember all elements ever seen."
2472... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
2473... # unique_everseen('ABBCcAD', str.lower) --> A B C D
2474... seen = set()
2475... seen_add = seen.add
2476... if key is None:
2477... for element in iterable:
2478... if element not in seen:
2479... seen_add(element)
2480... yield element
2481... else:
2482... for element in iterable:
2483... k = key(element)
2484... if k not in seen:
2485... seen_add(k)
2486... yield element
2487
2488>>> def unique_justseen(iterable, key=None):
2489... "List unique elements, preserving order. Remember only the element just seen."
2490... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
2491... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
2492... return map(next, map(itemgetter(1), groupby(iterable, key)))
2493
Raymond Hettinger31b26f62014-04-02 03:16:42 -07002494>>> def first_true(iterable, default=False, pred=None):
2495... '''Returns the first true value in the iterable.
2496...
2497... If no true value is found, returns *default*
2498...
2499... If *pred* is not None, returns the first item
2500... for which pred(item) is true.
2501...
2502... '''
2503... # first_true([a,b,c], x) --> a or b or c or x
2504... # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x
2505... return next(filter(pred, iterable), default)
2506
Raymond Hettingerd37258d2018-01-13 10:35:40 -08002507>>> def nth_combination(iterable, r, index):
2508... 'Equivalent to list(combinations(iterable, r))[index]'
2509... pool = tuple(iterable)
2510... n = len(pool)
2511... if r < 0 or r > n:
2512... raise ValueError
2513... c = 1
2514... k = min(r, n-r)
2515... for i in range(1, k+1):
2516... c = c * (n - k + i) // i
2517... if index < 0:
2518... index += c
2519... if index < 0 or index >= c:
2520... raise IndexError
2521... result = []
2522... while r:
2523... c, n, r = c*r//n, n-1, r-1
2524... while index >= c:
2525... index -= c
2526... c, n = c*(n-r)//n, n-1
2527... result.append(pool[-1-n])
2528... return tuple(result)
2529
2530
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002531This is not part of the examples but it tests to make sure the definitions
2532perform as purported.
2533
Raymond Hettingera098b332003-09-08 23:58:40 +00002534>>> take(10, count())
2535[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
2536
Raymond Hettinger9265dd72018-04-08 08:44:20 -07002537>>> list(prepend(1, [2, 3, 4]))
2538[1, 2, 3, 4]
2539
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002540>>> list(enumerate('abc'))
2541[(0, 'a'), (1, 'b'), (2, 'c')]
2542
2543>>> list(islice(tabulate(lambda x: 2*x), 4))
2544[0, 2, 4, 6]
2545
Cheryl Sabellada1734c2018-03-26 21:29:33 -04002546>>> it = iter(range(10))
2547>>> consume(it, 3)
2548>>> next(it)
25493
2550>>> consume(it)
2551>>> next(it, 'Done')
2552'Done'
2553
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002554>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00002555'd'
2556
2557>>> nth('abcde', 9) is None
2558True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002559
Raymond Hettingere525ee32016-03-06 18:11:38 -08002560>>> [all_equal(s) for s in ('', 'A', 'AAAA', 'AAAB', 'AAABA')]
2561[True, True, True, False, False]
2562
Guido van Rossum805365e2007-05-07 22:24:25 +00002563>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000256450
2565
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002566>>> a = [[1, 2, 3], [4, 5, 6]]
2567>>> flatten(a)
2568[1, 2, 3, 4, 5, 6]
2569
2570>>> list(repeatfunc(pow, 5, 2, 3))
2571[8, 8, 8, 8, 8]
2572
2573>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002574>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002575[0, 0, 0, 0, 0]
2576
Raymond Hettingerfc40b302020-11-29 10:47:22 -08002577>>> list(islice(pad_none('abc'), 0, 6))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002578['a', 'b', 'c', None, None, None]
2579
2580>>> list(ncycles('abc', 3))
2581['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
2582
2583>>> dotproduct([1,2,3], [4,5,6])
258432
2585
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002586>>> list(grouper(3, 'abcdefg', 'x'))
2587[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
2588
Miss Islington (bot)be33e582021-09-07 10:52:26 -07002589>>> list(triplewise('ABCDEFG'))
2590[('A', 'B', 'C'), ('B', 'C', 'D'), ('C', 'D', 'E'), ('D', 'E', 'F'), ('E', 'F', 'G')]
2591
2592>>> list(sliding_window('ABCDEFG', 4))
2593[('A', 'B', 'C', 'D'), ('B', 'C', 'D', 'E'), ('C', 'D', 'E', 'F'), ('D', 'E', 'F', 'G')]
2594
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002595>>> list(roundrobin('abc', 'd', 'ef'))
2596['a', 'd', 'e', 'b', 'f', 'c']
2597
Miss Islington (bot)656b0bd2021-09-04 22:30:37 -07002598>>> def is_odd(x):
2599... return x % 2 == 1
2600
2601>>> evens, odds = partition(is_odd, range(10))
2602>>> list(evens)
2603[0, 2, 4, 6, 8]
2604>>> list(odds)
2605[1, 3, 5, 7, 9]
2606
2607>>> all_upper, remainder = before_and_after(str.isupper, 'ABCdEfGhI')
2608>>> str.join('', all_upper)
2609'ABC'
2610>>> str.join('', remainder)
2611'dEfGhI'
2612
Raymond Hettingerace67332009-01-26 02:23:50 +00002613>>> list(powerset([1,2,3]))
2614[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002615
Raymond Hettinger191e8502009-01-27 13:29:43 +00002616>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
2617True
2618
2619>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
2620True
2621
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00002622>>> list(unique_everseen('AAAABBBCCDAABBB'))
2623['A', 'B', 'C', 'D']
2624
2625>>> list(unique_everseen('ABBCcAD', str.lower))
2626['A', 'B', 'C', 'D']
2627
2628>>> list(unique_justseen('AAAABBBCCDAABBB'))
2629['A', 'B', 'C', 'D', 'A', 'B']
2630
2631>>> list(unique_justseen('ABBCcAD', str.lower))
2632['A', 'B', 'C', 'A', 'D']
2633
Raymond Hettinger31b26f62014-04-02 03:16:42 -07002634>>> first_true('ABC0DEF1', '9', str.isdigit)
2635'0'
2636
Raymond Hettingerd37258d2018-01-13 10:35:40 -08002637>>> population = 'ABCDEFGH'
2638>>> for r in range(len(population) + 1):
2639... seq = list(combinations(population, r))
2640... for i in range(len(seq)):
2641... assert nth_combination(population, r, i) == seq[i]
2642... for i in range(-len(seq), 0):
2643... assert nth_combination(population, r, i) == seq[i]
2644
2645
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002646"""
2647
2648__test__ = {'libreftest' : libreftest}
2649
2650def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002651 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00002652 RegressionTests, LengthTransparency,
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002653 SubclassWithKwargsTest, TestExamples,
Cheryl Sabellada1734c2018-03-26 21:29:33 -04002654 TestPurePythonRoughEquivalents,
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002655 SizeofTest)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002656 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002657
2658 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002659 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00002660 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00002661 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00002662 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002663 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00002664 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00002665 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00002666 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002667
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002668 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002669 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002670
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002671if __name__ == "__main__":
2672 test_main(verbose=True)