blob: cbbb4c4f71d3b8d2e8fad564078ee14b0c467291 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001import unittest
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002from test import support
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003from itertools import *
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02004import weakref
Raymond Hettinger8a882a72009-02-12 12:08:18 +00005from decimal import Decimal
Raymond Hettingerde09acf2009-02-12 12:53:53 +00006from fractions import Fraction
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00007import operator
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00008import random
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00009import copy
10import pickle
Christian Heimesc3f30c42008-02-22 16:37:40 +000011from functools import reduce
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +020012import sys
13import struct
Benjamin Petersonee8712c2008-05-20 21:35:26 +000014maxsize = support.MAX_Py_ssize_t
Guido van Rossum360e4b82007-05-14 22:51:27 +000015minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000016
Guido van Rossum801f0d72006-08-24 19:48:10 +000017def lzip(*args):
18 return list(zip(*args))
19
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000020def onearg(x):
21 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000022 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000023
24def errfunc(*args):
25 'Test function that raises an error'
26 raise ValueError
27
28def gen3():
29 'Non-restartable source sequence'
30 for i in (0, 1, 2):
31 yield i
32
33def isEven(x):
34 'Test predicate'
35 return x%2==0
36
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000037def isOdd(x):
38 'Test predicate'
39 return x%2==1
40
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000041def tupleize(*args):
42 return args
43
44def irange(n):
45 for i in range(n):
46 yield i
47
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000048class StopNow:
49 'Class emulating an empty iterable.'
50 def __iter__(self):
51 return self
Georg Brandla18af4e2007-04-21 15:47:16 +000052 def __next__(self):
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000053 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000054
Raymond Hettinger02420702003-06-29 20:36:23 +000055def take(n, seq):
56 'Convenience function for partially consuming a long of infinite iterable'
57 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000058
Christian Heimes78644762008-03-04 23:39:23 +000059def prod(iterable):
60 return reduce(operator.mul, iterable, 1)
61
Christian Heimes380f7f22008-02-28 11:19:05 +000062def fact(n):
63 'Factorial'
Christian Heimes78644762008-03-04 23:39:23 +000064 return prod(range(1, n+1))
65
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000066# root level methods for pickling ability
67def testR(r):
68 return r[0]
69
70def testR2(r):
71 return r[2]
72
73def underten(x):
74 return x<10
75
Serhiy Storchakabad12572014-12-15 14:03:42 +020076picklecopiers = [lambda s, proto=proto: pickle.loads(pickle.dumps(s, proto))
77 for proto in range(pickle.HIGHEST_PROTOCOL + 1)]
78
Raymond Hettinger96ef8112003-02-01 00:10:11 +000079class TestBasicOps(unittest.TestCase):
Raymond Hettinger482ba772010-12-01 22:48:00 +000080
Serhiy Storchakabad12572014-12-15 14:03:42 +020081 def pickletest(self, protocol, it, stop=4, take=1, compare=None):
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000082 """Test that an iterator is the same after pickling, also when part-consumed"""
83 def expand(it, i=0):
84 # Recursively expand iterables, within sensible bounds
85 if i > 10:
86 raise RuntimeError("infinite recursion encountered")
87 if isinstance(it, str):
88 return it
89 try:
90 l = list(islice(it, stop))
91 except TypeError:
92 return it # can't expand it
93 return [expand(e, i+1) for e in l]
94
95 # Test the initial copy against the original
Serhiy Storchakabad12572014-12-15 14:03:42 +020096 dump = pickle.dumps(it, protocol)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000097 i2 = pickle.loads(dump)
98 self.assertEqual(type(it), type(i2))
99 a, b = expand(it), expand(i2)
100 self.assertEqual(a, b)
101 if compare:
102 c = expand(compare)
103 self.assertEqual(a, c)
104
105 # Take from the copy, and create another copy and compare them.
106 i3 = pickle.loads(dump)
107 took = 0
108 try:
109 for i in range(take):
110 next(i3)
111 took += 1
112 except StopIteration:
113 pass #in case there is less data than 'take'
Serhiy Storchakabad12572014-12-15 14:03:42 +0200114 dump = pickle.dumps(i3, protocol)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000115 i4 = pickle.loads(dump)
116 a, b = expand(i3), expand(i4)
117 self.assertEqual(a, b)
118 if compare:
119 c = expand(compare[took:])
120 self.assertEqual(a, c);
121
Raymond Hettinger482ba772010-12-01 22:48:00 +0000122 def test_accumulate(self):
123 self.assertEqual(list(accumulate(range(10))), # one positional arg
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000124 [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])
125 self.assertEqual(list(accumulate(iterable=range(10))), # kw arg
126 [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])
Raymond Hettinger482ba772010-12-01 22:48:00 +0000127 for typ in int, complex, Decimal, Fraction: # multiple types
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000128 self.assertEqual(
129 list(accumulate(map(typ, range(10)))),
Raymond Hettinger482ba772010-12-01 22:48:00 +0000130 list(map(typ, [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])))
Raymond Hettinger2d93e6e2010-12-03 02:33:53 +0000131 self.assertEqual(list(accumulate('abc')), ['a', 'ab', 'abc']) # works with non-numeric
Raymond Hettinger482ba772010-12-01 22:48:00 +0000132 self.assertEqual(list(accumulate([])), []) # empty iterable
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000133 self.assertEqual(list(accumulate([7])), [7]) # iterable of length one
Raymond Hettinger5d446132011-03-27 18:52:10 -0700134 self.assertRaises(TypeError, accumulate, range(10), 5, 6) # too many args
Raymond Hettinger482ba772010-12-01 22:48:00 +0000135 self.assertRaises(TypeError, accumulate) # too few args
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000136 self.assertRaises(TypeError, accumulate, x=range(10)) # unexpected kwd arg
Raymond Hettinger482ba772010-12-01 22:48:00 +0000137 self.assertRaises(TypeError, list, accumulate([1, []])) # args that don't add
138
Raymond Hettinger5d446132011-03-27 18:52:10 -0700139 s = [2, 8, 9, 5, 7, 0, 3, 4, 1, 6]
140 self.assertEqual(list(accumulate(s, min)),
141 [2, 2, 2, 2, 2, 0, 0, 0, 0, 0])
142 self.assertEqual(list(accumulate(s, max)),
143 [2, 8, 9, 9, 9, 9, 9, 9, 9, 9])
144 self.assertEqual(list(accumulate(s, operator.mul)),
145 [2, 16, 144, 720, 5040, 0, 0, 0, 0, 0])
146 with self.assertRaises(TypeError):
147 list(accumulate(s, chr)) # unary-operation
Serhiy Storchakabad12572014-12-15 14:03:42 +0200148 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
149 self.pickletest(proto, accumulate(range(10))) # test pickling
Raymond Hettinger5d446132011-03-27 18:52:10 -0700150
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000151 def test_chain(self):
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000152
153 def chain2(*iterables):
154 'Pure python version in the docs'
155 for it in iterables:
156 for element in it:
157 yield element
158
159 for c in (chain, chain2):
160 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
161 self.assertEqual(list(c('abc')), list('abc'))
162 self.assertEqual(list(c('')), [])
163 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
164 self.assertRaises(TypeError, list,c(2, 3))
Christian Heimesf16baeb2008-02-29 14:57:44 +0000165
166 def test_chain_from_iterable(self):
167 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
168 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
169 self.assertEqual(list(chain.from_iterable([''])), [])
170 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
171 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000172
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000173 def test_chain_reducible(self):
Serhiy Storchakabad12572014-12-15 14:03:42 +0200174 for oper in [copy.deepcopy] + picklecopiers:
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000175 it = chain('abc', 'def')
176 self.assertEqual(list(oper(it)), list('abcdef'))
177 self.assertEqual(next(it), 'a')
178 self.assertEqual(list(oper(it)), list('bcdef'))
179
180 self.assertEqual(list(oper(chain(''))), [])
181 self.assertEqual(take(4, oper(chain('abc', 'def'))), list('abcd'))
182 self.assertRaises(TypeError, list, oper(chain(2, 3)))
Serhiy Storchakabad12572014-12-15 14:03:42 +0200183 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
184 self.pickletest(proto, chain('abc', 'def'), compare=list('abcdef'))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000185
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300186 def test_chain_setstate(self):
187 self.assertRaises(TypeError, chain().__setstate__, ())
188 self.assertRaises(TypeError, chain().__setstate__, [])
189 self.assertRaises(TypeError, chain().__setstate__, 0)
190 self.assertRaises(TypeError, chain().__setstate__, ([],))
191 self.assertRaises(TypeError, chain().__setstate__, (iter([]), []))
192 it = chain()
193 it.__setstate__((iter(['abc', 'def']),))
194 self.assertEqual(list(it), ['a', 'b', 'c', 'd', 'e', 'f'])
195 it = chain()
196 it.__setstate__((iter(['abc', 'def']), iter(['ghi'])))
197 self.assertEqual(list(it), ['ghi', 'a', 'b', 'c', 'd', 'e', 'f'])
198
Christian Heimes380f7f22008-02-28 11:19:05 +0000199 def test_combinations(self):
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000200 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Christian Heimes380f7f22008-02-28 11:19:05 +0000201 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Christian Heimes78644762008-03-04 23:39:23 +0000202 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Christian Heimes380f7f22008-02-28 11:19:05 +0000203 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000204
Serhiy Storchakabad12572014-12-15 14:03:42 +0200205 for op in [lambda a:a] + picklecopiers:
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000206 self.assertEqual(list(op(combinations('abc', 32))), []) # r > n
207
208 self.assertEqual(list(op(combinations('ABCD', 2))),
209 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
210 testIntermediate = combinations('ABCD', 2)
211 next(testIntermediate)
212 self.assertEqual(list(op(testIntermediate)),
213 [('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
214
215 self.assertEqual(list(op(combinations(range(4), 3))),
216 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
217 testIntermediate = combinations(range(4), 3)
218 next(testIntermediate)
219 self.assertEqual(list(op(testIntermediate)),
220 [(0,1,3), (0,2,3), (1,2,3)])
221
Christian Heimes78644762008-03-04 23:39:23 +0000222
223 def combinations1(iterable, r):
224 'Pure python version shown in the docs'
225 pool = tuple(iterable)
226 n = len(pool)
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000227 if r > n:
228 return
Christian Heimes78644762008-03-04 23:39:23 +0000229 indices = list(range(r))
230 yield tuple(pool[i] for i in indices)
231 while 1:
232 for i in reversed(range(r)):
233 if indices[i] != i + n - r:
234 break
235 else:
236 return
237 indices[i] += 1
238 for j in range(i+1, r):
239 indices[j] = indices[j-1] + 1
240 yield tuple(pool[i] for i in indices)
241
242 def combinations2(iterable, r):
243 'Pure python version shown in the docs'
244 pool = tuple(iterable)
245 n = len(pool)
246 for indices in permutations(range(n), r):
247 if sorted(indices) == list(indices):
248 yield tuple(pool[i] for i in indices)
249
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000250 def combinations3(iterable, r):
251 'Pure python version from cwr()'
252 pool = tuple(iterable)
253 n = len(pool)
254 for indices in combinations_with_replacement(range(n), r):
255 if len(set(indices)) == r:
256 yield tuple(pool[i] for i in indices)
257
Christian Heimes78644762008-03-04 23:39:23 +0000258 for n in range(7):
Christian Heimes380f7f22008-02-28 11:19:05 +0000259 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000260 for r in range(n+2):
Christian Heimes380f7f22008-02-28 11:19:05 +0000261 result = list(combinations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000262 self.assertEqual(len(result), 0 if r>n else fact(n) / fact(r) / fact(n-r)) # right number of combs
Christian Heimes380f7f22008-02-28 11:19:05 +0000263 self.assertEqual(len(result), len(set(result))) # no repeats
264 self.assertEqual(result, sorted(result)) # lexicographic order
265 for c in result:
266 self.assertEqual(len(c), r) # r-length combinations
267 self.assertEqual(len(set(c)), r) # no duplicate elements
268 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000269 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000270 self.assertEqual(list(c),
271 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000272 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000273 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000274 self.assertEqual(result, list(combinations3(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000275
Serhiy Storchakabad12572014-12-15 14:03:42 +0200276 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
277 self.pickletest(proto, combinations(values, r)) # test pickling
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000278
Benjamin Peterson4b40eeb2015-02-01 20:59:00 -0500279 @support.bigaddrspacetest
280 def test_combinations_overflow(self):
Serhiy Storchakadee948b2015-02-03 01:34:09 +0200281 with self.assertRaises((OverflowError, MemoryError)):
Benjamin Peterson4b40eeb2015-02-01 20:59:00 -0500282 combinations("AA", 2**29)
283
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000284 # Test implementation detail: tuple re-use
Alex Gaynore151d212011-07-17 16:21:30 -0700285 @support.impl_detail("tuple reuse is specific to CPython")
286 def test_combinations_tuple_reuse(self):
Christian Heimes78644762008-03-04 23:39:23 +0000287 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
288 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
289
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000290 def test_combinations_with_replacement(self):
291 cwr = combinations_with_replacement
292 self.assertRaises(TypeError, cwr, 'abc') # missing r argument
293 self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
294 self.assertRaises(TypeError, cwr, None) # pool is not iterable
295 self.assertRaises(ValueError, cwr, 'abc', -2) # r is negative
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000296
Serhiy Storchakabad12572014-12-15 14:03:42 +0200297 for op in [lambda a:a] + picklecopiers:
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000298 self.assertEqual(list(op(cwr('ABC', 2))),
299 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
300 testIntermediate = cwr('ABC', 2)
301 next(testIntermediate)
302 self.assertEqual(list(op(testIntermediate)),
303 [('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
304
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000305
306 def cwr1(iterable, r):
307 'Pure python version shown in the docs'
308 # number items returned: (n+r-1)! / r! / (n-1)! when n>0
309 pool = tuple(iterable)
310 n = len(pool)
311 if not n and r:
312 return
313 indices = [0] * r
314 yield tuple(pool[i] for i in indices)
315 while 1:
316 for i in reversed(range(r)):
317 if indices[i] != n - 1:
318 break
319 else:
320 return
321 indices[i:] = [indices[i] + 1] * (r - i)
322 yield tuple(pool[i] for i in indices)
323
324 def cwr2(iterable, r):
325 'Pure python version shown in the docs'
326 pool = tuple(iterable)
327 n = len(pool)
328 for indices in product(range(n), repeat=r):
329 if sorted(indices) == list(indices):
330 yield tuple(pool[i] for i in indices)
331
332 def numcombs(n, r):
333 if not n:
334 return 0 if r else 1
335 return fact(n+r-1) / fact(r)/ fact(n-1)
336
337 for n in range(7):
338 values = [5*x-12 for x in range(n)]
339 for r in range(n+2):
340 result = list(cwr(values, r))
341
342 self.assertEqual(len(result), numcombs(n, r)) # right number of combs
343 self.assertEqual(len(result), len(set(result))) # no repeats
344 self.assertEqual(result, sorted(result)) # lexicographic order
345
346 regular_combs = list(combinations(values, r)) # compare to combs without replacement
347 if n == 0 or r <= 1:
Ezio Melottib3aedd42010-11-20 19:04:17 +0000348 self.assertEqual(result, regular_combs) # cases that should be identical
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000349 else:
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000350 self.assertTrue(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000351
352 for c in result:
353 self.assertEqual(len(c), r) # r-length combinations
354 noruns = [k for k,v in groupby(c)] # combo without consecutive repeats
355 self.assertEqual(len(noruns), len(set(noruns))) # no repeats other than consecutive
356 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000357 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000358 self.assertEqual(noruns,
359 [e for e in values if e in c]) # comb is a subsequence of the input iterable
360 self.assertEqual(result, list(cwr1(values, r))) # matches first pure python version
361 self.assertEqual(result, list(cwr2(values, r))) # matches second pure python version
362
Serhiy Storchakabad12572014-12-15 14:03:42 +0200363 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
364 self.pickletest(proto, cwr(values,r)) # test pickling
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000365
Benjamin Peterson6f082292015-02-01 21:10:47 -0500366 @support.bigaddrspacetest
367 def test_combinations_with_replacement_overflow(self):
Serhiy Storchakadee948b2015-02-03 01:34:09 +0200368 with self.assertRaises((OverflowError, MemoryError)):
Benjamin Peterson6f082292015-02-01 21:10:47 -0500369 combinations_with_replacement("AA", 2**30)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000370
Benjamin Peterson6f082292015-02-01 21:10:47 -0500371 # Test implementation detail: tuple re-use
Alex Gaynore151d212011-07-17 16:21:30 -0700372 @support.impl_detail("tuple reuse is specific to CPython")
373 def test_combinations_with_replacement_tuple_reuse(self):
374 cwr = combinations_with_replacement
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000375 self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
376 self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
377
Christian Heimes78644762008-03-04 23:39:23 +0000378 def test_permutations(self):
379 self.assertRaises(TypeError, permutations) # too few arguments
380 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000381 self.assertRaises(TypeError, permutations, None) # pool is not iterable
382 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000383 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000384 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Christian Heimes78644762008-03-04 23:39:23 +0000385 self.assertEqual(list(permutations(range(3), 2)),
386 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
387
388 def permutations1(iterable, r=None):
389 'Pure python version shown in the docs'
390 pool = tuple(iterable)
391 n = len(pool)
392 r = n if r is None else r
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000393 if r > n:
394 return
Christian Heimes78644762008-03-04 23:39:23 +0000395 indices = list(range(n))
396 cycles = list(range(n-r+1, n+1))[::-1]
397 yield tuple(pool[i] for i in indices[:r])
398 while n:
399 for i in reversed(range(r)):
400 cycles[i] -= 1
401 if cycles[i] == 0:
402 indices[i:] = indices[i+1:] + indices[i:i+1]
403 cycles[i] = n - i
404 else:
405 j = cycles[i]
406 indices[i], indices[-j] = indices[-j], indices[i]
407 yield tuple(pool[i] for i in indices[:r])
408 break
409 else:
410 return
411
412 def permutations2(iterable, r=None):
413 'Pure python version shown in the docs'
414 pool = tuple(iterable)
415 n = len(pool)
416 r = n if r is None else r
417 for indices in product(range(n), repeat=r):
418 if len(set(indices)) == r:
419 yield tuple(pool[i] for i in indices)
420
421 for n in range(7):
422 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000423 for r in range(n+2):
Christian Heimes78644762008-03-04 23:39:23 +0000424 result = list(permutations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000425 self.assertEqual(len(result), 0 if r>n else fact(n) / fact(n-r)) # right number of perms
Christian Heimes78644762008-03-04 23:39:23 +0000426 self.assertEqual(len(result), len(set(result))) # no repeats
427 self.assertEqual(result, sorted(result)) # lexicographic order
428 for p in result:
429 self.assertEqual(len(p), r) # r-length permutations
430 self.assertEqual(len(set(p)), r) # no duplicate elements
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000431 self.assertTrue(all(e in values for e in p)) # elements taken from input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000432 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000433 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000434 if r == n:
435 self.assertEqual(result, list(permutations(values, None))) # test r as None
436 self.assertEqual(result, list(permutations(values))) # test default r
437
Serhiy Storchakabad12572014-12-15 14:03:42 +0200438 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
439 self.pickletest(proto, permutations(values, r)) # test pickling
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000440
Benjamin Peterson0eaabf12015-02-01 21:34:07 -0500441 @support.bigaddrspacetest
442 def test_permutations_overflow(self):
Serhiy Storchakadee948b2015-02-03 01:34:09 +0200443 with self.assertRaises((OverflowError, MemoryError)):
Benjamin Peterson0eaabf12015-02-01 21:34:07 -0500444 permutations("A", 2**30)
Benjamin Peterson0eaabf12015-02-01 21:34:07 -0500445
Zachary Waredca807b2014-04-24 13:22:05 -0500446 @support.impl_detail("tuple reuse is specific to CPython")
Alex Gaynore151d212011-07-17 16:21:30 -0700447 def test_permutations_tuple_reuse(self):
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000448 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Christian Heimes78644762008-03-04 23:39:23 +0000449 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Christian Heimes380f7f22008-02-28 11:19:05 +0000450
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000451 def test_combinatorics(self):
452 # Test relationships between product(), permutations(),
453 # combinations() and combinations_with_replacement().
454
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000455 for n in range(6):
456 s = 'ABCDEFG'[:n]
457 for r in range(8):
458 prod = list(product(s, repeat=r))
459 cwr = list(combinations_with_replacement(s, r))
460 perm = list(permutations(s, r))
461 comb = list(combinations(s, r))
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000462
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000463 # Check size
Ezio Melottib3aedd42010-11-20 19:04:17 +0000464 self.assertEqual(len(prod), n**r)
465 self.assertEqual(len(cwr), (fact(n+r-1) / fact(r)/ fact(n-1)) if n else (not r))
466 self.assertEqual(len(perm), 0 if r>n else fact(n) / fact(n-r))
467 self.assertEqual(len(comb), 0 if r>n else fact(n) / fact(r) / fact(n-r))
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000468
469 # Check lexicographic order without repeated tuples
Ezio Melottib3aedd42010-11-20 19:04:17 +0000470 self.assertEqual(prod, sorted(set(prod)))
471 self.assertEqual(cwr, sorted(set(cwr)))
472 self.assertEqual(perm, sorted(set(perm)))
473 self.assertEqual(comb, sorted(set(comb)))
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000474
475 # Check interrelationships
Ezio Melottib3aedd42010-11-20 19:04:17 +0000476 self.assertEqual(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
477 self.assertEqual(perm, [t for t in prod if len(set(t))==r]) # perm: prods with no dups
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000478 self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
479 self.assertEqual(comb, [t for t in cwr if len(set(t))==r]) # comb: cwrs without dups
480 self.assertEqual(comb, list(filter(set(cwr).__contains__, perm))) # comb: perm that is a cwr
481 self.assertEqual(comb, list(filter(set(perm).__contains__, cwr))) # comb: cwr that is a perm
482 self.assertEqual(comb, sorted(set(cwr) & set(perm))) # comb: both a cwr and a perm
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000483
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000484 def test_compress(self):
Raymond Hettinger15a49502009-02-19 02:17:09 +0000485 self.assertEqual(list(compress(data='ABCDEF', selectors=[1,0,1,0,1,1])), list('ACEF'))
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000486 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
487 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
488 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
489 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
490 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
491 n = 10000
492 data = chain.from_iterable(repeat(range(6), n))
493 selectors = chain.from_iterable(repeat((0, 1)))
494 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
495 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
496 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
497 self.assertRaises(TypeError, compress, range(6)) # too few args
498 self.assertRaises(TypeError, compress, range(6), None) # too many args
499
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000500 # check copy, deepcopy, pickle
Serhiy Storchakabad12572014-12-15 14:03:42 +0200501 for op in [lambda a:copy.copy(a), lambda a:copy.deepcopy(a)] + picklecopiers:
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000502 for data, selectors, result1, result2 in [
503 ('ABCDEF', [1,0,1,0,1,1], 'ACEF', 'CEF'),
504 ('ABCDEF', [0,0,0,0,0,0], '', ''),
505 ('ABCDEF', [1,1,1,1,1,1], 'ABCDEF', 'BCDEF'),
506 ('ABCDEF', [1,0,1], 'AC', 'C'),
507 ('ABC', [0,1,1,1,1,1], 'BC', 'C'),
508 ]:
509
510 self.assertEqual(list(op(compress(data=data, selectors=selectors))), list(result1))
511 self.assertEqual(list(op(compress(data, selectors))), list(result1))
512 testIntermediate = compress(data, selectors)
513 if result1:
514 next(testIntermediate)
515 self.assertEqual(list(op(testIntermediate)), list(result2))
516
517
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000518 def test_count(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000519 self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
520 self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
521 self.assertEqual(take(2, lzip('abc',count(3))), [('a', 3), ('b', 4)])
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000522 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
523 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000524 self.assertRaises(TypeError, count, 2, 3, 4)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000525 self.assertRaises(TypeError, count, 'a')
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300526 self.assertEqual(take(10, count(maxsize-5)),
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000527 list(range(maxsize-5, maxsize+5)))
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300528 self.assertEqual(take(10, count(-maxsize-5)),
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000529 list(range(-maxsize-5, -maxsize+5)))
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300530 self.assertEqual(take(3, count(3.25)), [3.25, 4.25, 5.25])
531 self.assertEqual(take(3, count(3.25-4j)), [3.25-4j, 4.25-4j, 5.25-4j])
532 self.assertEqual(take(3, count(Decimal('1.1'))),
533 [Decimal('1.1'), Decimal('2.1'), Decimal('3.1')])
534 self.assertEqual(take(3, count(Fraction(2, 3))),
535 [Fraction(2, 3), Fraction(5, 3), Fraction(8, 3)])
536 BIGINT = 1<<1000
537 self.assertEqual(take(3, count(BIGINT)), [BIGINT, BIGINT+1, BIGINT+2])
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000538 c = count(3)
539 self.assertEqual(repr(c), 'count(3)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000540 next(c)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000541 self.assertEqual(repr(c), 'count(4)')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000542 c = count(-9)
543 self.assertEqual(repr(c), 'count(-9)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000544 next(c)
545 self.assertEqual(next(c), -8)
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300546 self.assertEqual(repr(count(10.25)), 'count(10.25)')
547 self.assertEqual(repr(count(10.0)), 'count(10.0)')
548 self.assertEqual(type(next(count(10.0))), float)
Christian Heimesa37d4c62007-12-04 23:02:19 +0000549 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
Serhiy Storchaka95949422013-08-27 19:40:23 +0300550 # Test repr
551 r1 = repr(count(i))
552 r2 = 'count(%r)'.__mod__(i)
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000553 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000554
Raymond Hettingera3e1ad22009-11-30 22:02:31 +0000555 # check copy, deepcopy, pickle
556 for value in -3, 3, maxsize-5, maxsize+5:
557 c = count(value)
558 self.assertEqual(next(copy.copy(c)), value)
559 self.assertEqual(next(copy.deepcopy(c)), value)
Serhiy Storchakabad12572014-12-15 14:03:42 +0200560 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
561 self.pickletest(proto, count(value))
Raymond Hettingera3e1ad22009-11-30 22:02:31 +0000562
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +0000563 #check proper internal error handling for large "step' sizes
564 count(1, maxsize+5); sys.exc_info()
565
Raymond Hettinger30729212009-02-12 06:28:27 +0000566 def test_count_with_stride(self):
567 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingereb13fdd2009-02-21 22:30:12 +0000568 self.assertEqual(lzip('abc',count(start=2,step=3)),
569 [('a', 2), ('b', 5), ('c', 8)])
570 self.assertEqual(lzip('abc',count(step=-1)),
571 [('a', 0), ('b', -1), ('c', -2)])
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300572 self.assertRaises(TypeError, count, 'a', 'b')
Raymond Hettinger30729212009-02-12 06:28:27 +0000573 self.assertEqual(lzip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
574 self.assertEqual(lzip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000575 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000576 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
577 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300578 self.assertEqual(take(3, count(10, maxsize+5)),
579 list(range(10, 10+3*(maxsize+5), maxsize+5)))
580 self.assertEqual(take(3, count(2, 1.25)), [2, 3.25, 4.5])
Raymond Hettinger30729212009-02-12 06:28:27 +0000581 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettinger8a882a72009-02-12 12:08:18 +0000582 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
583 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerde09acf2009-02-12 12:53:53 +0000584 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
585 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300586 BIGINT = 1<<1000
587 self.assertEqual(take(3, count(step=BIGINT)), [0, BIGINT, 2*BIGINT])
Raymond Hettinger30729212009-02-12 06:28:27 +0000588 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
589 c = count(3, 5)
590 self.assertEqual(repr(c), 'count(3, 5)')
591 next(c)
592 self.assertEqual(repr(c), 'count(8, 5)')
593 c = count(-9, 0)
594 self.assertEqual(repr(c), 'count(-9, 0)')
595 next(c)
596 self.assertEqual(repr(c), 'count(-9, 0)')
597 c = count(-9, -3)
598 self.assertEqual(repr(c), 'count(-9, -3)')
599 next(c)
600 self.assertEqual(repr(c), 'count(-12, -3)')
601 self.assertEqual(repr(c), 'count(-12, -3)')
602 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
603 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
604 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +0300605 self.assertEqual(repr(count(10, 1.00)), 'count(10, 1.0)')
606 c = count(10, 1.0)
607 self.assertEqual(type(next(c)), int)
608 self.assertEqual(type(next(c)), float)
Raymond Hettinger30729212009-02-12 06:28:27 +0000609 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
610 for j in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 1, 10, sys.maxsize-5, sys.maxsize+5):
Serhiy Storchaka95949422013-08-27 19:40:23 +0300611 # Test repr
612 r1 = repr(count(i, j))
Raymond Hettinger30729212009-02-12 06:28:27 +0000613 if j == 1:
Serhiy Storchaka95949422013-08-27 19:40:23 +0300614 r2 = ('count(%r)' % i)
Raymond Hettinger30729212009-02-12 06:28:27 +0000615 else:
Serhiy Storchaka95949422013-08-27 19:40:23 +0300616 r2 = ('count(%r, %r)' % (i, j))
Raymond Hettinger30729212009-02-12 06:28:27 +0000617 self.assertEqual(r1, r2)
Serhiy Storchakabad12572014-12-15 14:03:42 +0200618 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
619 self.pickletest(proto, count(i, j))
Raymond Hettinger30729212009-02-12 06:28:27 +0000620
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000621 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000622 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000623 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000624 self.assertRaises(TypeError, cycle)
625 self.assertRaises(TypeError, cycle, 5)
626 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000627
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000628 # check copy, deepcopy, pickle
629 c = cycle('abc')
630 self.assertEqual(next(c), 'a')
631 #simple copy currently not supported, because __reduce__ returns
632 #an internal iterator
633 #self.assertEqual(take(10, copy.copy(c)), list('bcabcabcab'))
634 self.assertEqual(take(10, copy.deepcopy(c)), list('bcabcabcab'))
Serhiy Storchakabad12572014-12-15 14:03:42 +0200635 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
636 self.assertEqual(take(10, pickle.loads(pickle.dumps(c, proto))),
637 list('bcabcabcab'))
638 next(c)
639 self.assertEqual(take(10, pickle.loads(pickle.dumps(c, proto))),
640 list('cabcabcabc'))
641 next(c)
642 next(c)
643 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
644 self.pickletest(proto, cycle('abc'))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000645
Raymond Hettingera166ce52015-08-15 14:45:49 -0700646 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
647 # test with partial consumed input iterable
648 it = iter('abcde')
649 c = cycle(it)
Raymond Hettinger1cadf762015-08-15 14:47:27 -0700650 _ = [next(c) for i in range(2)] # consume 2 of 5 inputs
Raymond Hettingera166ce52015-08-15 14:45:49 -0700651 p = pickle.dumps(c, proto)
652 d = pickle.loads(p) # rebuild the cycle object
653 self.assertEqual(take(20, d), list('cdeabcdeabcdeabcdeab'))
654
655 # test with completely consumed input iterable
656 it = iter('abcde')
657 c = cycle(it)
Raymond Hettinger1cadf762015-08-15 14:47:27 -0700658 _ = [next(c) for i in range(7)] # consume 7 of 5 inputs
Raymond Hettingera166ce52015-08-15 14:45:49 -0700659 p = pickle.dumps(c, proto)
660 d = pickle.loads(p) # rebuild the cycle object
661 self.assertEqual(take(20, d), list('cdeabcdeabcdeabcdeab'))
662
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700663 def test_cycle_setstate(self):
664 # Verify both modes for restoring state
665
666 # Mode 0 is efficient. It uses an incompletely consumed input
667 # iterator to build a cycle object and then passes in state with
668 # a list of previously consumed values. There is no data
Martin Panterf1579822016-05-26 06:03:33 +0000669 # overlap between the two.
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700670 c = cycle('defg')
671 c.__setstate__((list('abc'), 0))
672 self.assertEqual(take(20, c), list('defgabcdefgabcdefgab'))
673
674 # Mode 1 is inefficient. It starts with a cycle object built
675 # from an iterator over the remaining elements in a partial
676 # cycle and then passes in state with all of the previously
677 # seen values (this overlaps values included in the iterator).
678 c = cycle('defg')
679 c.__setstate__((list('abcdefg'), 1))
680 self.assertEqual(take(20, c), list('defgabcdefgabcdefgab'))
681
682 # The first argument to setstate needs to be a tuple
Serhiy Storchaka8f0f2052016-10-02 09:13:14 +0300683 with self.assertRaises(TypeError):
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700684 cycle('defg').__setstate__([list('abcdefg'), 0])
685
686 # The first argument in the setstate tuple must be a list
687 with self.assertRaises(TypeError):
688 c = cycle('defg')
Serhiy Storchaka8f0f2052016-10-02 09:13:14 +0300689 c.__setstate__((tuple('defg'), 0))
690 take(20, c)
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700691
Serhiy Storchaka8f0f2052016-10-02 09:13:14 +0300692 # The second argument in the setstate tuple must be an int
Raymond Hettinger79c878d2015-08-15 13:51:59 -0700693 with self.assertRaises(TypeError):
694 cycle('defg').__setstate__((list('abcdefg'), 'x'))
695
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300696 self.assertRaises(TypeError, cycle('').__setstate__, ())
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300697 self.assertRaises(TypeError, cycle('').__setstate__, ([],))
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300698
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000699 def test_groupby(self):
700 # Check whether it accepts arguments correctly
701 self.assertEqual([], list(groupby([])))
702 self.assertEqual([], list(groupby([], key=id)))
703 self.assertRaises(TypeError, list, groupby('abc', []))
704 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000705 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000706
707 # Check normal input
708 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
709 (2,15,22), (3,16,23), (3,17,23)]
710 dup = []
711 for k, g in groupby(s, lambda r:r[0]):
712 for elem in g:
713 self.assertEqual(k, elem[0])
714 dup.append(elem)
715 self.assertEqual(s, dup)
716
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000717 # Check normal pickled
Serhiy Storchakabad12572014-12-15 14:03:42 +0200718 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
719 dup = []
720 for k, g in pickle.loads(pickle.dumps(groupby(s, testR), proto)):
721 for elem in g:
722 self.assertEqual(k, elem[0])
723 dup.append(elem)
724 self.assertEqual(s, dup)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000725
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000726 # Check nested case
727 dup = []
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000728 for k, g in groupby(s, testR):
729 for ik, ig in groupby(g, testR2):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000730 for elem in ig:
731 self.assertEqual(k, elem[0])
732 self.assertEqual(ik, elem[2])
733 dup.append(elem)
734 self.assertEqual(s, dup)
735
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000736 # Check nested and pickled
Serhiy Storchakabad12572014-12-15 14:03:42 +0200737 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
738 dup = []
739 for k, g in pickle.loads(pickle.dumps(groupby(s, testR), proto)):
740 for ik, ig in pickle.loads(pickle.dumps(groupby(g, testR2), proto)):
741 for elem in ig:
742 self.assertEqual(k, elem[0])
743 self.assertEqual(ik, elem[2])
744 dup.append(elem)
745 self.assertEqual(s, dup)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000746
747
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000748 # Check case where inner iterator is not used
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000749 keys = [k for k, g in groupby(s, testR)]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000750 expectedkeys = set([r[0] for r in s])
751 self.assertEqual(set(keys), expectedkeys)
752 self.assertEqual(len(keys), len(expectedkeys))
753
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300754 # Check case where inner iterator is used after advancing the groupby
755 # iterator
756 s = list(zip('AABBBAAAA', range(9)))
757 it = groupby(s, testR)
758 _, g1 = next(it)
759 _, g2 = next(it)
760 _, g3 = next(it)
761 self.assertEqual(list(g1), [])
762 self.assertEqual(list(g2), [])
763 self.assertEqual(next(g3), ('A', 5))
764 list(it) # exhaust the groupby iterator
765 self.assertEqual(list(g3), [])
766
767 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
768 it = groupby(s, testR)
769 _, g = next(it)
770 next(it)
771 next(it)
772 self.assertEqual(list(pickle.loads(pickle.dumps(g, proto))), [])
773
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000774 # Exercise pipes and filters style
775 s = 'abracadabra'
776 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000777 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000778 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
779 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000780 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000781 self.assertEqual(r, ['a', 'b', 'r'])
782 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000783 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000784 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
785 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000786 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000787 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
788
Georg Brandla18af4e2007-04-21 15:47:16 +0000789 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000790 class ExpectedError(Exception):
791 pass
792 def delayed_raise(n=0):
793 for i in range(n):
794 yield 'yo'
795 raise ExpectedError
796 def gulp(iterable, keyp=None, func=list):
797 return [func(g) for k, g in groupby(iterable, keyp)]
798
Georg Brandla18af4e2007-04-21 15:47:16 +0000799 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000800 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000801 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000802 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
803
Serhiy Storchakaa60c2fe2015-03-12 21:56:08 +0200804 # __eq__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000805 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000806 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000807 raise ExpectedError
808 s = [DummyCmp(), DummyCmp(), None]
809
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000810 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000811 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000812 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000813 self.assertRaises(ExpectedError, gulp, s)
814
815 # keyfunc failure
816 def keyfunc(obj):
817 if keyfunc.skip > 0:
818 keyfunc.skip -= 1
819 return obj
820 else:
821 raise ExpectedError
822
823 # keyfunc failure on outer object
824 keyfunc.skip = 0
825 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
826 keyfunc.skip = 1
827 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
828
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000829 def test_filter(self):
830 self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
831 self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])
832 self.assertEqual(list(filter(bool, [0,1,0,2,0])), [1,2])
833 self.assertEqual(take(4, filter(isEven, count())), [0,2,4,6])
834 self.assertRaises(TypeError, filter)
835 self.assertRaises(TypeError, filter, lambda x:x)
836 self.assertRaises(TypeError, filter, lambda x:x, range(6), 7)
837 self.assertRaises(TypeError, filter, isEven, 3)
838 self.assertRaises(TypeError, next, filter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000839
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000840 # check copy, deepcopy, pickle
841 ans = [0,2,4]
842
843 c = filter(isEven, range(6))
844 self.assertEqual(list(copy.copy(c)), ans)
845 c = filter(isEven, range(6))
846 self.assertEqual(list(copy.deepcopy(c)), ans)
Serhiy Storchakabad12572014-12-15 14:03:42 +0200847 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
848 c = filter(isEven, range(6))
849 self.assertEqual(list(pickle.loads(pickle.dumps(c, proto))), ans)
850 next(c)
851 self.assertEqual(list(pickle.loads(pickle.dumps(c, proto))), ans[1:])
852 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
853 c = filter(isEven, range(6))
854 self.pickletest(proto, c)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000855
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000856 def test_filterfalse(self):
857 self.assertEqual(list(filterfalse(isEven, range(6))), [1,3,5])
858 self.assertEqual(list(filterfalse(None, [0,1,0,2,0])), [0,0,0])
859 self.assertEqual(list(filterfalse(bool, [0,1,0,2,0])), [0,0,0])
860 self.assertEqual(take(4, filterfalse(isEven, count())), [1,3,5,7])
861 self.assertRaises(TypeError, filterfalse)
862 self.assertRaises(TypeError, filterfalse, lambda x:x)
863 self.assertRaises(TypeError, filterfalse, lambda x:x, range(6), 7)
864 self.assertRaises(TypeError, filterfalse, isEven, 3)
865 self.assertRaises(TypeError, next, filterfalse(range(6), range(6)))
Serhiy Storchakabad12572014-12-15 14:03:42 +0200866 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
867 self.pickletest(proto, filterfalse(isEven, range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000868
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000869 def test_zip(self):
870 # XXX This is rather silly now that builtin zip() calls zip()...
871 ans = [(x,y) for x, y in zip('abc',count())]
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000872 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000873 self.assertEqual(list(zip('abc', range(6))), lzip('abc', range(6)))
874 self.assertEqual(list(zip('abcdef', range(3))), lzip('abcdef', range(3)))
875 self.assertEqual(take(3,zip('abcdef', count())), lzip('abcdef', range(3)))
876 self.assertEqual(list(zip('abcdef')), lzip('abcdef'))
877 self.assertEqual(list(zip()), lzip())
878 self.assertRaises(TypeError, zip, 3)
879 self.assertRaises(TypeError, zip, range(3), 3)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000880 self.assertEqual([tuple(list(pair)) for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000881 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000882 self.assertEqual([pair for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000883 lzip('abc', 'def'))
Alex Gaynore151d212011-07-17 16:21:30 -0700884
885 @support.impl_detail("tuple reuse is specific to CPython")
886 def test_zip_tuple_reuse(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000887 ids = list(map(id, zip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000888 self.assertEqual(min(ids), max(ids))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000889 ids = list(map(id, list(zip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000890 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000891
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000892 # check copy, deepcopy, pickle
893 ans = [(x,y) for x, y in copy.copy(zip('abc',count()))]
894 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
895
896 ans = [(x,y) for x, y in copy.deepcopy(zip('abc',count()))]
897 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
898
Serhiy Storchakabad12572014-12-15 14:03:42 +0200899 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
900 ans = [(x,y) for x, y in pickle.loads(pickle.dumps(zip('abc',count()), proto))]
901 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000902
Serhiy Storchakabad12572014-12-15 14:03:42 +0200903 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
904 testIntermediate = zip('abc',count())
905 next(testIntermediate)
906 ans = [(x,y) for x, y in pickle.loads(pickle.dumps(testIntermediate, proto))]
907 self.assertEqual(ans, [('b', 1), ('c', 2)])
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000908
Serhiy Storchakabad12572014-12-15 14:03:42 +0200909 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
910 self.pickletest(proto, zip('abc', count()))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000911
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000912 def test_ziplongest(self):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000913 for args in [
914 ['abc', range(6)],
915 [range(6), 'abc'],
916 [range(1000), range(2000,2100), range(3000,3050)],
917 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
918 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
919 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000920 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
921 for i in range(max(map(len, args)))]
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000922 self.assertEqual(list(zip_longest(*args)), target)
923 self.assertEqual(list(zip_longest(*args, **{})), target)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000924 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 +0000925 self.assertEqual(list(zip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000926
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000927 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 +0000928
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000929 self.assertEqual(list(zip_longest()), list(zip()))
930 self.assertEqual(list(zip_longest([])), list(zip([])))
931 self.assertEqual(list(zip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000932
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000933 self.assertEqual(list(zip_longest('abc', 'defg', **{})),
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000934 list(zip(list('abc')+[None], 'defg'))) # empty keyword dict
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000935 self.assertRaises(TypeError, zip_longest, 3)
936 self.assertRaises(TypeError, zip_longest, range(3), 3)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000937
938 for stmt in [
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000939 "zip_longest('abc', fv=1)",
940 "zip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000941 ]:
942 try:
943 eval(stmt, globals(), locals())
944 except TypeError:
945 pass
946 else:
947 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000948
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000949 self.assertEqual([tuple(list(pair)) for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000950 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000951 self.assertEqual([pair for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000952 list(zip('abc', 'def')))
Alex Gaynore151d212011-07-17 16:21:30 -0700953
954 @support.impl_detail("tuple reuse is specific to CPython")
955 def test_zip_longest_tuple_reuse(self):
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000956 ids = list(map(id, zip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000957 self.assertEqual(min(ids), max(ids))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000958 ids = list(map(id, list(zip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000959 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
960
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000961 def test_zip_longest_pickling(self):
Serhiy Storchakabad12572014-12-15 14:03:42 +0200962 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
963 self.pickletest(proto, zip_longest("abc", "def"))
964 self.pickletest(proto, zip_longest("abc", "defgh"))
965 self.pickletest(proto, zip_longest("abc", "defgh", fillvalue=1))
966 self.pickletest(proto, zip_longest("", "defgh"))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000967
Raymond Hettingerfc438512009-11-01 20:55:33 +0000968 def test_bug_7244(self):
969
970 class Repeater:
971 # this class is similar to itertools.repeat
972 def __init__(self, o, t, e):
973 self.o = o
974 self.t = int(t)
975 self.e = e
976 def __iter__(self): # its iterator is itself
977 return self
978 def __next__(self):
979 if self.t > 0:
980 self.t -= 1
981 return self.o
982 else:
983 raise self.e
984
985 # Formerly this code in would fail in debug mode
986 # with Undetected Error and Stop Iteration
987 r1 = Repeater(1, 3, StopIteration)
988 r2 = Repeater(2, 4, StopIteration)
989 def run(r1, r2):
990 result = []
991 for i, j in zip_longest(r1, r2, fillvalue=0):
992 with support.captured_output('stdout'):
993 print((i, j))
994 result.append((i, j))
995 return result
996 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
997
998 # Formerly, the RuntimeError would be lost
999 # and StopIteration would stop as expected
1000 r1 = Repeater(1, 3, RuntimeError)
1001 r2 = Repeater(2, 4, StopIteration)
1002 it = zip_longest(r1, r2, fillvalue=0)
1003 self.assertEqual(next(it), (1, 2))
1004 self.assertEqual(next(it), (1, 2))
1005 self.assertEqual(next(it), (1, 2))
1006 self.assertRaises(RuntimeError, next, it)
1007
Christian Heimesc3f30c42008-02-22 16:37:40 +00001008 def test_product(self):
1009 for args, result in [
Christian Heimes78644762008-03-04 23:39:23 +00001010 ([], [()]), # zero iterables
Christian Heimesc3f30c42008-02-22 16:37:40 +00001011 (['ab'], [('a',), ('b',)]), # one iterable
1012 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
1013 ([range(0), range(2), range(3)], []), # first iterable with zero length
1014 ([range(2), range(0), range(3)], []), # middle iterable with zero length
1015 ([range(2), range(3), range(0)], []), # last iterable with zero length
1016 ]:
1017 self.assertEqual(list(product(*args)), result)
Christian Heimesf16baeb2008-02-29 14:57:44 +00001018 for r in range(4):
1019 self.assertEqual(list(product(*(args*r))),
1020 list(product(*args, **dict(repeat=r))))
Christian Heimesc3f30c42008-02-22 16:37:40 +00001021 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
1022 self.assertRaises(TypeError, product, range(6), None)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001023
1024 def product1(*args, **kwds):
1025 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
1026 n = len(pools)
1027 if n == 0:
1028 yield ()
1029 return
1030 if any(len(pool) == 0 for pool in pools):
1031 return
1032 indices = [0] * n
1033 yield tuple(pool[i] for pool, i in zip(pools, indices))
1034 while 1:
1035 for i in reversed(range(n)): # right to left
1036 if indices[i] == len(pools[i]) - 1:
1037 continue
1038 indices[i] += 1
1039 for j in range(i+1, n):
1040 indices[j] = 0
1041 yield tuple(pool[i] for pool, i in zip(pools, indices))
1042 break
1043 else:
1044 return
1045
1046 def product2(*args, **kwds):
1047 'Pure python version used in docs'
1048 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
1049 result = [[]]
1050 for pool in pools:
1051 result = [x+[y] for x in result for y in pool]
1052 for prod in result:
1053 yield tuple(prod)
1054
Christian Heimesc3f30c42008-02-22 16:37:40 +00001055 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
1056 set('abcdefg'), range(11), tuple(range(13))]
1057 for i in range(100):
1058 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Christian Heimes78644762008-03-04 23:39:23 +00001059 expected_len = prod(map(len, args))
1060 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001061 self.assertEqual(list(product(*args)), list(product1(*args)))
1062 self.assertEqual(list(product(*args)), list(product2(*args)))
Christian Heimesc3f30c42008-02-22 16:37:40 +00001063 args = map(iter, args)
Christian Heimes78644762008-03-04 23:39:23 +00001064 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesc3f30c42008-02-22 16:37:40 +00001065
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05001066 @support.bigaddrspacetest
1067 def test_product_overflow(self):
Serhiy Storchakadee948b2015-02-03 01:34:09 +02001068 with self.assertRaises((OverflowError, MemoryError)):
1069 product(*(['ab']*2**5), repeat=2**25)
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05001070
Alex Gaynore151d212011-07-17 16:21:30 -07001071 @support.impl_detail("tuple reuse is specific to CPython")
1072 def test_product_tuple_reuse(self):
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001073 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
1074 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +00001075
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001076 def test_product_pickling(self):
1077 # check copy, deepcopy, pickle
1078 for args, result in [
1079 ([], [()]), # zero iterables
1080 (['ab'], [('a',), ('b',)]), # one iterable
1081 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
1082 ([range(0), range(2), range(3)], []), # first iterable with zero length
1083 ([range(2), range(0), range(3)], []), # middle iterable with zero length
1084 ([range(2), range(3), range(0)], []), # last iterable with zero length
1085 ]:
1086 self.assertEqual(list(copy.copy(product(*args))), result)
1087 self.assertEqual(list(copy.deepcopy(product(*args))), result)
Serhiy Storchakabad12572014-12-15 14:03:42 +02001088 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1089 self.pickletest(proto, product(*args))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001090
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00001091 def test_product_issue_25021(self):
1092 # test that indices are properly clamped to the length of the tuples
1093 p = product((1, 2),(3,))
1094 p.__setstate__((0, 0x1000)) # will access tuple element 1 if not clamped
1095 self.assertEqual(next(p), (2, 3))
1096 # test that empty tuple in the list will result in an immediate StopIteration
1097 p = product((1, 2), (), (3,))
1098 p.__setstate__((0, 0, 0x1000)) # will access tuple element 1 if not clamped
1099 self.assertRaises(StopIteration, next, p)
1100
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001101 def test_repeat(self):
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00001102 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
Guido van Rossum805365e2007-05-07 22:24:25 +00001103 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001104 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001105 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +00001106 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001107 self.assertEqual(list(repeat('a', 0)), [])
1108 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001109 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001110 self.assertRaises(TypeError, repeat, None, 3, 4)
1111 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001112 r = repeat(1+0j)
1113 self.assertEqual(repr(r), 'repeat((1+0j))')
1114 r = repeat(1+0j, 5)
1115 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
1116 list(r)
1117 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001118
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001119 # check copy, deepcopy, pickle
1120 c = repeat(object='a', times=10)
1121 self.assertEqual(next(c), 'a')
1122 self.assertEqual(take(2, copy.copy(c)), list('a' * 2))
1123 self.assertEqual(take(2, copy.deepcopy(c)), list('a' * 2))
Serhiy Storchakabad12572014-12-15 14:03:42 +02001124 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1125 self.pickletest(proto, repeat(object='a', times=10))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001126
Raymond Hettinger97d35552014-06-24 21:36:58 -07001127 def test_repeat_with_negative_times(self):
1128 self.assertEqual(repr(repeat('a', -1)), "repeat('a', 0)")
1129 self.assertEqual(repr(repeat('a', -2)), "repeat('a', 0)")
1130 self.assertEqual(repr(repeat('a', times=-1)), "repeat('a', 0)")
1131 self.assertEqual(repr(repeat('a', times=-2)), "repeat('a', 0)")
1132
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001133 def test_map(self):
1134 self.assertEqual(list(map(operator.pow, range(3), range(1,7))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001135 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001136 self.assertEqual(list(map(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001137 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001138 self.assertEqual(list(map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +00001139 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001140 self.assertEqual(take(2,map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +00001141 [('a',0),('b',1)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001142 self.assertEqual(list(map(operator.pow, [])), [])
1143 self.assertRaises(TypeError, map)
1144 self.assertRaises(TypeError, list, map(None, range(3), range(3)))
1145 self.assertRaises(TypeError, map, operator.neg)
1146 self.assertRaises(TypeError, next, map(10, range(5)))
1147 self.assertRaises(ValueError, next, map(errfunc, [4], [5]))
1148 self.assertRaises(TypeError, next, map(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001149
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001150 # check copy, deepcopy, pickle
1151 ans = [('a',0),('b',1),('c',2)]
1152
1153 c = map(tupleize, 'abc', count())
1154 self.assertEqual(list(copy.copy(c)), ans)
1155
1156 c = map(tupleize, 'abc', count())
1157 self.assertEqual(list(copy.deepcopy(c)), ans)
1158
Serhiy Storchakabad12572014-12-15 14:03:42 +02001159 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1160 c = map(tupleize, 'abc', count())
1161 self.pickletest(proto, c)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001162
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001163 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001164 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
1165 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001166 self.assertEqual(take(3, starmap(operator.pow, zip(count(), count(1)))),
Raymond Hettinger02420702003-06-29 20:36:23 +00001167 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001168 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +00001169 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
1170 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001171 self.assertRaises(TypeError, starmap)
1172 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001173 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
1174 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
1175 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001176
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001177 # check copy, deepcopy, pickle
1178 ans = [0**1, 1**2, 2**3]
1179
1180 c = starmap(operator.pow, zip(range(3), range(1,7)))
1181 self.assertEqual(list(copy.copy(c)), ans)
1182
1183 c = starmap(operator.pow, zip(range(3), range(1,7)))
1184 self.assertEqual(list(copy.deepcopy(c)), ans)
1185
Serhiy Storchakabad12572014-12-15 14:03:42 +02001186 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1187 c = starmap(operator.pow, zip(range(3), range(1,7)))
1188 self.pickletest(proto, c)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001189
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001190 def test_islice(self):
1191 for args in [ # islice(args) should agree with range(args)
1192 (10, 20, 3),
1193 (10, 3, 20),
1194 (10, 20),
Cheryl Sabellada1734c2018-03-26 21:29:33 -04001195 (10, 10),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001196 (10, 3),
1197 (20,)
1198 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +00001199 self.assertEqual(list(islice(range(100), *args)),
1200 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001201
1202 for args, tgtargs in [ # Stop when seqn is exhausted
1203 ((10, 110, 3), ((10, 100, 3))),
1204 ((10, 110), ((10, 100))),
1205 ((110,), (100,))
1206 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +00001207 self.assertEqual(list(islice(range(100), *args)),
1208 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001209
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001210 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +00001211 self.assertEqual(list(islice(range(10), None)), list(range(10)))
1212 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
1213 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
1214 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
1215 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001216
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +00001217 # Test number of items consumed SF #1171417
1218 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +00001219 self.assertEqual(list(islice(it, 3)), list(range(3)))
1220 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +00001221
Cheryl Sabellada1734c2018-03-26 21:29:33 -04001222 it = iter(range(10))
1223 self.assertEqual(list(islice(it, 3, 3)), [])
1224 self.assertEqual(list(it), list(range(3, 10)))
1225
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001226 # Test invalid arguments
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001227 ra = range(10)
1228 self.assertRaises(TypeError, islice, ra)
1229 self.assertRaises(TypeError, islice, ra, 1, 2, 3, 4)
1230 self.assertRaises(ValueError, islice, ra, -5, 10, 1)
1231 self.assertRaises(ValueError, islice, ra, 1, -5, -1)
1232 self.assertRaises(ValueError, islice, ra, 1, 10, -1)
1233 self.assertRaises(ValueError, islice, ra, 1, 10, 0)
1234 self.assertRaises(ValueError, islice, ra, 'a')
1235 self.assertRaises(ValueError, islice, ra, 'a', 1)
1236 self.assertRaises(ValueError, islice, ra, 1, 'a')
1237 self.assertRaises(ValueError, islice, ra, 'a', 1, 1)
1238 self.assertRaises(ValueError, islice, ra, 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +00001239 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001240
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001241 # Issue #10323: Less islice in a predictable state
1242 c = count()
1243 self.assertEqual(list(islice(c, 1, 3, 50)), [1])
1244 self.assertEqual(next(c), 3)
1245
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001246 # check copy, deepcopy, pickle
1247 for args in [ # islice(args) should agree with range(args)
1248 (10, 20, 3),
1249 (10, 3, 20),
1250 (10, 20),
1251 (10, 3),
1252 (20,)
1253 ]:
1254 self.assertEqual(list(copy.copy(islice(range(100), *args))),
1255 list(range(*args)))
1256 self.assertEqual(list(copy.deepcopy(islice(range(100), *args))),
1257 list(range(*args)))
Serhiy Storchakabad12572014-12-15 14:03:42 +02001258 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1259 self.pickletest(proto, islice(range(100), *args))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001260
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001261 # Issue #21321: check source iterator is not referenced
1262 # from islice() after the latter has been exhausted
1263 it = (x for x in (1, 2))
1264 wr = weakref.ref(it)
1265 it = islice(it, 1)
1266 self.assertIsNotNone(wr())
1267 list(it) # exhaust the iterator
Benjamin Peterson8e163512014-08-24 18:07:28 -05001268 support.gc_collect()
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001269 self.assertIsNone(wr())
1270
Will Roberts0ecdc522017-06-08 08:03:04 +02001271 # Issue #30537: islice can accept integer-like objects as
1272 # arguments
1273 class IntLike(object):
1274 def __init__(self, val):
1275 self.val = val
1276 def __index__(self):
1277 return self.val
1278 self.assertEqual(list(islice(range(100), IntLike(10))), list(range(10)))
1279 self.assertEqual(list(islice(range(100), IntLike(10), IntLike(50))),
1280 list(range(10, 50)))
1281 self.assertEqual(list(islice(range(100), IntLike(10), IntLike(50), IntLike(5))),
1282 list(range(10,50,5)))
1283
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001284 def test_takewhile(self):
1285 data = [1, 3, 5, 20, 2, 4, 6, 8]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001286 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001287 self.assertEqual(list(takewhile(underten, [])), [])
1288 self.assertRaises(TypeError, takewhile)
1289 self.assertRaises(TypeError, takewhile, operator.pow)
1290 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001291 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
1292 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001293 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
1294 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +00001295 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001296
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001297 # check copy, deepcopy, pickle
1298 self.assertEqual(list(copy.copy(takewhile(underten, data))), [1, 3, 5])
1299 self.assertEqual(list(copy.deepcopy(takewhile(underten, data))),
1300 [1, 3, 5])
Serhiy Storchakabad12572014-12-15 14:03:42 +02001301 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1302 self.pickletest(proto, takewhile(underten, data))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001303
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001304 def test_dropwhile(self):
1305 data = [1, 3, 5, 20, 2, 4, 6, 8]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001306 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001307 self.assertEqual(list(dropwhile(underten, [])), [])
1308 self.assertRaises(TypeError, dropwhile)
1309 self.assertRaises(TypeError, dropwhile, operator.pow)
1310 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001311 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
1312 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001313
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001314 # check copy, deepcopy, pickle
1315 self.assertEqual(list(copy.copy(dropwhile(underten, data))), [20, 2, 4, 6, 8])
1316 self.assertEqual(list(copy.deepcopy(dropwhile(underten, data))),
1317 [20, 2, 4, 6, 8])
Serhiy Storchakabad12572014-12-15 14:03:42 +02001318 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1319 self.pickletest(proto, dropwhile(underten, data))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001320
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001321 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +00001322 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001323
1324 a, b = tee([]) # test empty iterator
1325 self.assertEqual(list(a), [])
1326 self.assertEqual(list(b), [])
1327
1328 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +00001329 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001330
1331 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +00001332 self.assertEqual(list(a), list(range(n)))
1333 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001334
1335 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +00001336 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001337 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001338 del a
Guido van Rossum805365e2007-05-07 22:24:25 +00001339 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001340
1341 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +00001342 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001343 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001344 del b
Guido van Rossum805365e2007-05-07 22:24:25 +00001345 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001346
Guido van Rossum805365e2007-05-07 22:24:25 +00001347 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001348 order = [0]*n + [1]*n
1349 random.shuffle(order)
1350 lists = ([], [])
1351 its = tee(irange(n))
1352 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +00001353 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001354 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +00001355 self.assertEqual(lists[0], list(range(n)))
1356 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001357
Raymond Hettingerad983e72003-11-12 14:32:26 +00001358 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001359 self.assertRaises(TypeError, tee)
1360 self.assertRaises(TypeError, tee, 3)
1361 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +00001362 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001363
Raymond Hettingerad983e72003-11-12 14:32:26 +00001364 # tee object should be instantiable
1365 a, b = tee('abc')
1366 c = type(a)('def')
1367 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +00001368
Raymond Hettingerad983e72003-11-12 14:32:26 +00001369 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +00001370 a, b, c = tee(range(2000), 3)
1371 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001372 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +00001373 self.assertEqual(list(b), list(range(2000)))
1374 self.assertEqual([next(c), next(c)], list(range(2)))
1375 self.assertEqual(list(a), list(range(100,2000)))
1376 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +00001377
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001378 # test values of n
1379 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +00001380 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +00001381 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001382 result = tee('abc', n)
1383 self.assertEqual(type(result), tuple)
1384 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001385 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001386
Raymond Hettingerad983e72003-11-12 14:32:26 +00001387 # tee pass-through to copyable iterator
1388 a, b = tee('abc')
1389 c, d = tee(a)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +00001390 self.assertTrue(a is c)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001391
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001392 # test tee_new
1393 t1, t2 = tee('abc')
1394 tnew = type(t1)
1395 self.assertRaises(TypeError, tnew)
1396 self.assertRaises(TypeError, tnew, 10)
1397 t3 = tnew(t1)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +00001398 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +00001399
Raymond Hettingera9f60922004-10-17 16:40:14 +00001400 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +00001401 a, b = tee(range(10))
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001402 p = weakref.proxy(a)
Raymond Hettingera9f60922004-10-17 16:40:14 +00001403 self.assertEqual(getattr(p, '__class__'), type(b))
1404 del a
1405 self.assertRaises(ReferenceError, getattr, p, '__class__')
1406
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001407 ans = list('abc')
1408 long_ans = list(range(10000))
1409
1410 # check copy
1411 a, b = tee('abc')
1412 self.assertEqual(list(copy.copy(a)), ans)
1413 self.assertEqual(list(copy.copy(b)), ans)
1414 a, b = tee(list(range(10000)))
1415 self.assertEqual(list(copy.copy(a)), long_ans)
1416 self.assertEqual(list(copy.copy(b)), long_ans)
1417
1418 # check partially consumed copy
1419 a, b = tee('abc')
1420 take(2, a)
1421 take(1, b)
1422 self.assertEqual(list(copy.copy(a)), ans[2:])
1423 self.assertEqual(list(copy.copy(b)), ans[1:])
1424 self.assertEqual(list(a), ans[2:])
1425 self.assertEqual(list(b), ans[1:])
1426 a, b = tee(range(10000))
1427 take(100, a)
1428 take(60, b)
1429 self.assertEqual(list(copy.copy(a)), long_ans[100:])
1430 self.assertEqual(list(copy.copy(b)), long_ans[60:])
1431 self.assertEqual(list(a), long_ans[100:])
1432 self.assertEqual(list(b), long_ans[60:])
1433
1434 # check deepcopy
1435 a, b = tee('abc')
1436 self.assertEqual(list(copy.deepcopy(a)), ans)
1437 self.assertEqual(list(copy.deepcopy(b)), ans)
1438 self.assertEqual(list(a), ans)
1439 self.assertEqual(list(b), ans)
1440 a, b = tee(range(10000))
1441 self.assertEqual(list(copy.deepcopy(a)), long_ans)
1442 self.assertEqual(list(copy.deepcopy(b)), long_ans)
1443 self.assertEqual(list(a), long_ans)
1444 self.assertEqual(list(b), long_ans)
1445
1446 # check partially consumed deepcopy
1447 a, b = tee('abc')
1448 take(2, a)
1449 take(1, b)
1450 self.assertEqual(list(copy.deepcopy(a)), ans[2:])
1451 self.assertEqual(list(copy.deepcopy(b)), ans[1:])
1452 self.assertEqual(list(a), ans[2:])
1453 self.assertEqual(list(b), ans[1:])
1454 a, b = tee(range(10000))
1455 take(100, a)
1456 take(60, b)
1457 self.assertEqual(list(copy.deepcopy(a)), long_ans[100:])
1458 self.assertEqual(list(copy.deepcopy(b)), long_ans[60:])
1459 self.assertEqual(list(a), long_ans[100:])
1460 self.assertEqual(list(b), long_ans[60:])
1461
1462 # check pickle
Serhiy Storchakabad12572014-12-15 14:03:42 +02001463 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1464 self.pickletest(proto, iter(tee('abc')))
1465 a, b = tee('abc')
1466 self.pickletest(proto, a, compare=ans)
1467 self.pickletest(proto, b, compare=ans)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001468
Serhiy Storchakaa3e91282013-01-25 13:19:31 +02001469 # Issue 13454: Crash when deleting backward iterator from tee()
1470 def test_tee_del_backward(self):
Serhiy Storchaka5bb893c2013-01-26 11:52:06 +02001471 forward, backward = tee(repeat(None, 20000000))
Serhiy Storchaka9db55002015-03-28 20:38:37 +02001472 try:
1473 any(forward) # exhaust the iterator
1474 del backward
1475 except:
1476 del forward, backward
1477 raise
Serhiy Storchakaa3e91282013-01-25 13:19:31 +02001478
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001479 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001480 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +00001481
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001482 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +00001483 self.assertRaises(StopIteration, next, f([]))
1484 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001485
Georg Brandla18af4e2007-04-21 15:47:16 +00001486 self.assertRaises(StopIteration, next, islice([], None))
1487 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001488
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001489 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +00001490 self.assertRaises(StopIteration, next, p)
1491 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001492 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +00001493 self.assertRaises(StopIteration, next, p)
1494 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001495
Georg Brandla18af4e2007-04-21 15:47:16 +00001496 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001497
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001498 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +00001499 self.assertRaises(StopIteration, next, f(lambda x:x, []))
1500 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001501
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001502class TestExamples(unittest.TestCase):
1503
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001504 def test_accumulate(self):
Raymond Hettinger482ba772010-12-01 22:48:00 +00001505 self.assertEqual(list(accumulate([1,2,3,4,5])), [1, 3, 6, 10, 15])
1506
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001507 def test_accumulate_reducible(self):
1508 # check copy, deepcopy, pickle
1509 data = [1, 2, 3, 4, 5]
1510 accumulated = [1, 3, 6, 10, 15]
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001511
Serhiy Storchakabad12572014-12-15 14:03:42 +02001512 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1513 it = accumulate(data)
1514 self.assertEqual(list(pickle.loads(pickle.dumps(it, proto))), accumulated[:])
1515 self.assertEqual(next(it), 1)
1516 self.assertEqual(list(pickle.loads(pickle.dumps(it, proto))), accumulated[1:])
1517 it = accumulate(data)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001518 self.assertEqual(next(it), 1)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001519 self.assertEqual(list(copy.deepcopy(it)), accumulated[1:])
1520 self.assertEqual(list(copy.copy(it)), accumulated[1:])
1521
Serhiy Storchakad5516252016-03-06 14:00:45 +02001522 def test_accumulate_reducible_none(self):
1523 # Issue #25718: total is None
1524 it = accumulate([None, None, None], operator.is_)
1525 self.assertEqual(next(it), None)
1526 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1527 it_copy = pickle.loads(pickle.dumps(it, proto))
1528 self.assertEqual(list(it_copy), [True, False])
1529 self.assertEqual(list(copy.deepcopy(it)), [True, False])
1530 self.assertEqual(list(copy.copy(it)), [True, False])
1531
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001532 def test_chain(self):
1533 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
1534
1535 def test_chain_from_iterable(self):
1536 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
1537
1538 def test_combinations(self):
1539 self.assertEqual(list(combinations('ABCD', 2)),
1540 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
1541 self.assertEqual(list(combinations(range(4), 3)),
1542 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
1543
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001544 def test_combinations_with_replacement(self):
1545 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
1546 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
1547
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001548 def test_compress(self):
1549 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
1550
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001551 def test_count(self):
1552 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
1553
1554 def test_cycle(self):
1555 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
1556
1557 def test_dropwhile(self):
1558 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
1559
1560 def test_groupby(self):
1561 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
1562 list('ABCDAB'))
1563 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
1564 [list('AAAA'), list('BBB'), list('CC'), list('D')])
1565
1566 def test_filter(self):
1567 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
1568
1569 def test_filterfalse(self):
1570 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
1571
1572 def test_map(self):
1573 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
1574
1575 def test_islice(self):
1576 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
1577 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
1578 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
1579 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1580
1581 def test_zip(self):
1582 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
1583
1584 def test_zip_longest(self):
1585 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
1586 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
1587
1588 def test_permutations(self):
1589 self.assertEqual(list(permutations('ABCD', 2)),
1590 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
1591 self.assertEqual(list(permutations(range(3))),
1592 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
1593
1594 def test_product(self):
1595 self.assertEqual(list(product('ABCD', 'xy')),
1596 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
1597 self.assertEqual(list(product(range(2), repeat=3)),
1598 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
1599 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
1600
1601 def test_repeat(self):
1602 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
1603
1604 def test_stapmap(self):
1605 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
1606 [32, 9, 1000])
1607
1608 def test_takewhile(self):
1609 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
1610
1611
Cheryl Sabellada1734c2018-03-26 21:29:33 -04001612class TestPurePythonRoughEquivalents(unittest.TestCase):
1613
1614 @staticmethod
1615 def islice(iterable, *args):
1616 s = slice(*args)
1617 start, stop, step = s.start or 0, s.stop or sys.maxsize, s.step or 1
1618 it = iter(range(start, stop, step))
1619 try:
1620 nexti = next(it)
1621 except StopIteration:
1622 # Consume *iterable* up to the *start* position.
1623 for i, element in zip(range(start), iterable):
1624 pass
1625 return
1626 try:
1627 for i, element in enumerate(iterable):
1628 if i == nexti:
1629 yield element
1630 nexti = next(it)
1631 except StopIteration:
1632 # Consume to *stop*.
1633 for i, element in zip(range(i + 1, stop), iterable):
1634 pass
1635
1636 def test_islice_recipe(self):
1637 self.assertEqual(list(self.islice('ABCDEFG', 2)), list('AB'))
1638 self.assertEqual(list(self.islice('ABCDEFG', 2, 4)), list('CD'))
1639 self.assertEqual(list(self.islice('ABCDEFG', 2, None)), list('CDEFG'))
1640 self.assertEqual(list(self.islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1641 # Test items consumed.
1642 it = iter(range(10))
1643 self.assertEqual(list(self.islice(it, 3)), list(range(3)))
1644 self.assertEqual(list(it), list(range(3, 10)))
1645 it = iter(range(10))
1646 self.assertEqual(list(self.islice(it, 3, 3)), [])
1647 self.assertEqual(list(it), list(range(3, 10)))
1648 # Test that slice finishes in predictable state.
1649 c = count()
1650 self.assertEqual(list(self.islice(c, 1, 3, 50)), [1])
1651 self.assertEqual(next(c), 3)
1652
1653
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001654class TestGC(unittest.TestCase):
1655
1656 def makecycle(self, iterator, container):
1657 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +00001658 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001659 del container, iterator
1660
Raymond Hettinger482ba772010-12-01 22:48:00 +00001661 def test_accumulate(self):
1662 a = []
1663 self.makecycle(accumulate([1,2,a,3]), a)
1664
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001665 def test_chain(self):
1666 a = []
1667 self.makecycle(chain(a), a)
1668
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001669 def test_chain_from_iterable(self):
1670 a = []
1671 self.makecycle(chain.from_iterable([a]), a)
1672
1673 def test_combinations(self):
1674 a = []
1675 self.makecycle(combinations([1,2,a,3], 3), a)
1676
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001677 def test_combinations_with_replacement(self):
1678 a = []
1679 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1680
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001681 def test_compress(self):
1682 a = []
1683 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1684
Raymond Hettingerd6280f42009-02-16 20:50:56 +00001685 def test_count(self):
1686 a = []
1687 Int = type('Int', (int,), dict(x=a))
1688 self.makecycle(count(Int(0), Int(1)), a)
1689
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001690 def test_cycle(self):
1691 a = []
1692 self.makecycle(cycle([a]*2), a)
1693
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001694 def test_dropwhile(self):
1695 a = []
1696 self.makecycle(dropwhile(bool, [0, a, a]), a)
1697
1698 def test_groupby(self):
1699 a = []
1700 self.makecycle(groupby([a]*2, lambda x:x), a)
1701
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001702 def test_issue2246(self):
1703 # Issue 2246 -- the _grouper iterator was not included in GC
1704 n = 10
1705 keyfunc = lambda x: x
1706 for i, j in groupby(range(n), key=keyfunc):
1707 keyfunc.__dict__.setdefault('x',[]).append(j)
1708
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001709 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001710 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001711 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001712
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001713 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001714 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001715 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001716
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001717 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001718 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001719 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001720
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001721 def test_zip_longest(self):
1722 a = []
1723 self.makecycle(zip_longest([a]*2, [a]*3), a)
1724 b = [a, None]
1725 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1726
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001727 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001728 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001729 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001730
1731 def test_islice(self):
1732 a = []
1733 self.makecycle(islice([a]*2, None), a)
1734
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001735 def test_permutations(self):
1736 a = []
1737 self.makecycle(permutations([1,2,a,3], 3), a)
1738
1739 def test_product(self):
1740 a = []
1741 self.makecycle(product([1,2,a,3], repeat=3), a)
1742
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001743 def test_repeat(self):
1744 a = []
1745 self.makecycle(repeat(a), a)
1746
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001747 def test_starmap(self):
1748 a = []
1749 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1750
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001751 def test_takewhile(self):
1752 a = []
1753 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1754
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001755def R(seqn):
1756 'Regular generator'
1757 for i in seqn:
1758 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001759
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001760class G:
1761 'Sequence using __getitem__'
1762 def __init__(self, seqn):
1763 self.seqn = seqn
1764 def __getitem__(self, i):
1765 return self.seqn[i]
1766
1767class I:
1768 'Sequence using iterator protocol'
1769 def __init__(self, seqn):
1770 self.seqn = seqn
1771 self.i = 0
1772 def __iter__(self):
1773 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001774 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001775 if self.i >= len(self.seqn): raise StopIteration
1776 v = self.seqn[self.i]
1777 self.i += 1
1778 return v
1779
1780class Ig:
1781 'Sequence using iterator protocol defined with a generator'
1782 def __init__(self, seqn):
1783 self.seqn = seqn
1784 self.i = 0
1785 def __iter__(self):
1786 for val in self.seqn:
1787 yield val
1788
1789class X:
1790 'Missing __getitem__ and __iter__'
1791 def __init__(self, seqn):
1792 self.seqn = seqn
1793 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001794 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001795 if self.i >= len(self.seqn): raise StopIteration
1796 v = self.seqn[self.i]
1797 self.i += 1
1798 return v
1799
1800class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001801 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001802 def __init__(self, seqn):
1803 self.seqn = seqn
1804 self.i = 0
1805 def __iter__(self):
1806 return self
1807
1808class E:
1809 'Test propagation of exceptions'
1810 def __init__(self, seqn):
1811 self.seqn = seqn
1812 self.i = 0
1813 def __iter__(self):
1814 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001815 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001816 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001817
1818class S:
1819 'Test immediate stop'
1820 def __init__(self, seqn):
1821 pass
1822 def __iter__(self):
1823 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001824 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001825 raise StopIteration
1826
1827def L(seqn):
1828 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001829 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001830
1831
1832class TestVariousIteratorArgs(unittest.TestCase):
1833
Raymond Hettinger482ba772010-12-01 22:48:00 +00001834 def test_accumulate(self):
1835 s = [1,2,3,4,5]
1836 r = [1,3,6,10,15]
1837 n = len(s)
1838 for g in (G, I, Ig, L, R):
1839 self.assertEqual(list(accumulate(g(s))), r)
1840 self.assertEqual(list(accumulate(S(s))), [])
1841 self.assertRaises(TypeError, accumulate, X(s))
1842 self.assertRaises(TypeError, accumulate, N(s))
1843 self.assertRaises(ZeroDivisionError, list, accumulate(E(s)))
1844
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001845 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001846 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001847 for g in (G, I, Ig, S, L, R):
1848 self.assertEqual(list(chain(g(s))), list(g(s)))
1849 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001850 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001851 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001852 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1853
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001854 def test_compress(self):
1855 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1856 n = len(s)
1857 for g in (G, I, Ig, S, L, R):
1858 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1859 self.assertRaises(TypeError, compress, X(s), repeat(1))
1860 self.assertRaises(TypeError, compress, N(s), repeat(1))
1861 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1862
Christian Heimesc3f30c42008-02-22 16:37:40 +00001863 def test_product(self):
1864 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1865 self.assertRaises(TypeError, product, X(s))
1866 self.assertRaises(TypeError, product, N(s))
1867 self.assertRaises(ZeroDivisionError, product, E(s))
1868
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001869 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001870 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001871 for g in (G, I, Ig, S, L, R):
1872 tgtlen = len(s) * 3
1873 expected = list(g(s))*3
1874 actual = list(islice(cycle(g(s)), tgtlen))
1875 self.assertEqual(actual, expected)
1876 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001877 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001878 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1879
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001880 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001881 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001882 for g in (G, I, Ig, S, L, R):
1883 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1884 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001885 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001886 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1887
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001888 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001889 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001890 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001891 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001892 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001893 self.assertRaises(TypeError, filter, isEven, X(s))
1894 self.assertRaises(TypeError, filter, isEven, N(s))
1895 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001896
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001897 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001898 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001899 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001900 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001901 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001902 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1903 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1904 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001905
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001906 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001907 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001908 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001909 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1910 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1911 self.assertRaises(TypeError, zip, X(s))
1912 self.assertRaises(TypeError, zip, N(s))
1913 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001914
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001915 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001916 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001917 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001918 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1919 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1920 self.assertRaises(TypeError, zip_longest, X(s))
1921 self.assertRaises(TypeError, zip_longest, N(s))
1922 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001923
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001924 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001925 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001926 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001927 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001928 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001929 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001930 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001931 self.assertRaises(TypeError, map, onearg, X(s))
1932 self.assertRaises(TypeError, map, onearg, N(s))
1933 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001934
1935 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001936 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001937 for g in (G, I, Ig, S, L, R):
1938 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1939 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001940 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001941 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1942
1943 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001944 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001945 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001946 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001947 self.assertEqual(list(starmap(operator.pow, g(ss))),
1948 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001949 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001950 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001951 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1952
1953 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001954 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001955 for g in (G, I, Ig, S, L, R):
1956 tgt = []
1957 for elem in g(s):
1958 if not isEven(elem): break
1959 tgt.append(elem)
1960 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1961 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001962 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001963 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1964
1965 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001966 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001967 for g in (G, I, Ig, S, L, R):
1968 tgt = []
1969 for elem in g(s):
1970 if not tgt and isOdd(elem): continue
1971 tgt.append(elem)
1972 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1973 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001974 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001975 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1976
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001977 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001978 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001979 for g in (G, I, Ig, S, L, R):
1980 it1, it2 = tee(g(s))
1981 self.assertEqual(list(it1), list(g(s)))
1982 self.assertEqual(list(it2), list(g(s)))
1983 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001984 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001985 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1986
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001987class LengthTransparency(unittest.TestCase):
1988
1989 def test_repeat(self):
Armin Ronacheraa9a79d2012-10-06 14:03:24 +02001990 self.assertEqual(operator.length_hint(repeat(None, 50)), 50)
Raymond Hettinger97d35552014-06-24 21:36:58 -07001991 self.assertEqual(operator.length_hint(repeat(None, 0)), 0)
Armin Ronacheraa9a79d2012-10-06 14:03:24 +02001992 self.assertEqual(operator.length_hint(repeat(None), 12), 12)
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001993
Raymond Hettinger97d35552014-06-24 21:36:58 -07001994 def test_repeat_with_negative_times(self):
1995 self.assertEqual(operator.length_hint(repeat(None, -1)), 0)
1996 self.assertEqual(operator.length_hint(repeat(None, -2)), 0)
1997 self.assertEqual(operator.length_hint(repeat(None, times=-1)), 0)
1998 self.assertEqual(operator.length_hint(repeat(None, times=-2)), 0)
1999
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002000class RegressionTests(unittest.TestCase):
2001
2002 def test_sf_793826(self):
2003 # Fix Armin Rigo's successful efforts to wreak havoc
2004
2005 def mutatingtuple(tuple1, f, tuple2):
2006 # this builds a tuple t which is a copy of tuple1,
2007 # then calls f(t), then mutates t to be equal to tuple2
2008 # (needs len(tuple1) == len(tuple2)).
2009 def g(value, first=[1]):
2010 if first:
2011 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00002012 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002013 return value
2014 items = list(tuple2)
2015 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002016 gen = map(g, items)
2017 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00002018 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002019
2020 def f(t):
2021 global T
2022 T = t
2023 first[:] = list(T)
2024
2025 first = []
2026 mutatingtuple((1,2,3), f, (4,5,6))
2027 second = list(T)
2028 self.assertEqual(first, second)
2029
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002030
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00002031 def test_sf_950057(self):
2032 # Make sure that chain() and cycle() catch exceptions immediately
2033 # rather than when shifting between input sources
2034
2035 def gen1():
2036 hist.append(0)
2037 yield 1
2038 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00002039 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00002040 hist.append(2)
2041
2042 def gen2(x):
2043 hist.append(3)
2044 yield 2
2045 hist.append(4)
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00002046
2047 hist = []
2048 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
2049 self.assertEqual(hist, [0,1])
2050
2051 hist = []
2052 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
2053 self.assertEqual(hist, [0,1])
2054
2055 hist = []
2056 self.assertRaises(AssertionError, list, cycle(gen1()))
2057 self.assertEqual(hist, [0,1])
2058
T. Wouters5466d4a2017-03-30 09:58:35 -07002059 def test_long_chain_of_empty_iterables(self):
2060 # Make sure itertools.chain doesn't run into recursion limits when
2061 # dealing with long chains of empty iterables. Even with a high
2062 # number this would probably only fail in Py_DEBUG mode.
2063 it = chain.from_iterable(() for unused in range(10000000))
2064 with self.assertRaises(StopIteration):
2065 next(it)
2066
Serhiy Storchakac740e4f2017-09-26 21:47:56 +03002067 def test_issue30347_1(self):
2068 def f(n):
2069 if n == 5:
2070 list(b)
2071 return n != 6
2072 for (k, b) in groupby(range(10), f):
2073 list(b) # shouldn't crash
2074
2075 def test_issue30347_2(self):
2076 class K:
2077 def __init__(self, v):
2078 pass
2079 def __eq__(self, other):
2080 nonlocal i
2081 i += 1
2082 if i == 1:
2083 next(g, None)
2084 return True
2085 i = 0
2086 g = next(groupby(range(10), K))[1]
2087 for j in range(2):
2088 next(g, None) # shouldn't crash
2089
2090
Thomas Woutersb2137042007-02-01 18:02:27 +00002091class SubclassWithKwargsTest(unittest.TestCase):
2092 def test_keywords_in_subclass(self):
2093 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002094 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002095 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00002096 class Subclass(cls):
2097 def __init__(self, newarg=None, *args):
2098 cls.__init__(self, *args)
2099 try:
2100 Subclass(newarg=1)
2101 except TypeError as err:
2102 # we expect type errors because of wrong argument count
Serhiy Storchaka5eb788b2017-06-06 18:45:22 +03002103 self.assertNotIn("keyword arguments", err.args[0])
Thomas Woutersb2137042007-02-01 18:02:27 +00002104
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002105@support.cpython_only
2106class SizeofTest(unittest.TestCase):
2107 def setUp(self):
2108 self.ssize_t = struct.calcsize('n')
2109
2110 check_sizeof = support.check_sizeof
2111
2112 def test_product_sizeof(self):
2113 basesize = support.calcobjsize('3Pi')
2114 check = self.check_sizeof
2115 check(product('ab', '12'), basesize + 2 * self.ssize_t)
2116 check(product(*(('abc',) * 10)), basesize + 10 * self.ssize_t)
2117
2118 def test_combinations_sizeof(self):
2119 basesize = support.calcobjsize('3Pni')
2120 check = self.check_sizeof
2121 check(combinations('abcd', 3), basesize + 3 * self.ssize_t)
2122 check(combinations(range(10), 4), basesize + 4 * self.ssize_t)
2123
2124 def test_combinations_with_replacement_sizeof(self):
2125 cwr = combinations_with_replacement
2126 basesize = support.calcobjsize('3Pni')
2127 check = self.check_sizeof
2128 check(cwr('abcd', 3), basesize + 3 * self.ssize_t)
2129 check(cwr(range(10), 4), basesize + 4 * self.ssize_t)
2130
2131 def test_permutations_sizeof(self):
2132 basesize = support.calcobjsize('4Pni')
2133 check = self.check_sizeof
2134 check(permutations('abcd'),
2135 basesize + 4 * self.ssize_t + 4 * self.ssize_t)
2136 check(permutations('abcd', 3),
2137 basesize + 4 * self.ssize_t + 3 * self.ssize_t)
2138 check(permutations('abcde', 3),
2139 basesize + 5 * self.ssize_t + 3 * self.ssize_t)
2140 check(permutations(range(10), 4),
2141 basesize + 10 * self.ssize_t + 4 * self.ssize_t)
2142
Thomas Woutersb2137042007-02-01 18:02:27 +00002143
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002144libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002145
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002146
2147>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002148>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00002149... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00002150...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002151Check 1200 is for $120.15
2152Check 1201 is for $764.05
2153Check 1202 is for $823.14
2154
2155>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002156>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00002157... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00002158...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000021591
21608
216127
2162
2163>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00002164>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00002165... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00002166...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002167Alex
2168Laura
2169Martin
2170Walter
2171Samuele
2172
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002173>>> from operator import itemgetter
2174>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00002175>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002176>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00002177... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00002178...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000021791 ['a', 'c', 'e']
21802 ['b', 'd', 'f']
21813 ['g']
2182
Raymond Hettinger734fb572004-01-20 20:04:40 +00002183# Find runs of consecutive numbers using groupby. The key to the solution
2184# is differencing with a range so that consecutive numbers all appear in
2185# same group.
2186>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00002187>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00002188... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00002189...
Raymond Hettinger734fb572004-01-20 20:04:40 +00002190[1]
2191[4, 5, 6]
2192[10]
2193[15, 16, 17, 18]
2194[22]
2195[25, 26, 27, 28]
2196
Georg Brandl3dbca812008-07-23 16:10:53 +00002197>>> def take(n, iterable):
2198... "Return first n items of the iterable as a list"
2199... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00002200
Raymond Hettinger9265dd72018-04-08 08:44:20 -07002201>>> def prepend(value, iterator):
2202... "Prepend a single value in front of an iterator"
2203... # prepend(1, [2, 3, 4]) -> 1 2 3 4
2204... return chain([value], iterator)
2205
Georg Brandl3dbca812008-07-23 16:10:53 +00002206>>> def enumerate(iterable, start=0):
2207... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002208
Georg Brandl3dbca812008-07-23 16:10:53 +00002209>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002210... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00002211... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002212
Cheryl Sabellada1734c2018-03-26 21:29:33 -04002213>>> import collections
2214>>> def consume(iterator, n=None):
2215... "Advance the iterator n-steps ahead. If n is None, consume entirely."
2216... # Use functions that consume iterators at C speed.
2217... if n is None:
2218... # feed the entire iterator into a zero-length deque
2219... collections.deque(iterator, maxlen=0)
2220... else:
2221... # advance to the empty slice starting at position n
2222... next(islice(iterator, n, n), None)
2223
Raymond Hettingercdf8ba32009-02-19 04:45:07 +00002224>>> def nth(iterable, n, default=None):
2225... "Returns the nth item or a default value"
2226... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002227
Raymond Hettingere525ee32016-03-06 18:11:38 -08002228>>> def all_equal(iterable):
2229... "Returns True if all the elements are equal to each other"
2230... g = groupby(iterable)
2231... return next(g, True) and not next(g, False)
2232
Georg Brandl3dbca812008-07-23 16:10:53 +00002233>>> def quantify(iterable, pred=bool):
2234... "Count how many times the predicate is true"
2235... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00002236
Georg Brandl3dbca812008-07-23 16:10:53 +00002237>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002238... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00002239... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002240
Georg Brandl3dbca812008-07-23 16:10:53 +00002241>>> def ncycles(iterable, n):
Ezio Melotti13925002011-03-16 11:05:33 +02002242... "Returns the sequence elements n times"
Georg Brandl3dbca812008-07-23 16:10:53 +00002243... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002244
2245>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002246... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002247
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002248>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002249... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002250
2251>>> def repeatfunc(func, times=None, *args):
2252... "Repeat calls to func with specified arguments."
2253... " Example: repeatfunc(random.random)"
2254... if times is None:
2255... return starmap(func, repeat(args))
2256... else:
2257... return starmap(func, repeat(args, times))
2258
Raymond Hettingerd591f662003-10-26 15:34:50 +00002259>>> def pairwise(iterable):
2260... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
2261... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00002262... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00002263... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00002264... except StopIteration:
2265... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002266... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002267
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002268>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00002269... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002270... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00002271... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002272
2273>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00002274... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002275... # Recipe credited to George Sakkis
2276... pending = len(iterables)
2277... nexts = cycle(iter(it).__next__ for it in iterables)
2278... while pending:
2279... try:
2280... for next in nexts:
2281... yield next()
2282... except StopIteration:
2283... pending -= 1
2284... nexts = cycle(islice(nexts, pending))
2285
2286>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00002287... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
2288... s = list(iterable)
2289... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002290
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00002291>>> def unique_everseen(iterable, key=None):
2292... "List unique elements, preserving order. Remember all elements ever seen."
2293... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
2294... # unique_everseen('ABBCcAD', str.lower) --> A B C D
2295... seen = set()
2296... seen_add = seen.add
2297... if key is None:
2298... for element in iterable:
2299... if element not in seen:
2300... seen_add(element)
2301... yield element
2302... else:
2303... for element in iterable:
2304... k = key(element)
2305... if k not in seen:
2306... seen_add(k)
2307... yield element
2308
2309>>> def unique_justseen(iterable, key=None):
2310... "List unique elements, preserving order. Remember only the element just seen."
2311... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
2312... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
2313... return map(next, map(itemgetter(1), groupby(iterable, key)))
2314
Raymond Hettinger31b26f62014-04-02 03:16:42 -07002315>>> def first_true(iterable, default=False, pred=None):
2316... '''Returns the first true value in the iterable.
2317...
2318... If no true value is found, returns *default*
2319...
2320... If *pred* is not None, returns the first item
2321... for which pred(item) is true.
2322...
2323... '''
2324... # first_true([a,b,c], x) --> a or b or c or x
2325... # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x
2326... return next(filter(pred, iterable), default)
2327
Raymond Hettingerd37258d2018-01-13 10:35:40 -08002328>>> def nth_combination(iterable, r, index):
2329... 'Equivalent to list(combinations(iterable, r))[index]'
2330... pool = tuple(iterable)
2331... n = len(pool)
2332... if r < 0 or r > n:
2333... raise ValueError
2334... c = 1
2335... k = min(r, n-r)
2336... for i in range(1, k+1):
2337... c = c * (n - k + i) // i
2338... if index < 0:
2339... index += c
2340... if index < 0 or index >= c:
2341... raise IndexError
2342... result = []
2343... while r:
2344... c, n, r = c*r//n, n-1, r-1
2345... while index >= c:
2346... index -= c
2347... c, n = c*(n-r)//n, n-1
2348... result.append(pool[-1-n])
2349... return tuple(result)
2350
2351
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002352This is not part of the examples but it tests to make sure the definitions
2353perform as purported.
2354
Raymond Hettingera098b332003-09-08 23:58:40 +00002355>>> take(10, count())
2356[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
2357
Raymond Hettinger9265dd72018-04-08 08:44:20 -07002358>>> list(prepend(1, [2, 3, 4]))
2359[1, 2, 3, 4]
2360
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002361>>> list(enumerate('abc'))
2362[(0, 'a'), (1, 'b'), (2, 'c')]
2363
2364>>> list(islice(tabulate(lambda x: 2*x), 4))
2365[0, 2, 4, 6]
2366
Cheryl Sabellada1734c2018-03-26 21:29:33 -04002367>>> it = iter(range(10))
2368>>> consume(it, 3)
2369>>> next(it)
23703
2371>>> consume(it)
2372>>> next(it, 'Done')
2373'Done'
2374
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002375>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00002376'd'
2377
2378>>> nth('abcde', 9) is None
2379True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002380
Raymond Hettingere525ee32016-03-06 18:11:38 -08002381>>> [all_equal(s) for s in ('', 'A', 'AAAA', 'AAAB', 'AAABA')]
2382[True, True, True, False, False]
2383
Guido van Rossum805365e2007-05-07 22:24:25 +00002384>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000238550
2386
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002387>>> a = [[1, 2, 3], [4, 5, 6]]
2388>>> flatten(a)
2389[1, 2, 3, 4, 5, 6]
2390
2391>>> list(repeatfunc(pow, 5, 2, 3))
2392[8, 8, 8, 8, 8]
2393
2394>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002395>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002396[0, 0, 0, 0, 0]
2397
Raymond Hettingerd591f662003-10-26 15:34:50 +00002398>>> list(pairwise('abcd'))
2399[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002400
Raymond Hettingerd591f662003-10-26 15:34:50 +00002401>>> list(pairwise([]))
2402[]
2403
2404>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00002405[]
2406
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002407>>> list(islice(padnone('abc'), 0, 6))
2408['a', 'b', 'c', None, None, None]
2409
2410>>> list(ncycles('abc', 3))
2411['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
2412
2413>>> dotproduct([1,2,3], [4,5,6])
241432
2415
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002416>>> list(grouper(3, 'abcdefg', 'x'))
2417[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
2418
2419>>> list(roundrobin('abc', 'd', 'ef'))
2420['a', 'd', 'e', 'b', 'f', 'c']
2421
Raymond Hettingerace67332009-01-26 02:23:50 +00002422>>> list(powerset([1,2,3]))
2423[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002424
Raymond Hettinger191e8502009-01-27 13:29:43 +00002425>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
2426True
2427
2428>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
2429True
2430
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00002431>>> list(unique_everseen('AAAABBBCCDAABBB'))
2432['A', 'B', 'C', 'D']
2433
2434>>> list(unique_everseen('ABBCcAD', str.lower))
2435['A', 'B', 'C', 'D']
2436
2437>>> list(unique_justseen('AAAABBBCCDAABBB'))
2438['A', 'B', 'C', 'D', 'A', 'B']
2439
2440>>> list(unique_justseen('ABBCcAD', str.lower))
2441['A', 'B', 'C', 'A', 'D']
2442
Raymond Hettinger31b26f62014-04-02 03:16:42 -07002443>>> first_true('ABC0DEF1', '9', str.isdigit)
2444'0'
2445
Raymond Hettingerd37258d2018-01-13 10:35:40 -08002446>>> population = 'ABCDEFGH'
2447>>> for r in range(len(population) + 1):
2448... seq = list(combinations(population, r))
2449... for i in range(len(seq)):
2450... assert nth_combination(population, r, i) == seq[i]
2451... for i in range(-len(seq), 0):
2452... assert nth_combination(population, r, i) == seq[i]
2453
2454
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002455"""
2456
2457__test__ = {'libreftest' : libreftest}
2458
2459def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002460 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00002461 RegressionTests, LengthTransparency,
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002462 SubclassWithKwargsTest, TestExamples,
Cheryl Sabellada1734c2018-03-26 21:29:33 -04002463 TestPurePythonRoughEquivalents,
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002464 SizeofTest)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002465 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002466
2467 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002468 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00002469 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00002470 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00002471 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002472 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00002473 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00002474 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00002475 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002476
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002477 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002478 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002479
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002480if __name__ == "__main__":
2481 test_main(verbose=True)