blob: a359905e2e48cbd8b5fd52d33e2e52f213b9dc95 [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),
1195 (10, 3),
1196 (20,)
1197 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +00001198 self.assertEqual(list(islice(range(100), *args)),
1199 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001200
1201 for args, tgtargs in [ # Stop when seqn is exhausted
1202 ((10, 110, 3), ((10, 100, 3))),
1203 ((10, 110), ((10, 100))),
1204 ((110,), (100,))
1205 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +00001206 self.assertEqual(list(islice(range(100), *args)),
1207 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001208
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001209 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +00001210 self.assertEqual(list(islice(range(10), None)), list(range(10)))
1211 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
1212 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
1213 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
1214 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001215
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +00001216 # Test number of items consumed SF #1171417
1217 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +00001218 self.assertEqual(list(islice(it, 3)), list(range(3)))
1219 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +00001220
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001221 # Test invalid arguments
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001222 ra = range(10)
1223 self.assertRaises(TypeError, islice, ra)
1224 self.assertRaises(TypeError, islice, ra, 1, 2, 3, 4)
1225 self.assertRaises(ValueError, islice, ra, -5, 10, 1)
1226 self.assertRaises(ValueError, islice, ra, 1, -5, -1)
1227 self.assertRaises(ValueError, islice, ra, 1, 10, -1)
1228 self.assertRaises(ValueError, islice, ra, 1, 10, 0)
1229 self.assertRaises(ValueError, islice, ra, 'a')
1230 self.assertRaises(ValueError, islice, ra, 'a', 1)
1231 self.assertRaises(ValueError, islice, ra, 1, 'a')
1232 self.assertRaises(ValueError, islice, ra, 'a', 1, 1)
1233 self.assertRaises(ValueError, islice, ra, 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +00001234 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001235
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001236 # Issue #10323: Less islice in a predictable state
1237 c = count()
1238 self.assertEqual(list(islice(c, 1, 3, 50)), [1])
1239 self.assertEqual(next(c), 3)
1240
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001241 # check copy, deepcopy, pickle
1242 for args in [ # islice(args) should agree with range(args)
1243 (10, 20, 3),
1244 (10, 3, 20),
1245 (10, 20),
1246 (10, 3),
1247 (20,)
1248 ]:
1249 self.assertEqual(list(copy.copy(islice(range(100), *args))),
1250 list(range(*args)))
1251 self.assertEqual(list(copy.deepcopy(islice(range(100), *args))),
1252 list(range(*args)))
Serhiy Storchakabad12572014-12-15 14:03:42 +02001253 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1254 self.pickletest(proto, islice(range(100), *args))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001255
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001256 # Issue #21321: check source iterator is not referenced
1257 # from islice() after the latter has been exhausted
1258 it = (x for x in (1, 2))
1259 wr = weakref.ref(it)
1260 it = islice(it, 1)
1261 self.assertIsNotNone(wr())
1262 list(it) # exhaust the iterator
Benjamin Peterson8e163512014-08-24 18:07:28 -05001263 support.gc_collect()
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001264 self.assertIsNone(wr())
1265
Will Roberts0ecdc522017-06-08 08:03:04 +02001266 # Issue #30537: islice can accept integer-like objects as
1267 # arguments
1268 class IntLike(object):
1269 def __init__(self, val):
1270 self.val = val
1271 def __index__(self):
1272 return self.val
1273 self.assertEqual(list(islice(range(100), IntLike(10))), list(range(10)))
1274 self.assertEqual(list(islice(range(100), IntLike(10), IntLike(50))),
1275 list(range(10, 50)))
1276 self.assertEqual(list(islice(range(100), IntLike(10), IntLike(50), IntLike(5))),
1277 list(range(10,50,5)))
1278
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001279 def test_takewhile(self):
1280 data = [1, 3, 5, 20, 2, 4, 6, 8]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001281 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001282 self.assertEqual(list(takewhile(underten, [])), [])
1283 self.assertRaises(TypeError, takewhile)
1284 self.assertRaises(TypeError, takewhile, operator.pow)
1285 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001286 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
1287 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001288 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
1289 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +00001290 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001291
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001292 # check copy, deepcopy, pickle
1293 self.assertEqual(list(copy.copy(takewhile(underten, data))), [1, 3, 5])
1294 self.assertEqual(list(copy.deepcopy(takewhile(underten, data))),
1295 [1, 3, 5])
Serhiy Storchakabad12572014-12-15 14:03:42 +02001296 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1297 self.pickletest(proto, takewhile(underten, data))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001298
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001299 def test_dropwhile(self):
1300 data = [1, 3, 5, 20, 2, 4, 6, 8]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001301 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001302 self.assertEqual(list(dropwhile(underten, [])), [])
1303 self.assertRaises(TypeError, dropwhile)
1304 self.assertRaises(TypeError, dropwhile, operator.pow)
1305 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001306 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
1307 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001308
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001309 # check copy, deepcopy, pickle
1310 self.assertEqual(list(copy.copy(dropwhile(underten, data))), [20, 2, 4, 6, 8])
1311 self.assertEqual(list(copy.deepcopy(dropwhile(underten, data))),
1312 [20, 2, 4, 6, 8])
Serhiy Storchakabad12572014-12-15 14:03:42 +02001313 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1314 self.pickletest(proto, dropwhile(underten, data))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001315
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001316 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +00001317 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001318
1319 a, b = tee([]) # test empty iterator
1320 self.assertEqual(list(a), [])
1321 self.assertEqual(list(b), [])
1322
1323 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +00001324 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001325
1326 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +00001327 self.assertEqual(list(a), list(range(n)))
1328 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001329
1330 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +00001331 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001332 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001333 del a
Guido van Rossum805365e2007-05-07 22:24:25 +00001334 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001335
1336 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +00001337 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001338 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001339 del b
Guido van Rossum805365e2007-05-07 22:24:25 +00001340 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001341
Guido van Rossum805365e2007-05-07 22:24:25 +00001342 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001343 order = [0]*n + [1]*n
1344 random.shuffle(order)
1345 lists = ([], [])
1346 its = tee(irange(n))
1347 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +00001348 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001349 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +00001350 self.assertEqual(lists[0], list(range(n)))
1351 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001352
Raymond Hettingerad983e72003-11-12 14:32:26 +00001353 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001354 self.assertRaises(TypeError, tee)
1355 self.assertRaises(TypeError, tee, 3)
1356 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +00001357 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001358
Raymond Hettingerad983e72003-11-12 14:32:26 +00001359 # tee object should be instantiable
1360 a, b = tee('abc')
1361 c = type(a)('def')
1362 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +00001363
Raymond Hettingerad983e72003-11-12 14:32:26 +00001364 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +00001365 a, b, c = tee(range(2000), 3)
1366 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001367 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +00001368 self.assertEqual(list(b), list(range(2000)))
1369 self.assertEqual([next(c), next(c)], list(range(2)))
1370 self.assertEqual(list(a), list(range(100,2000)))
1371 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +00001372
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001373 # test values of n
1374 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +00001375 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +00001376 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001377 result = tee('abc', n)
1378 self.assertEqual(type(result), tuple)
1379 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001380 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001381
Raymond Hettingerad983e72003-11-12 14:32:26 +00001382 # tee pass-through to copyable iterator
1383 a, b = tee('abc')
1384 c, d = tee(a)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +00001385 self.assertTrue(a is c)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001386
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001387 # test tee_new
1388 t1, t2 = tee('abc')
1389 tnew = type(t1)
1390 self.assertRaises(TypeError, tnew)
1391 self.assertRaises(TypeError, tnew, 10)
1392 t3 = tnew(t1)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +00001393 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +00001394
Raymond Hettingera9f60922004-10-17 16:40:14 +00001395 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +00001396 a, b = tee(range(10))
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001397 p = weakref.proxy(a)
Raymond Hettingera9f60922004-10-17 16:40:14 +00001398 self.assertEqual(getattr(p, '__class__'), type(b))
1399 del a
1400 self.assertRaises(ReferenceError, getattr, p, '__class__')
1401
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001402 ans = list('abc')
1403 long_ans = list(range(10000))
1404
1405 # check copy
1406 a, b = tee('abc')
1407 self.assertEqual(list(copy.copy(a)), ans)
1408 self.assertEqual(list(copy.copy(b)), ans)
1409 a, b = tee(list(range(10000)))
1410 self.assertEqual(list(copy.copy(a)), long_ans)
1411 self.assertEqual(list(copy.copy(b)), long_ans)
1412
1413 # check partially consumed copy
1414 a, b = tee('abc')
1415 take(2, a)
1416 take(1, b)
1417 self.assertEqual(list(copy.copy(a)), ans[2:])
1418 self.assertEqual(list(copy.copy(b)), ans[1:])
1419 self.assertEqual(list(a), ans[2:])
1420 self.assertEqual(list(b), ans[1:])
1421 a, b = tee(range(10000))
1422 take(100, a)
1423 take(60, b)
1424 self.assertEqual(list(copy.copy(a)), long_ans[100:])
1425 self.assertEqual(list(copy.copy(b)), long_ans[60:])
1426 self.assertEqual(list(a), long_ans[100:])
1427 self.assertEqual(list(b), long_ans[60:])
1428
1429 # check deepcopy
1430 a, b = tee('abc')
1431 self.assertEqual(list(copy.deepcopy(a)), ans)
1432 self.assertEqual(list(copy.deepcopy(b)), ans)
1433 self.assertEqual(list(a), ans)
1434 self.assertEqual(list(b), ans)
1435 a, b = tee(range(10000))
1436 self.assertEqual(list(copy.deepcopy(a)), long_ans)
1437 self.assertEqual(list(copy.deepcopy(b)), long_ans)
1438 self.assertEqual(list(a), long_ans)
1439 self.assertEqual(list(b), long_ans)
1440
1441 # check partially consumed deepcopy
1442 a, b = tee('abc')
1443 take(2, a)
1444 take(1, b)
1445 self.assertEqual(list(copy.deepcopy(a)), ans[2:])
1446 self.assertEqual(list(copy.deepcopy(b)), ans[1:])
1447 self.assertEqual(list(a), ans[2:])
1448 self.assertEqual(list(b), ans[1:])
1449 a, b = tee(range(10000))
1450 take(100, a)
1451 take(60, b)
1452 self.assertEqual(list(copy.deepcopy(a)), long_ans[100:])
1453 self.assertEqual(list(copy.deepcopy(b)), long_ans[60:])
1454 self.assertEqual(list(a), long_ans[100:])
1455 self.assertEqual(list(b), long_ans[60:])
1456
1457 # check pickle
Serhiy Storchakabad12572014-12-15 14:03:42 +02001458 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1459 self.pickletest(proto, iter(tee('abc')))
1460 a, b = tee('abc')
1461 self.pickletest(proto, a, compare=ans)
1462 self.pickletest(proto, b, compare=ans)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001463
Serhiy Storchakaa3e91282013-01-25 13:19:31 +02001464 # Issue 13454: Crash when deleting backward iterator from tee()
1465 def test_tee_del_backward(self):
Serhiy Storchaka5bb893c2013-01-26 11:52:06 +02001466 forward, backward = tee(repeat(None, 20000000))
Serhiy Storchaka9db55002015-03-28 20:38:37 +02001467 try:
1468 any(forward) # exhaust the iterator
1469 del backward
1470 except:
1471 del forward, backward
1472 raise
Serhiy Storchakaa3e91282013-01-25 13:19:31 +02001473
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001474 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001475 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +00001476
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001477 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +00001478 self.assertRaises(StopIteration, next, f([]))
1479 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001480
Georg Brandla18af4e2007-04-21 15:47:16 +00001481 self.assertRaises(StopIteration, next, islice([], None))
1482 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001483
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001484 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +00001485 self.assertRaises(StopIteration, next, p)
1486 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001487 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +00001488 self.assertRaises(StopIteration, next, p)
1489 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001490
Georg Brandla18af4e2007-04-21 15:47:16 +00001491 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001492
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001493 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +00001494 self.assertRaises(StopIteration, next, f(lambda x:x, []))
1495 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001496
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001497class TestExamples(unittest.TestCase):
1498
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001499 def test_accumulate(self):
Raymond Hettinger482ba772010-12-01 22:48:00 +00001500 self.assertEqual(list(accumulate([1,2,3,4,5])), [1, 3, 6, 10, 15])
1501
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001502 def test_accumulate_reducible(self):
1503 # check copy, deepcopy, pickle
1504 data = [1, 2, 3, 4, 5]
1505 accumulated = [1, 3, 6, 10, 15]
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001506
Serhiy Storchakabad12572014-12-15 14:03:42 +02001507 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1508 it = accumulate(data)
1509 self.assertEqual(list(pickle.loads(pickle.dumps(it, proto))), accumulated[:])
1510 self.assertEqual(next(it), 1)
1511 self.assertEqual(list(pickle.loads(pickle.dumps(it, proto))), accumulated[1:])
1512 it = accumulate(data)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001513 self.assertEqual(next(it), 1)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001514 self.assertEqual(list(copy.deepcopy(it)), accumulated[1:])
1515 self.assertEqual(list(copy.copy(it)), accumulated[1:])
1516
Serhiy Storchakad5516252016-03-06 14:00:45 +02001517 def test_accumulate_reducible_none(self):
1518 # Issue #25718: total is None
1519 it = accumulate([None, None, None], operator.is_)
1520 self.assertEqual(next(it), None)
1521 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1522 it_copy = pickle.loads(pickle.dumps(it, proto))
1523 self.assertEqual(list(it_copy), [True, False])
1524 self.assertEqual(list(copy.deepcopy(it)), [True, False])
1525 self.assertEqual(list(copy.copy(it)), [True, False])
1526
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001527 def test_chain(self):
1528 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
1529
1530 def test_chain_from_iterable(self):
1531 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
1532
1533 def test_combinations(self):
1534 self.assertEqual(list(combinations('ABCD', 2)),
1535 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
1536 self.assertEqual(list(combinations(range(4), 3)),
1537 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
1538
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001539 def test_combinations_with_replacement(self):
1540 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
1541 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
1542
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001543 def test_compress(self):
1544 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
1545
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001546 def test_count(self):
1547 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
1548
1549 def test_cycle(self):
1550 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
1551
1552 def test_dropwhile(self):
1553 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
1554
1555 def test_groupby(self):
1556 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
1557 list('ABCDAB'))
1558 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
1559 [list('AAAA'), list('BBB'), list('CC'), list('D')])
1560
1561 def test_filter(self):
1562 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
1563
1564 def test_filterfalse(self):
1565 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
1566
1567 def test_map(self):
1568 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
1569
1570 def test_islice(self):
1571 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
1572 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
1573 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
1574 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1575
1576 def test_zip(self):
1577 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
1578
1579 def test_zip_longest(self):
1580 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
1581 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
1582
1583 def test_permutations(self):
1584 self.assertEqual(list(permutations('ABCD', 2)),
1585 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
1586 self.assertEqual(list(permutations(range(3))),
1587 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
1588
1589 def test_product(self):
1590 self.assertEqual(list(product('ABCD', 'xy')),
1591 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
1592 self.assertEqual(list(product(range(2), repeat=3)),
1593 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
1594 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
1595
1596 def test_repeat(self):
1597 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
1598
1599 def test_stapmap(self):
1600 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
1601 [32, 9, 1000])
1602
1603 def test_takewhile(self):
1604 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
1605
1606
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001607class TestGC(unittest.TestCase):
1608
1609 def makecycle(self, iterator, container):
1610 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +00001611 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001612 del container, iterator
1613
Raymond Hettinger482ba772010-12-01 22:48:00 +00001614 def test_accumulate(self):
1615 a = []
1616 self.makecycle(accumulate([1,2,a,3]), a)
1617
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001618 def test_chain(self):
1619 a = []
1620 self.makecycle(chain(a), a)
1621
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001622 def test_chain_from_iterable(self):
1623 a = []
1624 self.makecycle(chain.from_iterable([a]), a)
1625
1626 def test_combinations(self):
1627 a = []
1628 self.makecycle(combinations([1,2,a,3], 3), a)
1629
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001630 def test_combinations_with_replacement(self):
1631 a = []
1632 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1633
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001634 def test_compress(self):
1635 a = []
1636 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1637
Raymond Hettingerd6280f42009-02-16 20:50:56 +00001638 def test_count(self):
1639 a = []
1640 Int = type('Int', (int,), dict(x=a))
1641 self.makecycle(count(Int(0), Int(1)), a)
1642
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001643 def test_cycle(self):
1644 a = []
1645 self.makecycle(cycle([a]*2), a)
1646
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001647 def test_dropwhile(self):
1648 a = []
1649 self.makecycle(dropwhile(bool, [0, a, a]), a)
1650
1651 def test_groupby(self):
1652 a = []
1653 self.makecycle(groupby([a]*2, lambda x:x), a)
1654
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001655 def test_issue2246(self):
1656 # Issue 2246 -- the _grouper iterator was not included in GC
1657 n = 10
1658 keyfunc = lambda x: x
1659 for i, j in groupby(range(n), key=keyfunc):
1660 keyfunc.__dict__.setdefault('x',[]).append(j)
1661
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001662 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001663 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001664 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001665
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001666 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001667 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001668 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001669
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001670 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001671 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001672 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001673
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001674 def test_zip_longest(self):
1675 a = []
1676 self.makecycle(zip_longest([a]*2, [a]*3), a)
1677 b = [a, None]
1678 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1679
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001680 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001681 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001682 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001683
1684 def test_islice(self):
1685 a = []
1686 self.makecycle(islice([a]*2, None), a)
1687
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001688 def test_permutations(self):
1689 a = []
1690 self.makecycle(permutations([1,2,a,3], 3), a)
1691
1692 def test_product(self):
1693 a = []
1694 self.makecycle(product([1,2,a,3], repeat=3), a)
1695
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001696 def test_repeat(self):
1697 a = []
1698 self.makecycle(repeat(a), a)
1699
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001700 def test_starmap(self):
1701 a = []
1702 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1703
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001704 def test_takewhile(self):
1705 a = []
1706 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1707
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001708def R(seqn):
1709 'Regular generator'
1710 for i in seqn:
1711 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001712
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001713class G:
1714 'Sequence using __getitem__'
1715 def __init__(self, seqn):
1716 self.seqn = seqn
1717 def __getitem__(self, i):
1718 return self.seqn[i]
1719
1720class I:
1721 'Sequence using iterator protocol'
1722 def __init__(self, seqn):
1723 self.seqn = seqn
1724 self.i = 0
1725 def __iter__(self):
1726 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001727 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001728 if self.i >= len(self.seqn): raise StopIteration
1729 v = self.seqn[self.i]
1730 self.i += 1
1731 return v
1732
1733class Ig:
1734 'Sequence using iterator protocol defined with a generator'
1735 def __init__(self, seqn):
1736 self.seqn = seqn
1737 self.i = 0
1738 def __iter__(self):
1739 for val in self.seqn:
1740 yield val
1741
1742class X:
1743 'Missing __getitem__ and __iter__'
1744 def __init__(self, seqn):
1745 self.seqn = seqn
1746 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001747 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001748 if self.i >= len(self.seqn): raise StopIteration
1749 v = self.seqn[self.i]
1750 self.i += 1
1751 return v
1752
1753class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001754 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001755 def __init__(self, seqn):
1756 self.seqn = seqn
1757 self.i = 0
1758 def __iter__(self):
1759 return self
1760
1761class E:
1762 'Test propagation of exceptions'
1763 def __init__(self, seqn):
1764 self.seqn = seqn
1765 self.i = 0
1766 def __iter__(self):
1767 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001768 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001769 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001770
1771class S:
1772 'Test immediate stop'
1773 def __init__(self, seqn):
1774 pass
1775 def __iter__(self):
1776 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001777 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001778 raise StopIteration
1779
1780def L(seqn):
1781 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001782 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001783
1784
1785class TestVariousIteratorArgs(unittest.TestCase):
1786
Raymond Hettinger482ba772010-12-01 22:48:00 +00001787 def test_accumulate(self):
1788 s = [1,2,3,4,5]
1789 r = [1,3,6,10,15]
1790 n = len(s)
1791 for g in (G, I, Ig, L, R):
1792 self.assertEqual(list(accumulate(g(s))), r)
1793 self.assertEqual(list(accumulate(S(s))), [])
1794 self.assertRaises(TypeError, accumulate, X(s))
1795 self.assertRaises(TypeError, accumulate, N(s))
1796 self.assertRaises(ZeroDivisionError, list, accumulate(E(s)))
1797
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001798 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001799 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001800 for g in (G, I, Ig, S, L, R):
1801 self.assertEqual(list(chain(g(s))), list(g(s)))
1802 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001803 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001804 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001805 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1806
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001807 def test_compress(self):
1808 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1809 n = len(s)
1810 for g in (G, I, Ig, S, L, R):
1811 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1812 self.assertRaises(TypeError, compress, X(s), repeat(1))
1813 self.assertRaises(TypeError, compress, N(s), repeat(1))
1814 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1815
Christian Heimesc3f30c42008-02-22 16:37:40 +00001816 def test_product(self):
1817 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1818 self.assertRaises(TypeError, product, X(s))
1819 self.assertRaises(TypeError, product, N(s))
1820 self.assertRaises(ZeroDivisionError, product, E(s))
1821
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001822 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001823 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001824 for g in (G, I, Ig, S, L, R):
1825 tgtlen = len(s) * 3
1826 expected = list(g(s))*3
1827 actual = list(islice(cycle(g(s)), tgtlen))
1828 self.assertEqual(actual, expected)
1829 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001830 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001831 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1832
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001833 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001834 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001835 for g in (G, I, Ig, S, L, R):
1836 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1837 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001838 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001839 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1840
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001841 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001842 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001843 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001844 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001845 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001846 self.assertRaises(TypeError, filter, isEven, X(s))
1847 self.assertRaises(TypeError, filter, isEven, N(s))
1848 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001849
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001850 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001851 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001852 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001853 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001854 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001855 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1856 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1857 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001858
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001859 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001860 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001861 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001862 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1863 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1864 self.assertRaises(TypeError, zip, X(s))
1865 self.assertRaises(TypeError, zip, N(s))
1866 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001867
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001868 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001869 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001870 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001871 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1872 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1873 self.assertRaises(TypeError, zip_longest, X(s))
1874 self.assertRaises(TypeError, zip_longest, N(s))
1875 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001876
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001877 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001878 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001879 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001880 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001881 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001882 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001883 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001884 self.assertRaises(TypeError, map, onearg, X(s))
1885 self.assertRaises(TypeError, map, onearg, N(s))
1886 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001887
1888 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001889 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001890 for g in (G, I, Ig, S, L, R):
1891 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1892 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001893 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001894 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1895
1896 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001897 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001898 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001899 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001900 self.assertEqual(list(starmap(operator.pow, g(ss))),
1901 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001902 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001903 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001904 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1905
1906 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001907 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001908 for g in (G, I, Ig, S, L, R):
1909 tgt = []
1910 for elem in g(s):
1911 if not isEven(elem): break
1912 tgt.append(elem)
1913 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1914 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001915 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001916 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1917
1918 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001919 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001920 for g in (G, I, Ig, S, L, R):
1921 tgt = []
1922 for elem in g(s):
1923 if not tgt and isOdd(elem): continue
1924 tgt.append(elem)
1925 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1926 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001927 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001928 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1929
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001930 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001931 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001932 for g in (G, I, Ig, S, L, R):
1933 it1, it2 = tee(g(s))
1934 self.assertEqual(list(it1), list(g(s)))
1935 self.assertEqual(list(it2), list(g(s)))
1936 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001937 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001938 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1939
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001940class LengthTransparency(unittest.TestCase):
1941
1942 def test_repeat(self):
Armin Ronacheraa9a79d2012-10-06 14:03:24 +02001943 self.assertEqual(operator.length_hint(repeat(None, 50)), 50)
Raymond Hettinger97d35552014-06-24 21:36:58 -07001944 self.assertEqual(operator.length_hint(repeat(None, 0)), 0)
Armin Ronacheraa9a79d2012-10-06 14:03:24 +02001945 self.assertEqual(operator.length_hint(repeat(None), 12), 12)
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001946
Raymond Hettinger97d35552014-06-24 21:36:58 -07001947 def test_repeat_with_negative_times(self):
1948 self.assertEqual(operator.length_hint(repeat(None, -1)), 0)
1949 self.assertEqual(operator.length_hint(repeat(None, -2)), 0)
1950 self.assertEqual(operator.length_hint(repeat(None, times=-1)), 0)
1951 self.assertEqual(operator.length_hint(repeat(None, times=-2)), 0)
1952
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001953class RegressionTests(unittest.TestCase):
1954
1955 def test_sf_793826(self):
1956 # Fix Armin Rigo's successful efforts to wreak havoc
1957
1958 def mutatingtuple(tuple1, f, tuple2):
1959 # this builds a tuple t which is a copy of tuple1,
1960 # then calls f(t), then mutates t to be equal to tuple2
1961 # (needs len(tuple1) == len(tuple2)).
1962 def g(value, first=[1]):
1963 if first:
1964 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00001965 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001966 return value
1967 items = list(tuple2)
1968 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001969 gen = map(g, items)
1970 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00001971 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001972
1973 def f(t):
1974 global T
1975 T = t
1976 first[:] = list(T)
1977
1978 first = []
1979 mutatingtuple((1,2,3), f, (4,5,6))
1980 second = list(T)
1981 self.assertEqual(first, second)
1982
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001983
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001984 def test_sf_950057(self):
1985 # Make sure that chain() and cycle() catch exceptions immediately
1986 # rather than when shifting between input sources
1987
1988 def gen1():
1989 hist.append(0)
1990 yield 1
1991 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001992 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001993 hist.append(2)
1994
1995 def gen2(x):
1996 hist.append(3)
1997 yield 2
1998 hist.append(4)
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001999
2000 hist = []
2001 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
2002 self.assertEqual(hist, [0,1])
2003
2004 hist = []
2005 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
2006 self.assertEqual(hist, [0,1])
2007
2008 hist = []
2009 self.assertRaises(AssertionError, list, cycle(gen1()))
2010 self.assertEqual(hist, [0,1])
2011
T. Wouters5466d4a2017-03-30 09:58:35 -07002012 def test_long_chain_of_empty_iterables(self):
2013 # Make sure itertools.chain doesn't run into recursion limits when
2014 # dealing with long chains of empty iterables. Even with a high
2015 # number this would probably only fail in Py_DEBUG mode.
2016 it = chain.from_iterable(() for unused in range(10000000))
2017 with self.assertRaises(StopIteration):
2018 next(it)
2019
Serhiy Storchakac740e4f2017-09-26 21:47:56 +03002020 def test_issue30347_1(self):
2021 def f(n):
2022 if n == 5:
2023 list(b)
2024 return n != 6
2025 for (k, b) in groupby(range(10), f):
2026 list(b) # shouldn't crash
2027
2028 def test_issue30347_2(self):
2029 class K:
2030 def __init__(self, v):
2031 pass
2032 def __eq__(self, other):
2033 nonlocal i
2034 i += 1
2035 if i == 1:
2036 next(g, None)
2037 return True
2038 i = 0
2039 g = next(groupby(range(10), K))[1]
2040 for j in range(2):
2041 next(g, None) # shouldn't crash
2042
2043
Thomas Woutersb2137042007-02-01 18:02:27 +00002044class SubclassWithKwargsTest(unittest.TestCase):
2045 def test_keywords_in_subclass(self):
2046 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002047 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002048 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00002049 class Subclass(cls):
2050 def __init__(self, newarg=None, *args):
2051 cls.__init__(self, *args)
2052 try:
2053 Subclass(newarg=1)
2054 except TypeError as err:
2055 # we expect type errors because of wrong argument count
Serhiy Storchaka5eb788b2017-06-06 18:45:22 +03002056 self.assertNotIn("keyword arguments", err.args[0])
Thomas Woutersb2137042007-02-01 18:02:27 +00002057
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002058@support.cpython_only
2059class SizeofTest(unittest.TestCase):
2060 def setUp(self):
2061 self.ssize_t = struct.calcsize('n')
2062
2063 check_sizeof = support.check_sizeof
2064
2065 def test_product_sizeof(self):
2066 basesize = support.calcobjsize('3Pi')
2067 check = self.check_sizeof
2068 check(product('ab', '12'), basesize + 2 * self.ssize_t)
2069 check(product(*(('abc',) * 10)), basesize + 10 * self.ssize_t)
2070
2071 def test_combinations_sizeof(self):
2072 basesize = support.calcobjsize('3Pni')
2073 check = self.check_sizeof
2074 check(combinations('abcd', 3), basesize + 3 * self.ssize_t)
2075 check(combinations(range(10), 4), basesize + 4 * self.ssize_t)
2076
2077 def test_combinations_with_replacement_sizeof(self):
2078 cwr = combinations_with_replacement
2079 basesize = support.calcobjsize('3Pni')
2080 check = self.check_sizeof
2081 check(cwr('abcd', 3), basesize + 3 * self.ssize_t)
2082 check(cwr(range(10), 4), basesize + 4 * self.ssize_t)
2083
2084 def test_permutations_sizeof(self):
2085 basesize = support.calcobjsize('4Pni')
2086 check = self.check_sizeof
2087 check(permutations('abcd'),
2088 basesize + 4 * self.ssize_t + 4 * self.ssize_t)
2089 check(permutations('abcd', 3),
2090 basesize + 4 * self.ssize_t + 3 * self.ssize_t)
2091 check(permutations('abcde', 3),
2092 basesize + 5 * self.ssize_t + 3 * self.ssize_t)
2093 check(permutations(range(10), 4),
2094 basesize + 10 * self.ssize_t + 4 * self.ssize_t)
2095
Thomas Woutersb2137042007-02-01 18:02:27 +00002096
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002097libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002098
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002099
2100>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002101>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00002102... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00002103...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002104Check 1200 is for $120.15
2105Check 1201 is for $764.05
2106Check 1202 is for $823.14
2107
2108>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002109>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00002110... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00002111...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000021121
21138
211427
2115
2116>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00002117>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00002118... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00002119...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002120Alex
2121Laura
2122Martin
2123Walter
2124Samuele
2125
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002126>>> from operator import itemgetter
2127>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00002128>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002129>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00002130... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00002131...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000021321 ['a', 'c', 'e']
21332 ['b', 'd', 'f']
21343 ['g']
2135
Raymond Hettinger734fb572004-01-20 20:04:40 +00002136# Find runs of consecutive numbers using groupby. The key to the solution
2137# is differencing with a range so that consecutive numbers all appear in
2138# same group.
2139>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00002140>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00002141... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00002142...
Raymond Hettinger734fb572004-01-20 20:04:40 +00002143[1]
2144[4, 5, 6]
2145[10]
2146[15, 16, 17, 18]
2147[22]
2148[25, 26, 27, 28]
2149
Georg Brandl3dbca812008-07-23 16:10:53 +00002150>>> def take(n, iterable):
2151... "Return first n items of the iterable as a list"
2152... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00002153
Georg Brandl3dbca812008-07-23 16:10:53 +00002154>>> def enumerate(iterable, start=0):
2155... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002156
Georg Brandl3dbca812008-07-23 16:10:53 +00002157>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002158... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00002159... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002160
Raymond Hettingercdf8ba32009-02-19 04:45:07 +00002161>>> def nth(iterable, n, default=None):
2162... "Returns the nth item or a default value"
2163... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002164
Raymond Hettingere525ee32016-03-06 18:11:38 -08002165>>> def all_equal(iterable):
2166... "Returns True if all the elements are equal to each other"
2167... g = groupby(iterable)
2168... return next(g, True) and not next(g, False)
2169
Georg Brandl3dbca812008-07-23 16:10:53 +00002170>>> def quantify(iterable, pred=bool):
2171... "Count how many times the predicate is true"
2172... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00002173
Georg Brandl3dbca812008-07-23 16:10:53 +00002174>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002175... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00002176... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002177
Georg Brandl3dbca812008-07-23 16:10:53 +00002178>>> def ncycles(iterable, n):
Ezio Melotti13925002011-03-16 11:05:33 +02002179... "Returns the sequence elements n times"
Georg Brandl3dbca812008-07-23 16:10:53 +00002180... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002181
2182>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002183... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002184
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002185>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002186... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002187
2188>>> def repeatfunc(func, times=None, *args):
2189... "Repeat calls to func with specified arguments."
2190... " Example: repeatfunc(random.random)"
2191... if times is None:
2192... return starmap(func, repeat(args))
2193... else:
2194... return starmap(func, repeat(args, times))
2195
Raymond Hettingerd591f662003-10-26 15:34:50 +00002196>>> def pairwise(iterable):
2197... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
2198... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00002199... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00002200... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00002201... except StopIteration:
2202... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002203... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002204
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002205>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00002206... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002207... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00002208... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002209
2210>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00002211... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002212... # Recipe credited to George Sakkis
2213... pending = len(iterables)
2214... nexts = cycle(iter(it).__next__ for it in iterables)
2215... while pending:
2216... try:
2217... for next in nexts:
2218... yield next()
2219... except StopIteration:
2220... pending -= 1
2221... nexts = cycle(islice(nexts, pending))
2222
2223>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00002224... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
2225... s = list(iterable)
2226... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002227
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00002228>>> def unique_everseen(iterable, key=None):
2229... "List unique elements, preserving order. Remember all elements ever seen."
2230... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
2231... # unique_everseen('ABBCcAD', str.lower) --> A B C D
2232... seen = set()
2233... seen_add = seen.add
2234... if key is None:
2235... for element in iterable:
2236... if element not in seen:
2237... seen_add(element)
2238... yield element
2239... else:
2240... for element in iterable:
2241... k = key(element)
2242... if k not in seen:
2243... seen_add(k)
2244... yield element
2245
2246>>> def unique_justseen(iterable, key=None):
2247... "List unique elements, preserving order. Remember only the element just seen."
2248... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
2249... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
2250... return map(next, map(itemgetter(1), groupby(iterable, key)))
2251
Raymond Hettinger31b26f62014-04-02 03:16:42 -07002252>>> def first_true(iterable, default=False, pred=None):
2253... '''Returns the first true value in the iterable.
2254...
2255... If no true value is found, returns *default*
2256...
2257... If *pred* is not None, returns the first item
2258... for which pred(item) is true.
2259...
2260... '''
2261... # first_true([a,b,c], x) --> a or b or c or x
2262... # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x
2263... return next(filter(pred, iterable), default)
2264
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002265This is not part of the examples but it tests to make sure the definitions
2266perform as purported.
2267
Raymond Hettingera098b332003-09-08 23:58:40 +00002268>>> take(10, count())
2269[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
2270
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002271>>> list(enumerate('abc'))
2272[(0, 'a'), (1, 'b'), (2, 'c')]
2273
2274>>> list(islice(tabulate(lambda x: 2*x), 4))
2275[0, 2, 4, 6]
2276
2277>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00002278'd'
2279
2280>>> nth('abcde', 9) is None
2281True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002282
Raymond Hettingere525ee32016-03-06 18:11:38 -08002283>>> [all_equal(s) for s in ('', 'A', 'AAAA', 'AAAB', 'AAABA')]
2284[True, True, True, False, False]
2285
Guido van Rossum805365e2007-05-07 22:24:25 +00002286>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000228750
2288
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002289>>> a = [[1, 2, 3], [4, 5, 6]]
2290>>> flatten(a)
2291[1, 2, 3, 4, 5, 6]
2292
2293>>> list(repeatfunc(pow, 5, 2, 3))
2294[8, 8, 8, 8, 8]
2295
2296>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002297>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002298[0, 0, 0, 0, 0]
2299
Raymond Hettingerd591f662003-10-26 15:34:50 +00002300>>> list(pairwise('abcd'))
2301[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002302
Raymond Hettingerd591f662003-10-26 15:34:50 +00002303>>> list(pairwise([]))
2304[]
2305
2306>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00002307[]
2308
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002309>>> list(islice(padnone('abc'), 0, 6))
2310['a', 'b', 'c', None, None, None]
2311
2312>>> list(ncycles('abc', 3))
2313['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
2314
2315>>> dotproduct([1,2,3], [4,5,6])
231632
2317
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002318>>> list(grouper(3, 'abcdefg', 'x'))
2319[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
2320
2321>>> list(roundrobin('abc', 'd', 'ef'))
2322['a', 'd', 'e', 'b', 'f', 'c']
2323
Raymond Hettingerace67332009-01-26 02:23:50 +00002324>>> list(powerset([1,2,3]))
2325[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002326
Raymond Hettinger191e8502009-01-27 13:29:43 +00002327>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
2328True
2329
2330>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
2331True
2332
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00002333>>> list(unique_everseen('AAAABBBCCDAABBB'))
2334['A', 'B', 'C', 'D']
2335
2336>>> list(unique_everseen('ABBCcAD', str.lower))
2337['A', 'B', 'C', 'D']
2338
2339>>> list(unique_justseen('AAAABBBCCDAABBB'))
2340['A', 'B', 'C', 'D', 'A', 'B']
2341
2342>>> list(unique_justseen('ABBCcAD', str.lower))
2343['A', 'B', 'C', 'A', 'D']
2344
Raymond Hettinger31b26f62014-04-02 03:16:42 -07002345>>> first_true('ABC0DEF1', '9', str.isdigit)
2346'0'
2347
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002348"""
2349
2350__test__ = {'libreftest' : libreftest}
2351
2352def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002353 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00002354 RegressionTests, LengthTransparency,
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002355 SubclassWithKwargsTest, TestExamples,
2356 SizeofTest)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002357 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002358
2359 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002360 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00002361 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00002362 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00002363 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002364 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00002365 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00002366 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00002367 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002368
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002369 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002370 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002371
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002372if __name__ == "__main__":
2373 test_main(verbose=True)