blob: 50cf1488ec6184b24ff368a4c329ddf4ded6fc85 [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
754 # Exercise pipes and filters style
755 s = 'abracadabra'
756 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000757 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000758 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
759 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000760 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000761 self.assertEqual(r, ['a', 'b', 'r'])
762 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000763 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000764 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
765 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000766 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000767 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
768
Georg Brandla18af4e2007-04-21 15:47:16 +0000769 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000770 class ExpectedError(Exception):
771 pass
772 def delayed_raise(n=0):
773 for i in range(n):
774 yield 'yo'
775 raise ExpectedError
776 def gulp(iterable, keyp=None, func=list):
777 return [func(g) for k, g in groupby(iterable, keyp)]
778
Georg Brandla18af4e2007-04-21 15:47:16 +0000779 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000780 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000781 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000782 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
783
Serhiy Storchakaa60c2fe2015-03-12 21:56:08 +0200784 # __eq__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000785 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000786 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000787 raise ExpectedError
788 s = [DummyCmp(), DummyCmp(), None]
789
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000790 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000791 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000792 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000793 self.assertRaises(ExpectedError, gulp, s)
794
795 # keyfunc failure
796 def keyfunc(obj):
797 if keyfunc.skip > 0:
798 keyfunc.skip -= 1
799 return obj
800 else:
801 raise ExpectedError
802
803 # keyfunc failure on outer object
804 keyfunc.skip = 0
805 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
806 keyfunc.skip = 1
807 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
808
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000809 def test_filter(self):
810 self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
811 self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])
812 self.assertEqual(list(filter(bool, [0,1,0,2,0])), [1,2])
813 self.assertEqual(take(4, filter(isEven, count())), [0,2,4,6])
814 self.assertRaises(TypeError, filter)
815 self.assertRaises(TypeError, filter, lambda x:x)
816 self.assertRaises(TypeError, filter, lambda x:x, range(6), 7)
817 self.assertRaises(TypeError, filter, isEven, 3)
818 self.assertRaises(TypeError, next, filter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000819
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000820 # check copy, deepcopy, pickle
821 ans = [0,2,4]
822
823 c = filter(isEven, range(6))
824 self.assertEqual(list(copy.copy(c)), ans)
825 c = filter(isEven, range(6))
826 self.assertEqual(list(copy.deepcopy(c)), ans)
Serhiy Storchakabad12572014-12-15 14:03:42 +0200827 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
828 c = filter(isEven, range(6))
829 self.assertEqual(list(pickle.loads(pickle.dumps(c, proto))), ans)
830 next(c)
831 self.assertEqual(list(pickle.loads(pickle.dumps(c, proto))), ans[1:])
832 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
833 c = filter(isEven, range(6))
834 self.pickletest(proto, c)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000835
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000836 def test_filterfalse(self):
837 self.assertEqual(list(filterfalse(isEven, range(6))), [1,3,5])
838 self.assertEqual(list(filterfalse(None, [0,1,0,2,0])), [0,0,0])
839 self.assertEqual(list(filterfalse(bool, [0,1,0,2,0])), [0,0,0])
840 self.assertEqual(take(4, filterfalse(isEven, count())), [1,3,5,7])
841 self.assertRaises(TypeError, filterfalse)
842 self.assertRaises(TypeError, filterfalse, lambda x:x)
843 self.assertRaises(TypeError, filterfalse, lambda x:x, range(6), 7)
844 self.assertRaises(TypeError, filterfalse, isEven, 3)
845 self.assertRaises(TypeError, next, filterfalse(range(6), range(6)))
Serhiy Storchakabad12572014-12-15 14:03:42 +0200846 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
847 self.pickletest(proto, filterfalse(isEven, range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000848
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000849 def test_zip(self):
850 # XXX This is rather silly now that builtin zip() calls zip()...
851 ans = [(x,y) for x, y in zip('abc',count())]
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000852 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000853 self.assertEqual(list(zip('abc', range(6))), lzip('abc', range(6)))
854 self.assertEqual(list(zip('abcdef', range(3))), lzip('abcdef', range(3)))
855 self.assertEqual(take(3,zip('abcdef', count())), lzip('abcdef', range(3)))
856 self.assertEqual(list(zip('abcdef')), lzip('abcdef'))
857 self.assertEqual(list(zip()), lzip())
858 self.assertRaises(TypeError, zip, 3)
859 self.assertRaises(TypeError, zip, range(3), 3)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000860 self.assertEqual([tuple(list(pair)) for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000861 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000862 self.assertEqual([pair for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000863 lzip('abc', 'def'))
Alex Gaynore151d212011-07-17 16:21:30 -0700864
865 @support.impl_detail("tuple reuse is specific to CPython")
866 def test_zip_tuple_reuse(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000867 ids = list(map(id, zip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000868 self.assertEqual(min(ids), max(ids))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000869 ids = list(map(id, list(zip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000870 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000871
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000872 # check copy, deepcopy, pickle
873 ans = [(x,y) for x, y in copy.copy(zip('abc',count()))]
874 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
875
876 ans = [(x,y) for x, y in copy.deepcopy(zip('abc',count()))]
877 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
878
Serhiy Storchakabad12572014-12-15 14:03:42 +0200879 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
880 ans = [(x,y) for x, y in pickle.loads(pickle.dumps(zip('abc',count()), proto))]
881 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000882
Serhiy Storchakabad12572014-12-15 14:03:42 +0200883 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
884 testIntermediate = zip('abc',count())
885 next(testIntermediate)
886 ans = [(x,y) for x, y in pickle.loads(pickle.dumps(testIntermediate, proto))]
887 self.assertEqual(ans, [('b', 1), ('c', 2)])
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000888
Serhiy Storchakabad12572014-12-15 14:03:42 +0200889 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
890 self.pickletest(proto, zip('abc', count()))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000891
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000892 def test_ziplongest(self):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000893 for args in [
894 ['abc', range(6)],
895 [range(6), 'abc'],
896 [range(1000), range(2000,2100), range(3000,3050)],
897 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
898 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
899 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000900 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
901 for i in range(max(map(len, args)))]
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000902 self.assertEqual(list(zip_longest(*args)), target)
903 self.assertEqual(list(zip_longest(*args, **{})), target)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000904 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 +0000905 self.assertEqual(list(zip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000906
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000907 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 +0000908
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000909 self.assertEqual(list(zip_longest()), list(zip()))
910 self.assertEqual(list(zip_longest([])), list(zip([])))
911 self.assertEqual(list(zip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000912
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000913 self.assertEqual(list(zip_longest('abc', 'defg', **{})),
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000914 list(zip(list('abc')+[None], 'defg'))) # empty keyword dict
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000915 self.assertRaises(TypeError, zip_longest, 3)
916 self.assertRaises(TypeError, zip_longest, range(3), 3)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000917
918 for stmt in [
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000919 "zip_longest('abc', fv=1)",
920 "zip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000921 ]:
922 try:
923 eval(stmt, globals(), locals())
924 except TypeError:
925 pass
926 else:
927 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000928
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000929 self.assertEqual([tuple(list(pair)) for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000930 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000931 self.assertEqual([pair for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000932 list(zip('abc', 'def')))
Alex Gaynore151d212011-07-17 16:21:30 -0700933
934 @support.impl_detail("tuple reuse is specific to CPython")
935 def test_zip_longest_tuple_reuse(self):
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000936 ids = list(map(id, zip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000937 self.assertEqual(min(ids), max(ids))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000938 ids = list(map(id, list(zip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000939 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
940
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000941 def test_zip_longest_pickling(self):
Serhiy Storchakabad12572014-12-15 14:03:42 +0200942 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
943 self.pickletest(proto, zip_longest("abc", "def"))
944 self.pickletest(proto, zip_longest("abc", "defgh"))
945 self.pickletest(proto, zip_longest("abc", "defgh", fillvalue=1))
946 self.pickletest(proto, zip_longest("", "defgh"))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000947
Raymond Hettingerfc438512009-11-01 20:55:33 +0000948 def test_bug_7244(self):
949
950 class Repeater:
951 # this class is similar to itertools.repeat
952 def __init__(self, o, t, e):
953 self.o = o
954 self.t = int(t)
955 self.e = e
956 def __iter__(self): # its iterator is itself
957 return self
958 def __next__(self):
959 if self.t > 0:
960 self.t -= 1
961 return self.o
962 else:
963 raise self.e
964
965 # Formerly this code in would fail in debug mode
966 # with Undetected Error and Stop Iteration
967 r1 = Repeater(1, 3, StopIteration)
968 r2 = Repeater(2, 4, StopIteration)
969 def run(r1, r2):
970 result = []
971 for i, j in zip_longest(r1, r2, fillvalue=0):
972 with support.captured_output('stdout'):
973 print((i, j))
974 result.append((i, j))
975 return result
976 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
977
978 # Formerly, the RuntimeError would be lost
979 # and StopIteration would stop as expected
980 r1 = Repeater(1, 3, RuntimeError)
981 r2 = Repeater(2, 4, StopIteration)
982 it = zip_longest(r1, r2, fillvalue=0)
983 self.assertEqual(next(it), (1, 2))
984 self.assertEqual(next(it), (1, 2))
985 self.assertEqual(next(it), (1, 2))
986 self.assertRaises(RuntimeError, next, it)
987
Christian Heimesc3f30c42008-02-22 16:37:40 +0000988 def test_product(self):
989 for args, result in [
Christian Heimes78644762008-03-04 23:39:23 +0000990 ([], [()]), # zero iterables
Christian Heimesc3f30c42008-02-22 16:37:40 +0000991 (['ab'], [('a',), ('b',)]), # one iterable
992 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
993 ([range(0), range(2), range(3)], []), # first iterable with zero length
994 ([range(2), range(0), range(3)], []), # middle iterable with zero length
995 ([range(2), range(3), range(0)], []), # last iterable with zero length
996 ]:
997 self.assertEqual(list(product(*args)), result)
Christian Heimesf16baeb2008-02-29 14:57:44 +0000998 for r in range(4):
999 self.assertEqual(list(product(*(args*r))),
1000 list(product(*args, **dict(repeat=r))))
Christian Heimesc3f30c42008-02-22 16:37:40 +00001001 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
1002 self.assertRaises(TypeError, product, range(6), None)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001003
1004 def product1(*args, **kwds):
1005 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
1006 n = len(pools)
1007 if n == 0:
1008 yield ()
1009 return
1010 if any(len(pool) == 0 for pool in pools):
1011 return
1012 indices = [0] * n
1013 yield tuple(pool[i] for pool, i in zip(pools, indices))
1014 while 1:
1015 for i in reversed(range(n)): # right to left
1016 if indices[i] == len(pools[i]) - 1:
1017 continue
1018 indices[i] += 1
1019 for j in range(i+1, n):
1020 indices[j] = 0
1021 yield tuple(pool[i] for pool, i in zip(pools, indices))
1022 break
1023 else:
1024 return
1025
1026 def product2(*args, **kwds):
1027 'Pure python version used in docs'
1028 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
1029 result = [[]]
1030 for pool in pools:
1031 result = [x+[y] for x in result for y in pool]
1032 for prod in result:
1033 yield tuple(prod)
1034
Christian Heimesc3f30c42008-02-22 16:37:40 +00001035 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
1036 set('abcdefg'), range(11), tuple(range(13))]
1037 for i in range(100):
1038 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Christian Heimes78644762008-03-04 23:39:23 +00001039 expected_len = prod(map(len, args))
1040 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001041 self.assertEqual(list(product(*args)), list(product1(*args)))
1042 self.assertEqual(list(product(*args)), list(product2(*args)))
Christian Heimesc3f30c42008-02-22 16:37:40 +00001043 args = map(iter, args)
Christian Heimes78644762008-03-04 23:39:23 +00001044 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesc3f30c42008-02-22 16:37:40 +00001045
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05001046 @support.bigaddrspacetest
1047 def test_product_overflow(self):
Serhiy Storchakadee948b2015-02-03 01:34:09 +02001048 with self.assertRaises((OverflowError, MemoryError)):
1049 product(*(['ab']*2**5), repeat=2**25)
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05001050
Alex Gaynore151d212011-07-17 16:21:30 -07001051 @support.impl_detail("tuple reuse is specific to CPython")
1052 def test_product_tuple_reuse(self):
Christian Heimes90c3d9b2008-02-23 13:18:03 +00001053 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
1054 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +00001055
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001056 def test_product_pickling(self):
1057 # check copy, deepcopy, pickle
1058 for args, result in [
1059 ([], [()]), # zero iterables
1060 (['ab'], [('a',), ('b',)]), # one iterable
1061 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
1062 ([range(0), range(2), range(3)], []), # first iterable with zero length
1063 ([range(2), range(0), range(3)], []), # middle iterable with zero length
1064 ([range(2), range(3), range(0)], []), # last iterable with zero length
1065 ]:
1066 self.assertEqual(list(copy.copy(product(*args))), result)
1067 self.assertEqual(list(copy.deepcopy(product(*args))), result)
Serhiy Storchakabad12572014-12-15 14:03:42 +02001068 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1069 self.pickletest(proto, product(*args))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001070
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00001071 def test_product_issue_25021(self):
1072 # test that indices are properly clamped to the length of the tuples
1073 p = product((1, 2),(3,))
1074 p.__setstate__((0, 0x1000)) # will access tuple element 1 if not clamped
1075 self.assertEqual(next(p), (2, 3))
1076 # test that empty tuple in the list will result in an immediate StopIteration
1077 p = product((1, 2), (), (3,))
1078 p.__setstate__((0, 0, 0x1000)) # will access tuple element 1 if not clamped
1079 self.assertRaises(StopIteration, next, p)
1080
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001081 def test_repeat(self):
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00001082 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
Guido van Rossum805365e2007-05-07 22:24:25 +00001083 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001084 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001085 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +00001086 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001087 self.assertEqual(list(repeat('a', 0)), [])
1088 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001089 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001090 self.assertRaises(TypeError, repeat, None, 3, 4)
1091 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001092 r = repeat(1+0j)
1093 self.assertEqual(repr(r), 'repeat((1+0j))')
1094 r = repeat(1+0j, 5)
1095 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
1096 list(r)
1097 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001098
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001099 # check copy, deepcopy, pickle
1100 c = repeat(object='a', times=10)
1101 self.assertEqual(next(c), 'a')
1102 self.assertEqual(take(2, copy.copy(c)), list('a' * 2))
1103 self.assertEqual(take(2, copy.deepcopy(c)), list('a' * 2))
Serhiy Storchakabad12572014-12-15 14:03:42 +02001104 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1105 self.pickletest(proto, repeat(object='a', times=10))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001106
Raymond Hettinger97d35552014-06-24 21:36:58 -07001107 def test_repeat_with_negative_times(self):
1108 self.assertEqual(repr(repeat('a', -1)), "repeat('a', 0)")
1109 self.assertEqual(repr(repeat('a', -2)), "repeat('a', 0)")
1110 self.assertEqual(repr(repeat('a', times=-1)), "repeat('a', 0)")
1111 self.assertEqual(repr(repeat('a', times=-2)), "repeat('a', 0)")
1112
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001113 def test_map(self):
1114 self.assertEqual(list(map(operator.pow, range(3), range(1,7))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001115 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001116 self.assertEqual(list(map(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001117 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001118 self.assertEqual(list(map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +00001119 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001120 self.assertEqual(take(2,map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +00001121 [('a',0),('b',1)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001122 self.assertEqual(list(map(operator.pow, [])), [])
1123 self.assertRaises(TypeError, map)
1124 self.assertRaises(TypeError, list, map(None, range(3), range(3)))
1125 self.assertRaises(TypeError, map, operator.neg)
1126 self.assertRaises(TypeError, next, map(10, range(5)))
1127 self.assertRaises(ValueError, next, map(errfunc, [4], [5]))
1128 self.assertRaises(TypeError, next, map(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001129
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001130 # check copy, deepcopy, pickle
1131 ans = [('a',0),('b',1),('c',2)]
1132
1133 c = map(tupleize, 'abc', count())
1134 self.assertEqual(list(copy.copy(c)), ans)
1135
1136 c = map(tupleize, 'abc', count())
1137 self.assertEqual(list(copy.deepcopy(c)), ans)
1138
Serhiy Storchakabad12572014-12-15 14:03:42 +02001139 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1140 c = map(tupleize, 'abc', count())
1141 self.pickletest(proto, c)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001142
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001143 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001144 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
1145 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001146 self.assertEqual(take(3, starmap(operator.pow, zip(count(), count(1)))),
Raymond Hettinger02420702003-06-29 20:36:23 +00001147 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001148 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +00001149 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
1150 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001151 self.assertRaises(TypeError, starmap)
1152 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001153 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
1154 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
1155 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001156
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001157 # check copy, deepcopy, pickle
1158 ans = [0**1, 1**2, 2**3]
1159
1160 c = starmap(operator.pow, zip(range(3), range(1,7)))
1161 self.assertEqual(list(copy.copy(c)), ans)
1162
1163 c = starmap(operator.pow, zip(range(3), range(1,7)))
1164 self.assertEqual(list(copy.deepcopy(c)), ans)
1165
Serhiy Storchakabad12572014-12-15 14:03:42 +02001166 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1167 c = starmap(operator.pow, zip(range(3), range(1,7)))
1168 self.pickletest(proto, c)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001169
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001170 def test_islice(self):
1171 for args in [ # islice(args) should agree with range(args)
1172 (10, 20, 3),
1173 (10, 3, 20),
1174 (10, 20),
1175 (10, 3),
1176 (20,)
1177 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +00001178 self.assertEqual(list(islice(range(100), *args)),
1179 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001180
1181 for args, tgtargs in [ # Stop when seqn is exhausted
1182 ((10, 110, 3), ((10, 100, 3))),
1183 ((10, 110), ((10, 100))),
1184 ((110,), (100,))
1185 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +00001186 self.assertEqual(list(islice(range(100), *args)),
1187 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001188
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001189 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +00001190 self.assertEqual(list(islice(range(10), None)), list(range(10)))
1191 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
1192 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
1193 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
1194 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001195
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +00001196 # Test number of items consumed SF #1171417
1197 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +00001198 self.assertEqual(list(islice(it, 3)), list(range(3)))
1199 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +00001200
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001201 # Test invalid arguments
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001202 ra = range(10)
1203 self.assertRaises(TypeError, islice, ra)
1204 self.assertRaises(TypeError, islice, ra, 1, 2, 3, 4)
1205 self.assertRaises(ValueError, islice, ra, -5, 10, 1)
1206 self.assertRaises(ValueError, islice, ra, 1, -5, -1)
1207 self.assertRaises(ValueError, islice, ra, 1, 10, -1)
1208 self.assertRaises(ValueError, islice, ra, 1, 10, 0)
1209 self.assertRaises(ValueError, islice, ra, 'a')
1210 self.assertRaises(ValueError, islice, ra, 'a', 1)
1211 self.assertRaises(ValueError, islice, ra, 1, 'a')
1212 self.assertRaises(ValueError, islice, ra, 'a', 1, 1)
1213 self.assertRaises(ValueError, islice, ra, 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +00001214 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001215
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001216 # Issue #10323: Less islice in a predictable state
1217 c = count()
1218 self.assertEqual(list(islice(c, 1, 3, 50)), [1])
1219 self.assertEqual(next(c), 3)
1220
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001221 # check copy, deepcopy, pickle
1222 for args in [ # islice(args) should agree with range(args)
1223 (10, 20, 3),
1224 (10, 3, 20),
1225 (10, 20),
1226 (10, 3),
1227 (20,)
1228 ]:
1229 self.assertEqual(list(copy.copy(islice(range(100), *args))),
1230 list(range(*args)))
1231 self.assertEqual(list(copy.deepcopy(islice(range(100), *args))),
1232 list(range(*args)))
Serhiy Storchakabad12572014-12-15 14:03:42 +02001233 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1234 self.pickletest(proto, islice(range(100), *args))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001235
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001236 # Issue #21321: check source iterator is not referenced
1237 # from islice() after the latter has been exhausted
1238 it = (x for x in (1, 2))
1239 wr = weakref.ref(it)
1240 it = islice(it, 1)
1241 self.assertIsNotNone(wr())
1242 list(it) # exhaust the iterator
Benjamin Peterson8e163512014-08-24 18:07:28 -05001243 support.gc_collect()
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001244 self.assertIsNone(wr())
1245
Will Roberts0ecdc522017-06-08 08:03:04 +02001246 # Issue #30537: islice can accept integer-like objects as
1247 # arguments
1248 class IntLike(object):
1249 def __init__(self, val):
1250 self.val = val
1251 def __index__(self):
1252 return self.val
1253 self.assertEqual(list(islice(range(100), IntLike(10))), list(range(10)))
1254 self.assertEqual(list(islice(range(100), IntLike(10), IntLike(50))),
1255 list(range(10, 50)))
1256 self.assertEqual(list(islice(range(100), IntLike(10), IntLike(50), IntLike(5))),
1257 list(range(10,50,5)))
1258
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001259 def test_takewhile(self):
1260 data = [1, 3, 5, 20, 2, 4, 6, 8]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001261 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001262 self.assertEqual(list(takewhile(underten, [])), [])
1263 self.assertRaises(TypeError, takewhile)
1264 self.assertRaises(TypeError, takewhile, operator.pow)
1265 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001266 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
1267 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001268 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
1269 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +00001270 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001271
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001272 # check copy, deepcopy, pickle
1273 self.assertEqual(list(copy.copy(takewhile(underten, data))), [1, 3, 5])
1274 self.assertEqual(list(copy.deepcopy(takewhile(underten, data))),
1275 [1, 3, 5])
Serhiy Storchakabad12572014-12-15 14:03:42 +02001276 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1277 self.pickletest(proto, takewhile(underten, data))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001278
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001279 def test_dropwhile(self):
1280 data = [1, 3, 5, 20, 2, 4, 6, 8]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001281 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001282 self.assertEqual(list(dropwhile(underten, [])), [])
1283 self.assertRaises(TypeError, dropwhile)
1284 self.assertRaises(TypeError, dropwhile, operator.pow)
1285 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001286 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
1287 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001288
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001289 # check copy, deepcopy, pickle
1290 self.assertEqual(list(copy.copy(dropwhile(underten, data))), [20, 2, 4, 6, 8])
1291 self.assertEqual(list(copy.deepcopy(dropwhile(underten, data))),
1292 [20, 2, 4, 6, 8])
Serhiy Storchakabad12572014-12-15 14:03:42 +02001293 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1294 self.pickletest(proto, dropwhile(underten, data))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001295
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001296 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +00001297 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001298
1299 a, b = tee([]) # test empty iterator
1300 self.assertEqual(list(a), [])
1301 self.assertEqual(list(b), [])
1302
1303 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +00001304 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001305
1306 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +00001307 self.assertEqual(list(a), list(range(n)))
1308 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001309
1310 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +00001311 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001312 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001313 del a
Guido van Rossum805365e2007-05-07 22:24:25 +00001314 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001315
1316 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +00001317 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001318 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001319 del b
Guido van Rossum805365e2007-05-07 22:24:25 +00001320 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001321
Guido van Rossum805365e2007-05-07 22:24:25 +00001322 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001323 order = [0]*n + [1]*n
1324 random.shuffle(order)
1325 lists = ([], [])
1326 its = tee(irange(n))
1327 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +00001328 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001329 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +00001330 self.assertEqual(lists[0], list(range(n)))
1331 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001332
Raymond Hettingerad983e72003-11-12 14:32:26 +00001333 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001334 self.assertRaises(TypeError, tee)
1335 self.assertRaises(TypeError, tee, 3)
1336 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +00001337 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001338
Raymond Hettingerad983e72003-11-12 14:32:26 +00001339 # tee object should be instantiable
1340 a, b = tee('abc')
1341 c = type(a)('def')
1342 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +00001343
Raymond Hettingerad983e72003-11-12 14:32:26 +00001344 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +00001345 a, b, c = tee(range(2000), 3)
1346 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001347 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +00001348 self.assertEqual(list(b), list(range(2000)))
1349 self.assertEqual([next(c), next(c)], list(range(2)))
1350 self.assertEqual(list(a), list(range(100,2000)))
1351 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +00001352
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001353 # test values of n
1354 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +00001355 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +00001356 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001357 result = tee('abc', n)
1358 self.assertEqual(type(result), tuple)
1359 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001360 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001361
Raymond Hettingerad983e72003-11-12 14:32:26 +00001362 # tee pass-through to copyable iterator
1363 a, b = tee('abc')
1364 c, d = tee(a)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +00001365 self.assertTrue(a is c)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001366
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001367 # test tee_new
1368 t1, t2 = tee('abc')
1369 tnew = type(t1)
1370 self.assertRaises(TypeError, tnew)
1371 self.assertRaises(TypeError, tnew, 10)
1372 t3 = tnew(t1)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +00001373 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +00001374
Raymond Hettingera9f60922004-10-17 16:40:14 +00001375 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +00001376 a, b = tee(range(10))
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001377 p = weakref.proxy(a)
Raymond Hettingera9f60922004-10-17 16:40:14 +00001378 self.assertEqual(getattr(p, '__class__'), type(b))
1379 del a
1380 self.assertRaises(ReferenceError, getattr, p, '__class__')
1381
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001382 ans = list('abc')
1383 long_ans = list(range(10000))
1384
1385 # check copy
1386 a, b = tee('abc')
1387 self.assertEqual(list(copy.copy(a)), ans)
1388 self.assertEqual(list(copy.copy(b)), ans)
1389 a, b = tee(list(range(10000)))
1390 self.assertEqual(list(copy.copy(a)), long_ans)
1391 self.assertEqual(list(copy.copy(b)), long_ans)
1392
1393 # check partially consumed copy
1394 a, b = tee('abc')
1395 take(2, a)
1396 take(1, b)
1397 self.assertEqual(list(copy.copy(a)), ans[2:])
1398 self.assertEqual(list(copy.copy(b)), ans[1:])
1399 self.assertEqual(list(a), ans[2:])
1400 self.assertEqual(list(b), ans[1:])
1401 a, b = tee(range(10000))
1402 take(100, a)
1403 take(60, b)
1404 self.assertEqual(list(copy.copy(a)), long_ans[100:])
1405 self.assertEqual(list(copy.copy(b)), long_ans[60:])
1406 self.assertEqual(list(a), long_ans[100:])
1407 self.assertEqual(list(b), long_ans[60:])
1408
1409 # check deepcopy
1410 a, b = tee('abc')
1411 self.assertEqual(list(copy.deepcopy(a)), ans)
1412 self.assertEqual(list(copy.deepcopy(b)), ans)
1413 self.assertEqual(list(a), ans)
1414 self.assertEqual(list(b), ans)
1415 a, b = tee(range(10000))
1416 self.assertEqual(list(copy.deepcopy(a)), long_ans)
1417 self.assertEqual(list(copy.deepcopy(b)), long_ans)
1418 self.assertEqual(list(a), long_ans)
1419 self.assertEqual(list(b), long_ans)
1420
1421 # check partially consumed deepcopy
1422 a, b = tee('abc')
1423 take(2, a)
1424 take(1, b)
1425 self.assertEqual(list(copy.deepcopy(a)), ans[2:])
1426 self.assertEqual(list(copy.deepcopy(b)), ans[1:])
1427 self.assertEqual(list(a), ans[2:])
1428 self.assertEqual(list(b), ans[1:])
1429 a, b = tee(range(10000))
1430 take(100, a)
1431 take(60, b)
1432 self.assertEqual(list(copy.deepcopy(a)), long_ans[100:])
1433 self.assertEqual(list(copy.deepcopy(b)), long_ans[60:])
1434 self.assertEqual(list(a), long_ans[100:])
1435 self.assertEqual(list(b), long_ans[60:])
1436
1437 # check pickle
Serhiy Storchakabad12572014-12-15 14:03:42 +02001438 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1439 self.pickletest(proto, iter(tee('abc')))
1440 a, b = tee('abc')
1441 self.pickletest(proto, a, compare=ans)
1442 self.pickletest(proto, b, compare=ans)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001443
Serhiy Storchakaa3e91282013-01-25 13:19:31 +02001444 # Issue 13454: Crash when deleting backward iterator from tee()
1445 def test_tee_del_backward(self):
Serhiy Storchaka5bb893c2013-01-26 11:52:06 +02001446 forward, backward = tee(repeat(None, 20000000))
Serhiy Storchaka9db55002015-03-28 20:38:37 +02001447 try:
1448 any(forward) # exhaust the iterator
1449 del backward
1450 except:
1451 del forward, backward
1452 raise
Serhiy Storchakaa3e91282013-01-25 13:19:31 +02001453
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001454 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001455 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +00001456
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001457 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +00001458 self.assertRaises(StopIteration, next, f([]))
1459 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001460
Georg Brandla18af4e2007-04-21 15:47:16 +00001461 self.assertRaises(StopIteration, next, islice([], None))
1462 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001463
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001464 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +00001465 self.assertRaises(StopIteration, next, p)
1466 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001467 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +00001468 self.assertRaises(StopIteration, next, p)
1469 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001470
Georg Brandla18af4e2007-04-21 15:47:16 +00001471 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001472
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001473 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +00001474 self.assertRaises(StopIteration, next, f(lambda x:x, []))
1475 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001476
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001477class TestExamples(unittest.TestCase):
1478
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001479 def test_accumulate(self):
Raymond Hettinger482ba772010-12-01 22:48:00 +00001480 self.assertEqual(list(accumulate([1,2,3,4,5])), [1, 3, 6, 10, 15])
1481
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001482 def test_accumulate_reducible(self):
1483 # check copy, deepcopy, pickle
1484 data = [1, 2, 3, 4, 5]
1485 accumulated = [1, 3, 6, 10, 15]
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001486
Serhiy Storchakabad12572014-12-15 14:03:42 +02001487 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1488 it = accumulate(data)
1489 self.assertEqual(list(pickle.loads(pickle.dumps(it, proto))), accumulated[:])
1490 self.assertEqual(next(it), 1)
1491 self.assertEqual(list(pickle.loads(pickle.dumps(it, proto))), accumulated[1:])
1492 it = accumulate(data)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001493 self.assertEqual(next(it), 1)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001494 self.assertEqual(list(copy.deepcopy(it)), accumulated[1:])
1495 self.assertEqual(list(copy.copy(it)), accumulated[1:])
1496
Serhiy Storchakad5516252016-03-06 14:00:45 +02001497 def test_accumulate_reducible_none(self):
1498 # Issue #25718: total is None
1499 it = accumulate([None, None, None], operator.is_)
1500 self.assertEqual(next(it), None)
1501 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
1502 it_copy = pickle.loads(pickle.dumps(it, proto))
1503 self.assertEqual(list(it_copy), [True, False])
1504 self.assertEqual(list(copy.deepcopy(it)), [True, False])
1505 self.assertEqual(list(copy.copy(it)), [True, False])
1506
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001507 def test_chain(self):
1508 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
1509
1510 def test_chain_from_iterable(self):
1511 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
1512
1513 def test_combinations(self):
1514 self.assertEqual(list(combinations('ABCD', 2)),
1515 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
1516 self.assertEqual(list(combinations(range(4), 3)),
1517 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
1518
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001519 def test_combinations_with_replacement(self):
1520 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
1521 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
1522
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001523 def test_compress(self):
1524 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
1525
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001526 def test_count(self):
1527 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
1528
1529 def test_cycle(self):
1530 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
1531
1532 def test_dropwhile(self):
1533 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
1534
1535 def test_groupby(self):
1536 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
1537 list('ABCDAB'))
1538 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
1539 [list('AAAA'), list('BBB'), list('CC'), list('D')])
1540
1541 def test_filter(self):
1542 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
1543
1544 def test_filterfalse(self):
1545 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
1546
1547 def test_map(self):
1548 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
1549
1550 def test_islice(self):
1551 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
1552 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
1553 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
1554 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1555
1556 def test_zip(self):
1557 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
1558
1559 def test_zip_longest(self):
1560 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
1561 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
1562
1563 def test_permutations(self):
1564 self.assertEqual(list(permutations('ABCD', 2)),
1565 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
1566 self.assertEqual(list(permutations(range(3))),
1567 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
1568
1569 def test_product(self):
1570 self.assertEqual(list(product('ABCD', 'xy')),
1571 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
1572 self.assertEqual(list(product(range(2), repeat=3)),
1573 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
1574 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
1575
1576 def test_repeat(self):
1577 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
1578
1579 def test_stapmap(self):
1580 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
1581 [32, 9, 1000])
1582
1583 def test_takewhile(self):
1584 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
1585
1586
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001587class TestGC(unittest.TestCase):
1588
1589 def makecycle(self, iterator, container):
1590 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +00001591 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001592 del container, iterator
1593
Raymond Hettinger482ba772010-12-01 22:48:00 +00001594 def test_accumulate(self):
1595 a = []
1596 self.makecycle(accumulate([1,2,a,3]), a)
1597
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001598 def test_chain(self):
1599 a = []
1600 self.makecycle(chain(a), a)
1601
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001602 def test_chain_from_iterable(self):
1603 a = []
1604 self.makecycle(chain.from_iterable([a]), a)
1605
1606 def test_combinations(self):
1607 a = []
1608 self.makecycle(combinations([1,2,a,3], 3), a)
1609
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001610 def test_combinations_with_replacement(self):
1611 a = []
1612 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1613
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001614 def test_compress(self):
1615 a = []
1616 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1617
Raymond Hettingerd6280f42009-02-16 20:50:56 +00001618 def test_count(self):
1619 a = []
1620 Int = type('Int', (int,), dict(x=a))
1621 self.makecycle(count(Int(0), Int(1)), a)
1622
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001623 def test_cycle(self):
1624 a = []
1625 self.makecycle(cycle([a]*2), a)
1626
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001627 def test_dropwhile(self):
1628 a = []
1629 self.makecycle(dropwhile(bool, [0, a, a]), a)
1630
1631 def test_groupby(self):
1632 a = []
1633 self.makecycle(groupby([a]*2, lambda x:x), a)
1634
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001635 def test_issue2246(self):
1636 # Issue 2246 -- the _grouper iterator was not included in GC
1637 n = 10
1638 keyfunc = lambda x: x
1639 for i, j in groupby(range(n), key=keyfunc):
1640 keyfunc.__dict__.setdefault('x',[]).append(j)
1641
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001642 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001643 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001644 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001645
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001646 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001647 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001648 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001649
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001650 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001651 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001652 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001653
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001654 def test_zip_longest(self):
1655 a = []
1656 self.makecycle(zip_longest([a]*2, [a]*3), a)
1657 b = [a, None]
1658 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1659
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001660 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001661 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001662 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001663
1664 def test_islice(self):
1665 a = []
1666 self.makecycle(islice([a]*2, None), a)
1667
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001668 def test_permutations(self):
1669 a = []
1670 self.makecycle(permutations([1,2,a,3], 3), a)
1671
1672 def test_product(self):
1673 a = []
1674 self.makecycle(product([1,2,a,3], repeat=3), a)
1675
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001676 def test_repeat(self):
1677 a = []
1678 self.makecycle(repeat(a), a)
1679
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001680 def test_starmap(self):
1681 a = []
1682 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1683
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001684 def test_takewhile(self):
1685 a = []
1686 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1687
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001688def R(seqn):
1689 'Regular generator'
1690 for i in seqn:
1691 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001692
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001693class G:
1694 'Sequence using __getitem__'
1695 def __init__(self, seqn):
1696 self.seqn = seqn
1697 def __getitem__(self, i):
1698 return self.seqn[i]
1699
1700class I:
1701 'Sequence using iterator protocol'
1702 def __init__(self, seqn):
1703 self.seqn = seqn
1704 self.i = 0
1705 def __iter__(self):
1706 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001707 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001708 if self.i >= len(self.seqn): raise StopIteration
1709 v = self.seqn[self.i]
1710 self.i += 1
1711 return v
1712
1713class Ig:
1714 'Sequence using iterator protocol defined with a generator'
1715 def __init__(self, seqn):
1716 self.seqn = seqn
1717 self.i = 0
1718 def __iter__(self):
1719 for val in self.seqn:
1720 yield val
1721
1722class X:
1723 'Missing __getitem__ and __iter__'
1724 def __init__(self, seqn):
1725 self.seqn = seqn
1726 self.i = 0
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 N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001734 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001735 def __init__(self, seqn):
1736 self.seqn = seqn
1737 self.i = 0
1738 def __iter__(self):
1739 return self
1740
1741class E:
1742 'Test propagation of exceptions'
1743 def __init__(self, seqn):
1744 self.seqn = seqn
1745 self.i = 0
1746 def __iter__(self):
1747 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001748 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001749 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001750
1751class S:
1752 'Test immediate stop'
1753 def __init__(self, seqn):
1754 pass
1755 def __iter__(self):
1756 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001757 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001758 raise StopIteration
1759
1760def L(seqn):
1761 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001762 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001763
1764
1765class TestVariousIteratorArgs(unittest.TestCase):
1766
Raymond Hettinger482ba772010-12-01 22:48:00 +00001767 def test_accumulate(self):
1768 s = [1,2,3,4,5]
1769 r = [1,3,6,10,15]
1770 n = len(s)
1771 for g in (G, I, Ig, L, R):
1772 self.assertEqual(list(accumulate(g(s))), r)
1773 self.assertEqual(list(accumulate(S(s))), [])
1774 self.assertRaises(TypeError, accumulate, X(s))
1775 self.assertRaises(TypeError, accumulate, N(s))
1776 self.assertRaises(ZeroDivisionError, list, accumulate(E(s)))
1777
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001778 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001779 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001780 for g in (G, I, Ig, S, L, R):
1781 self.assertEqual(list(chain(g(s))), list(g(s)))
1782 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001783 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001784 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001785 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1786
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001787 def test_compress(self):
1788 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1789 n = len(s)
1790 for g in (G, I, Ig, S, L, R):
1791 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1792 self.assertRaises(TypeError, compress, X(s), repeat(1))
1793 self.assertRaises(TypeError, compress, N(s), repeat(1))
1794 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1795
Christian Heimesc3f30c42008-02-22 16:37:40 +00001796 def test_product(self):
1797 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1798 self.assertRaises(TypeError, product, X(s))
1799 self.assertRaises(TypeError, product, N(s))
1800 self.assertRaises(ZeroDivisionError, product, E(s))
1801
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001802 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001803 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001804 for g in (G, I, Ig, S, L, R):
1805 tgtlen = len(s) * 3
1806 expected = list(g(s))*3
1807 actual = list(islice(cycle(g(s)), tgtlen))
1808 self.assertEqual(actual, expected)
1809 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001810 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001811 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1812
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001813 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001814 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001815 for g in (G, I, Ig, S, L, R):
1816 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1817 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001818 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001819 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1820
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001821 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001822 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001823 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001824 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001825 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001826 self.assertRaises(TypeError, filter, isEven, X(s))
1827 self.assertRaises(TypeError, filter, isEven, N(s))
1828 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001829
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001830 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001831 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001832 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001833 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001834 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001835 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1836 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1837 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001838
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001839 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001840 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001841 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001842 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1843 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1844 self.assertRaises(TypeError, zip, X(s))
1845 self.assertRaises(TypeError, zip, N(s))
1846 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001847
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001848 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001849 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001850 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001851 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1852 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1853 self.assertRaises(TypeError, zip_longest, X(s))
1854 self.assertRaises(TypeError, zip_longest, N(s))
1855 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001856
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001857 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001858 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001859 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001860 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001861 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001862 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001863 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001864 self.assertRaises(TypeError, map, onearg, X(s))
1865 self.assertRaises(TypeError, map, onearg, N(s))
1866 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001867
1868 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001869 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001870 for g in (G, I, Ig, S, L, R):
1871 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1872 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001873 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001874 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1875
1876 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001877 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001878 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001879 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001880 self.assertEqual(list(starmap(operator.pow, g(ss))),
1881 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001882 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001883 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001884 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1885
1886 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001887 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001888 for g in (G, I, Ig, S, L, R):
1889 tgt = []
1890 for elem in g(s):
1891 if not isEven(elem): break
1892 tgt.append(elem)
1893 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1894 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001895 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001896 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1897
1898 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001899 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001900 for g in (G, I, Ig, S, L, R):
1901 tgt = []
1902 for elem in g(s):
1903 if not tgt and isOdd(elem): continue
1904 tgt.append(elem)
1905 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1906 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001907 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001908 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1909
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001910 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001911 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001912 for g in (G, I, Ig, S, L, R):
1913 it1, it2 = tee(g(s))
1914 self.assertEqual(list(it1), list(g(s)))
1915 self.assertEqual(list(it2), list(g(s)))
1916 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001917 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001918 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1919
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001920class LengthTransparency(unittest.TestCase):
1921
1922 def test_repeat(self):
Armin Ronacheraa9a79d2012-10-06 14:03:24 +02001923 self.assertEqual(operator.length_hint(repeat(None, 50)), 50)
Raymond Hettinger97d35552014-06-24 21:36:58 -07001924 self.assertEqual(operator.length_hint(repeat(None, 0)), 0)
Armin Ronacheraa9a79d2012-10-06 14:03:24 +02001925 self.assertEqual(operator.length_hint(repeat(None), 12), 12)
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001926
Raymond Hettinger97d35552014-06-24 21:36:58 -07001927 def test_repeat_with_negative_times(self):
1928 self.assertEqual(operator.length_hint(repeat(None, -1)), 0)
1929 self.assertEqual(operator.length_hint(repeat(None, -2)), 0)
1930 self.assertEqual(operator.length_hint(repeat(None, times=-1)), 0)
1931 self.assertEqual(operator.length_hint(repeat(None, times=-2)), 0)
1932
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001933class RegressionTests(unittest.TestCase):
1934
1935 def test_sf_793826(self):
1936 # Fix Armin Rigo's successful efforts to wreak havoc
1937
1938 def mutatingtuple(tuple1, f, tuple2):
1939 # this builds a tuple t which is a copy of tuple1,
1940 # then calls f(t), then mutates t to be equal to tuple2
1941 # (needs len(tuple1) == len(tuple2)).
1942 def g(value, first=[1]):
1943 if first:
1944 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00001945 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001946 return value
1947 items = list(tuple2)
1948 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001949 gen = map(g, items)
1950 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00001951 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001952
1953 def f(t):
1954 global T
1955 T = t
1956 first[:] = list(T)
1957
1958 first = []
1959 mutatingtuple((1,2,3), f, (4,5,6))
1960 second = list(T)
1961 self.assertEqual(first, second)
1962
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001963
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001964 def test_sf_950057(self):
1965 # Make sure that chain() and cycle() catch exceptions immediately
1966 # rather than when shifting between input sources
1967
1968 def gen1():
1969 hist.append(0)
1970 yield 1
1971 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001972 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001973 hist.append(2)
1974
1975 def gen2(x):
1976 hist.append(3)
1977 yield 2
1978 hist.append(4)
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001979
1980 hist = []
1981 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1982 self.assertEqual(hist, [0,1])
1983
1984 hist = []
1985 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1986 self.assertEqual(hist, [0,1])
1987
1988 hist = []
1989 self.assertRaises(AssertionError, list, cycle(gen1()))
1990 self.assertEqual(hist, [0,1])
1991
T. Wouters5466d4a2017-03-30 09:58:35 -07001992 def test_long_chain_of_empty_iterables(self):
1993 # Make sure itertools.chain doesn't run into recursion limits when
1994 # dealing with long chains of empty iterables. Even with a high
1995 # number this would probably only fail in Py_DEBUG mode.
1996 it = chain.from_iterable(() for unused in range(10000000))
1997 with self.assertRaises(StopIteration):
1998 next(it)
1999
Thomas Woutersb2137042007-02-01 18:02:27 +00002000class SubclassWithKwargsTest(unittest.TestCase):
2001 def test_keywords_in_subclass(self):
2002 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002003 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00002004 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00002005 class Subclass(cls):
2006 def __init__(self, newarg=None, *args):
2007 cls.__init__(self, *args)
2008 try:
2009 Subclass(newarg=1)
2010 except TypeError as err:
2011 # we expect type errors because of wrong argument count
Serhiy Storchaka5eb788b2017-06-06 18:45:22 +03002012 self.assertNotIn("keyword arguments", err.args[0])
Thomas Woutersb2137042007-02-01 18:02:27 +00002013
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002014@support.cpython_only
2015class SizeofTest(unittest.TestCase):
2016 def setUp(self):
2017 self.ssize_t = struct.calcsize('n')
2018
2019 check_sizeof = support.check_sizeof
2020
2021 def test_product_sizeof(self):
2022 basesize = support.calcobjsize('3Pi')
2023 check = self.check_sizeof
2024 check(product('ab', '12'), basesize + 2 * self.ssize_t)
2025 check(product(*(('abc',) * 10)), basesize + 10 * self.ssize_t)
2026
2027 def test_combinations_sizeof(self):
2028 basesize = support.calcobjsize('3Pni')
2029 check = self.check_sizeof
2030 check(combinations('abcd', 3), basesize + 3 * self.ssize_t)
2031 check(combinations(range(10), 4), basesize + 4 * self.ssize_t)
2032
2033 def test_combinations_with_replacement_sizeof(self):
2034 cwr = combinations_with_replacement
2035 basesize = support.calcobjsize('3Pni')
2036 check = self.check_sizeof
2037 check(cwr('abcd', 3), basesize + 3 * self.ssize_t)
2038 check(cwr(range(10), 4), basesize + 4 * self.ssize_t)
2039
2040 def test_permutations_sizeof(self):
2041 basesize = support.calcobjsize('4Pni')
2042 check = self.check_sizeof
2043 check(permutations('abcd'),
2044 basesize + 4 * self.ssize_t + 4 * self.ssize_t)
2045 check(permutations('abcd', 3),
2046 basesize + 4 * self.ssize_t + 3 * self.ssize_t)
2047 check(permutations('abcde', 3),
2048 basesize + 5 * self.ssize_t + 3 * self.ssize_t)
2049 check(permutations(range(10), 4),
2050 basesize + 10 * self.ssize_t + 4 * self.ssize_t)
2051
Thomas Woutersb2137042007-02-01 18:02:27 +00002052
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00002053libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002054
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002055
2056>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002057>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00002058... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00002059...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002060Check 1200 is for $120.15
2061Check 1201 is for $764.05
2062Check 1202 is for $823.14
2063
2064>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002065>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00002066... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00002067...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000020681
20698
207027
2071
2072>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00002073>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00002074... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00002075...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002076Alex
2077Laura
2078Martin
2079Walter
2080Samuele
2081
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002082>>> from operator import itemgetter
2083>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00002084>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002085>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00002086... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00002087...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000020881 ['a', 'c', 'e']
20892 ['b', 'd', 'f']
20903 ['g']
2091
Raymond Hettinger734fb572004-01-20 20:04:40 +00002092# Find runs of consecutive numbers using groupby. The key to the solution
2093# is differencing with a range so that consecutive numbers all appear in
2094# same group.
2095>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00002096>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00002097... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00002098...
Raymond Hettinger734fb572004-01-20 20:04:40 +00002099[1]
2100[4, 5, 6]
2101[10]
2102[15, 16, 17, 18]
2103[22]
2104[25, 26, 27, 28]
2105
Georg Brandl3dbca812008-07-23 16:10:53 +00002106>>> def take(n, iterable):
2107... "Return first n items of the iterable as a list"
2108... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00002109
Georg Brandl3dbca812008-07-23 16:10:53 +00002110>>> def enumerate(iterable, start=0):
2111... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002112
Georg Brandl3dbca812008-07-23 16:10:53 +00002113>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002114... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00002115... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002116
Raymond Hettingercdf8ba32009-02-19 04:45:07 +00002117>>> def nth(iterable, n, default=None):
2118... "Returns the nth item or a default value"
2119... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002120
Raymond Hettingere525ee32016-03-06 18:11:38 -08002121>>> def all_equal(iterable):
2122... "Returns True if all the elements are equal to each other"
2123... g = groupby(iterable)
2124... return next(g, True) and not next(g, False)
2125
Georg Brandl3dbca812008-07-23 16:10:53 +00002126>>> def quantify(iterable, pred=bool):
2127... "Count how many times the predicate is true"
2128... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00002129
Georg Brandl3dbca812008-07-23 16:10:53 +00002130>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002131... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00002132... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002133
Georg Brandl3dbca812008-07-23 16:10:53 +00002134>>> def ncycles(iterable, n):
Ezio Melotti13925002011-03-16 11:05:33 +02002135... "Returns the sequence elements n times"
Georg Brandl3dbca812008-07-23 16:10:53 +00002136... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002137
2138>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002139... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002140
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002141>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002142... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002143
2144>>> def repeatfunc(func, times=None, *args):
2145... "Repeat calls to func with specified arguments."
2146... " Example: repeatfunc(random.random)"
2147... if times is None:
2148... return starmap(func, repeat(args))
2149... else:
2150... return starmap(func, repeat(args, times))
2151
Raymond Hettingerd591f662003-10-26 15:34:50 +00002152>>> def pairwise(iterable):
2153... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
2154... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00002155... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00002156... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00002157... except StopIteration:
2158... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002159... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002160
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002161>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00002162... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002163... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00002164... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002165
2166>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00002167... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002168... # Recipe credited to George Sakkis
2169... pending = len(iterables)
2170... nexts = cycle(iter(it).__next__ for it in iterables)
2171... while pending:
2172... try:
2173... for next in nexts:
2174... yield next()
2175... except StopIteration:
2176... pending -= 1
2177... nexts = cycle(islice(nexts, pending))
2178
2179>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00002180... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
2181... s = list(iterable)
2182... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002183
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00002184>>> def unique_everseen(iterable, key=None):
2185... "List unique elements, preserving order. Remember all elements ever seen."
2186... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
2187... # unique_everseen('ABBCcAD', str.lower) --> A B C D
2188... seen = set()
2189... seen_add = seen.add
2190... if key is None:
2191... for element in iterable:
2192... if element not in seen:
2193... seen_add(element)
2194... yield element
2195... else:
2196... for element in iterable:
2197... k = key(element)
2198... if k not in seen:
2199... seen_add(k)
2200... yield element
2201
2202>>> def unique_justseen(iterable, key=None):
2203... "List unique elements, preserving order. Remember only the element just seen."
2204... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
2205... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
2206... return map(next, map(itemgetter(1), groupby(iterable, key)))
2207
Raymond Hettinger31b26f62014-04-02 03:16:42 -07002208>>> def first_true(iterable, default=False, pred=None):
2209... '''Returns the first true value in the iterable.
2210...
2211... If no true value is found, returns *default*
2212...
2213... If *pred* is not None, returns the first item
2214... for which pred(item) is true.
2215...
2216... '''
2217... # first_true([a,b,c], x) --> a or b or c or x
2218... # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x
2219... return next(filter(pred, iterable), default)
2220
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002221This is not part of the examples but it tests to make sure the definitions
2222perform as purported.
2223
Raymond Hettingera098b332003-09-08 23:58:40 +00002224>>> take(10, count())
2225[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
2226
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002227>>> list(enumerate('abc'))
2228[(0, 'a'), (1, 'b'), (2, 'c')]
2229
2230>>> list(islice(tabulate(lambda x: 2*x), 4))
2231[0, 2, 4, 6]
2232
2233>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00002234'd'
2235
2236>>> nth('abcde', 9) is None
2237True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002238
Raymond Hettingere525ee32016-03-06 18:11:38 -08002239>>> [all_equal(s) for s in ('', 'A', 'AAAA', 'AAAB', 'AAABA')]
2240[True, True, True, False, False]
2241
Guido van Rossum805365e2007-05-07 22:24:25 +00002242>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000224350
2244
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002245>>> a = [[1, 2, 3], [4, 5, 6]]
2246>>> flatten(a)
2247[1, 2, 3, 4, 5, 6]
2248
2249>>> list(repeatfunc(pow, 5, 2, 3))
2250[8, 8, 8, 8, 8]
2251
2252>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002253>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002254[0, 0, 0, 0, 0]
2255
Raymond Hettingerd591f662003-10-26 15:34:50 +00002256>>> list(pairwise('abcd'))
2257[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002258
Raymond Hettingerd591f662003-10-26 15:34:50 +00002259>>> list(pairwise([]))
2260[]
2261
2262>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00002263[]
2264
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002265>>> list(islice(padnone('abc'), 0, 6))
2266['a', 'b', 'c', None, None, None]
2267
2268>>> list(ncycles('abc', 3))
2269['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
2270
2271>>> dotproduct([1,2,3], [4,5,6])
227232
2273
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002274>>> list(grouper(3, 'abcdefg', 'x'))
2275[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
2276
2277>>> list(roundrobin('abc', 'd', 'ef'))
2278['a', 'd', 'e', 'b', 'f', 'c']
2279
Raymond Hettingerace67332009-01-26 02:23:50 +00002280>>> list(powerset([1,2,3]))
2281[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002282
Raymond Hettinger191e8502009-01-27 13:29:43 +00002283>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
2284True
2285
2286>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
2287True
2288
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00002289>>> list(unique_everseen('AAAABBBCCDAABBB'))
2290['A', 'B', 'C', 'D']
2291
2292>>> list(unique_everseen('ABBCcAD', str.lower))
2293['A', 'B', 'C', 'D']
2294
2295>>> list(unique_justseen('AAAABBBCCDAABBB'))
2296['A', 'B', 'C', 'D', 'A', 'B']
2297
2298>>> list(unique_justseen('ABBCcAD', str.lower))
2299['A', 'B', 'C', 'A', 'D']
2300
Raymond Hettinger31b26f62014-04-02 03:16:42 -07002301>>> first_true('ABC0DEF1', '9', str.isdigit)
2302'0'
2303
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002304"""
2305
2306__test__ = {'libreftest' : libreftest}
2307
2308def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002309 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00002310 RegressionTests, LengthTransparency,
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002311 SubclassWithKwargsTest, TestExamples,
2312 SizeofTest)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002313 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002314
2315 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002316 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00002317 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00002318 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00002319 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002320 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00002321 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00002322 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00002323 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002324
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002325 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002326 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002327
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002328if __name__ == "__main__":
2329 test_main(verbose=True)