blob: fea68163b3a51beb1669dbe2d7acb84d7b38d111 [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 Hettinger2012f172003-02-07 05:32:58 +00007import sys
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00008import operator
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00009import random
Raymond Hettingera3e1ad22009-11-30 22:02:31 +000010import copy
11import pickle
Christian Heimesc3f30c42008-02-22 16:37:40 +000012from functools import reduce
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +020013import sys
14import struct
Benjamin Petersonee8712c2008-05-20 21:35:26 +000015maxsize = support.MAX_Py_ssize_t
Guido van Rossum360e4b82007-05-14 22:51:27 +000016minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000017
Guido van Rossum801f0d72006-08-24 19:48:10 +000018def lzip(*args):
19 return list(zip(*args))
20
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000021def onearg(x):
22 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000023 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000024
25def errfunc(*args):
26 'Test function that raises an error'
27 raise ValueError
28
29def gen3():
30 'Non-restartable source sequence'
31 for i in (0, 1, 2):
32 yield i
33
34def isEven(x):
35 'Test predicate'
36 return x%2==0
37
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000038def isOdd(x):
39 'Test predicate'
40 return x%2==1
41
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000042def tupleize(*args):
43 return args
44
45def irange(n):
46 for i in range(n):
47 yield i
48
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000049class StopNow:
50 'Class emulating an empty iterable.'
51 def __iter__(self):
52 return self
Georg Brandla18af4e2007-04-21 15:47:16 +000053 def __next__(self):
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000054 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000055
Raymond Hettinger02420702003-06-29 20:36:23 +000056def take(n, seq):
57 'Convenience function for partially consuming a long of infinite iterable'
58 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000059
Christian Heimes78644762008-03-04 23:39:23 +000060def prod(iterable):
61 return reduce(operator.mul, iterable, 1)
62
Christian Heimes380f7f22008-02-28 11:19:05 +000063def fact(n):
64 'Factorial'
Christian Heimes78644762008-03-04 23:39:23 +000065 return prod(range(1, n+1))
66
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000067# root level methods for pickling ability
68def testR(r):
69 return r[0]
70
71def testR2(r):
72 return r[2]
73
74def underten(x):
75 return x<10
76
Raymond Hettinger96ef8112003-02-01 00:10:11 +000077class TestBasicOps(unittest.TestCase):
Raymond Hettinger482ba772010-12-01 22:48:00 +000078
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000079 def pickletest(self, it, stop=4, take=1, compare=None):
80 """Test that an iterator is the same after pickling, also when part-consumed"""
81 def expand(it, i=0):
82 # Recursively expand iterables, within sensible bounds
83 if i > 10:
84 raise RuntimeError("infinite recursion encountered")
85 if isinstance(it, str):
86 return it
87 try:
88 l = list(islice(it, stop))
89 except TypeError:
90 return it # can't expand it
91 return [expand(e, i+1) for e in l]
92
93 # Test the initial copy against the original
94 dump = pickle.dumps(it)
95 i2 = pickle.loads(dump)
96 self.assertEqual(type(it), type(i2))
97 a, b = expand(it), expand(i2)
98 self.assertEqual(a, b)
99 if compare:
100 c = expand(compare)
101 self.assertEqual(a, c)
102
103 # Take from the copy, and create another copy and compare them.
104 i3 = pickle.loads(dump)
105 took = 0
106 try:
107 for i in range(take):
108 next(i3)
109 took += 1
110 except StopIteration:
111 pass #in case there is less data than 'take'
112 dump = pickle.dumps(i3)
113 i4 = pickle.loads(dump)
114 a, b = expand(i3), expand(i4)
115 self.assertEqual(a, b)
116 if compare:
117 c = expand(compare[took:])
118 self.assertEqual(a, c);
119
Raymond Hettinger482ba772010-12-01 22:48:00 +0000120 def test_accumulate(self):
121 self.assertEqual(list(accumulate(range(10))), # one positional arg
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000122 [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])
123 self.assertEqual(list(accumulate(iterable=range(10))), # kw arg
124 [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])
Raymond Hettinger482ba772010-12-01 22:48:00 +0000125 for typ in int, complex, Decimal, Fraction: # multiple types
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000126 self.assertEqual(
127 list(accumulate(map(typ, range(10)))),
Raymond Hettinger482ba772010-12-01 22:48:00 +0000128 list(map(typ, [0, 1, 3, 6, 10, 15, 21, 28, 36, 45])))
Raymond Hettinger2d93e6e2010-12-03 02:33:53 +0000129 self.assertEqual(list(accumulate('abc')), ['a', 'ab', 'abc']) # works with non-numeric
Raymond Hettinger482ba772010-12-01 22:48:00 +0000130 self.assertEqual(list(accumulate([])), []) # empty iterable
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000131 self.assertEqual(list(accumulate([7])), [7]) # iterable of length one
Raymond Hettinger5d446132011-03-27 18:52:10 -0700132 self.assertRaises(TypeError, accumulate, range(10), 5, 6) # too many args
Raymond Hettinger482ba772010-12-01 22:48:00 +0000133 self.assertRaises(TypeError, accumulate) # too few args
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000134 self.assertRaises(TypeError, accumulate, x=range(10)) # unexpected kwd arg
Raymond Hettinger482ba772010-12-01 22:48:00 +0000135 self.assertRaises(TypeError, list, accumulate([1, []])) # args that don't add
136
Raymond Hettinger5d446132011-03-27 18:52:10 -0700137 s = [2, 8, 9, 5, 7, 0, 3, 4, 1, 6]
138 self.assertEqual(list(accumulate(s, min)),
139 [2, 2, 2, 2, 2, 0, 0, 0, 0, 0])
140 self.assertEqual(list(accumulate(s, max)),
141 [2, 8, 9, 9, 9, 9, 9, 9, 9, 9])
142 self.assertEqual(list(accumulate(s, operator.mul)),
143 [2, 16, 144, 720, 5040, 0, 0, 0, 0, 0])
144 with self.assertRaises(TypeError):
145 list(accumulate(s, chr)) # unary-operation
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000146 self.pickletest(accumulate(range(10))) # test pickling
Raymond Hettinger5d446132011-03-27 18:52:10 -0700147
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000148 def test_chain(self):
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000149
150 def chain2(*iterables):
151 'Pure python version in the docs'
152 for it in iterables:
153 for element in it:
154 yield element
155
156 for c in (chain, chain2):
157 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
158 self.assertEqual(list(c('abc')), list('abc'))
159 self.assertEqual(list(c('')), [])
160 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
161 self.assertRaises(TypeError, list,c(2, 3))
Christian Heimesf16baeb2008-02-29 14:57:44 +0000162
163 def test_chain_from_iterable(self):
164 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
165 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
166 self.assertEqual(list(chain.from_iterable([''])), [])
167 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
168 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000169
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000170 def test_chain_reducible(self):
171 operators = [copy.deepcopy,
172 lambda s: pickle.loads(pickle.dumps(s))]
173 for oper in operators:
174 it = chain('abc', 'def')
175 self.assertEqual(list(oper(it)), list('abcdef'))
176 self.assertEqual(next(it), 'a')
177 self.assertEqual(list(oper(it)), list('bcdef'))
178
179 self.assertEqual(list(oper(chain(''))), [])
180 self.assertEqual(take(4, oper(chain('abc', 'def'))), list('abcd'))
181 self.assertRaises(TypeError, list, oper(chain(2, 3)))
182 self.pickletest(chain('abc', 'def'), compare=list('abcdef'))
183
Christian Heimes380f7f22008-02-28 11:19:05 +0000184 def test_combinations(self):
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000185 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Christian Heimes380f7f22008-02-28 11:19:05 +0000186 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Christian Heimes78644762008-03-04 23:39:23 +0000187 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Christian Heimes380f7f22008-02-28 11:19:05 +0000188 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000189
190 for op in (lambda a:a, lambda a:pickle.loads(pickle.dumps(a))):
191 self.assertEqual(list(op(combinations('abc', 32))), []) # r > n
192
193 self.assertEqual(list(op(combinations('ABCD', 2))),
194 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
195 testIntermediate = combinations('ABCD', 2)
196 next(testIntermediate)
197 self.assertEqual(list(op(testIntermediate)),
198 [('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
199
200 self.assertEqual(list(op(combinations(range(4), 3))),
201 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
202 testIntermediate = combinations(range(4), 3)
203 next(testIntermediate)
204 self.assertEqual(list(op(testIntermediate)),
205 [(0,1,3), (0,2,3), (1,2,3)])
206
Christian Heimes78644762008-03-04 23:39:23 +0000207
208 def combinations1(iterable, r):
209 'Pure python version shown in the docs'
210 pool = tuple(iterable)
211 n = len(pool)
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000212 if r > n:
213 return
Christian Heimes78644762008-03-04 23:39:23 +0000214 indices = list(range(r))
215 yield tuple(pool[i] for i in indices)
216 while 1:
217 for i in reversed(range(r)):
218 if indices[i] != i + n - r:
219 break
220 else:
221 return
222 indices[i] += 1
223 for j in range(i+1, r):
224 indices[j] = indices[j-1] + 1
225 yield tuple(pool[i] for i in indices)
226
227 def combinations2(iterable, r):
228 'Pure python version shown in the docs'
229 pool = tuple(iterable)
230 n = len(pool)
231 for indices in permutations(range(n), r):
232 if sorted(indices) == list(indices):
233 yield tuple(pool[i] for i in indices)
234
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000235 def combinations3(iterable, r):
236 'Pure python version from cwr()'
237 pool = tuple(iterable)
238 n = len(pool)
239 for indices in combinations_with_replacement(range(n), r):
240 if len(set(indices)) == r:
241 yield tuple(pool[i] for i in indices)
242
Christian Heimes78644762008-03-04 23:39:23 +0000243 for n in range(7):
Christian Heimes380f7f22008-02-28 11:19:05 +0000244 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000245 for r in range(n+2):
Christian Heimes380f7f22008-02-28 11:19:05 +0000246 result = list(combinations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000247 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 +0000248 self.assertEqual(len(result), len(set(result))) # no repeats
249 self.assertEqual(result, sorted(result)) # lexicographic order
250 for c in result:
251 self.assertEqual(len(c), r) # r-length combinations
252 self.assertEqual(len(set(c)), r) # no duplicate elements
253 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000254 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000255 self.assertEqual(list(c),
256 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000257 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000258 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000259 self.assertEqual(result, list(combinations3(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000260
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000261 self.pickletest(combinations(values, r)) # test pickling
262
263 # Test implementation detail: tuple re-use
Alex Gaynore151d212011-07-17 16:21:30 -0700264 @support.impl_detail("tuple reuse is specific to CPython")
265 def test_combinations_tuple_reuse(self):
Christian Heimes78644762008-03-04 23:39:23 +0000266 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
267 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
268
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000269 def test_combinations_with_replacement(self):
270 cwr = combinations_with_replacement
271 self.assertRaises(TypeError, cwr, 'abc') # missing r argument
272 self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
273 self.assertRaises(TypeError, cwr, None) # pool is not iterable
274 self.assertRaises(ValueError, cwr, 'abc', -2) # r is negative
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000275
276 for op in (lambda a:a, lambda a:pickle.loads(pickle.dumps(a))):
277 self.assertEqual(list(op(cwr('ABC', 2))),
278 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
279 testIntermediate = cwr('ABC', 2)
280 next(testIntermediate)
281 self.assertEqual(list(op(testIntermediate)),
282 [('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
283
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000284
285 def cwr1(iterable, r):
286 'Pure python version shown in the docs'
287 # number items returned: (n+r-1)! / r! / (n-1)! when n>0
288 pool = tuple(iterable)
289 n = len(pool)
290 if not n and r:
291 return
292 indices = [0] * r
293 yield tuple(pool[i] for i in indices)
294 while 1:
295 for i in reversed(range(r)):
296 if indices[i] != n - 1:
297 break
298 else:
299 return
300 indices[i:] = [indices[i] + 1] * (r - i)
301 yield tuple(pool[i] for i in indices)
302
303 def cwr2(iterable, r):
304 'Pure python version shown in the docs'
305 pool = tuple(iterable)
306 n = len(pool)
307 for indices in product(range(n), repeat=r):
308 if sorted(indices) == list(indices):
309 yield tuple(pool[i] for i in indices)
310
311 def numcombs(n, r):
312 if not n:
313 return 0 if r else 1
314 return fact(n+r-1) / fact(r)/ fact(n-1)
315
316 for n in range(7):
317 values = [5*x-12 for x in range(n)]
318 for r in range(n+2):
319 result = list(cwr(values, r))
320
321 self.assertEqual(len(result), numcombs(n, r)) # right number of combs
322 self.assertEqual(len(result), len(set(result))) # no repeats
323 self.assertEqual(result, sorted(result)) # lexicographic order
324
325 regular_combs = list(combinations(values, r)) # compare to combs without replacement
326 if n == 0 or r <= 1:
Ezio Melottib3aedd42010-11-20 19:04:17 +0000327 self.assertEqual(result, regular_combs) # cases that should be identical
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000328 else:
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000329 self.assertTrue(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000330
331 for c in result:
332 self.assertEqual(len(c), r) # r-length combinations
333 noruns = [k for k,v in groupby(c)] # combo without consecutive repeats
334 self.assertEqual(len(noruns), len(set(noruns))) # no repeats other than consecutive
335 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000336 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000337 self.assertEqual(noruns,
338 [e for e in values if e in c]) # comb is a subsequence of the input iterable
339 self.assertEqual(result, list(cwr1(values, r))) # matches first pure python version
340 self.assertEqual(result, list(cwr2(values, r))) # matches second pure python version
341
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000342 self.pickletest(cwr(values,r)) # test pickling
343
344 # Test implementation detail: tuple re-use
345
Alex Gaynore151d212011-07-17 16:21:30 -0700346 @support.impl_detail("tuple reuse is specific to CPython")
347 def test_combinations_with_replacement_tuple_reuse(self):
348 cwr = combinations_with_replacement
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000349 self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
350 self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
351
Christian Heimes78644762008-03-04 23:39:23 +0000352 def test_permutations(self):
353 self.assertRaises(TypeError, permutations) # too few arguments
354 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000355 self.assertRaises(TypeError, permutations, None) # pool is not iterable
356 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000357 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000358 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Christian Heimes78644762008-03-04 23:39:23 +0000359 self.assertEqual(list(permutations(range(3), 2)),
360 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
361
362 def permutations1(iterable, r=None):
363 'Pure python version shown in the docs'
364 pool = tuple(iterable)
365 n = len(pool)
366 r = n if r is None else r
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000367 if r > n:
368 return
Christian Heimes78644762008-03-04 23:39:23 +0000369 indices = list(range(n))
370 cycles = list(range(n-r+1, n+1))[::-1]
371 yield tuple(pool[i] for i in indices[:r])
372 while n:
373 for i in reversed(range(r)):
374 cycles[i] -= 1
375 if cycles[i] == 0:
376 indices[i:] = indices[i+1:] + indices[i:i+1]
377 cycles[i] = n - i
378 else:
379 j = cycles[i]
380 indices[i], indices[-j] = indices[-j], indices[i]
381 yield tuple(pool[i] for i in indices[:r])
382 break
383 else:
384 return
385
386 def permutations2(iterable, r=None):
387 'Pure python version shown in the docs'
388 pool = tuple(iterable)
389 n = len(pool)
390 r = n if r is None else r
391 for indices in product(range(n), repeat=r):
392 if len(set(indices)) == r:
393 yield tuple(pool[i] for i in indices)
394
395 for n in range(7):
396 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000397 for r in range(n+2):
Christian Heimes78644762008-03-04 23:39:23 +0000398 result = list(permutations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000399 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 +0000400 self.assertEqual(len(result), len(set(result))) # no repeats
401 self.assertEqual(result, sorted(result)) # lexicographic order
402 for p in result:
403 self.assertEqual(len(p), r) # r-length permutations
404 self.assertEqual(len(set(p)), r) # no duplicate elements
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000405 self.assertTrue(all(e in values for e in p)) # elements taken from input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000406 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000407 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000408 if r == n:
409 self.assertEqual(result, list(permutations(values, None))) # test r as None
410 self.assertEqual(result, list(permutations(values))) # test default r
411
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000412 self.pickletest(permutations(values, r)) # test pickling
413
Zachary Waredca807b2014-04-24 13:22:05 -0500414 @support.impl_detail("tuple reuse is specific to CPython")
Alex Gaynore151d212011-07-17 16:21:30 -0700415 def test_permutations_tuple_reuse(self):
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000416 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Christian Heimes78644762008-03-04 23:39:23 +0000417 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Christian Heimes380f7f22008-02-28 11:19:05 +0000418
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000419 def test_combinatorics(self):
420 # Test relationships between product(), permutations(),
421 # combinations() and combinations_with_replacement().
422
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000423 for n in range(6):
424 s = 'ABCDEFG'[:n]
425 for r in range(8):
426 prod = list(product(s, repeat=r))
427 cwr = list(combinations_with_replacement(s, r))
428 perm = list(permutations(s, r))
429 comb = list(combinations(s, r))
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000430
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000431 # Check size
Ezio Melottib3aedd42010-11-20 19:04:17 +0000432 self.assertEqual(len(prod), n**r)
433 self.assertEqual(len(cwr), (fact(n+r-1) / fact(r)/ fact(n-1)) if n else (not r))
434 self.assertEqual(len(perm), 0 if r>n else fact(n) / fact(n-r))
435 self.assertEqual(len(comb), 0 if r>n else fact(n) / fact(r) / fact(n-r))
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000436
437 # Check lexicographic order without repeated tuples
Ezio Melottib3aedd42010-11-20 19:04:17 +0000438 self.assertEqual(prod, sorted(set(prod)))
439 self.assertEqual(cwr, sorted(set(cwr)))
440 self.assertEqual(perm, sorted(set(perm)))
441 self.assertEqual(comb, sorted(set(comb)))
Raymond Hettingerda6bc522009-01-27 10:39:42 +0000442
443 # Check interrelationships
Ezio Melottib3aedd42010-11-20 19:04:17 +0000444 self.assertEqual(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
445 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 +0000446 self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
447 self.assertEqual(comb, [t for t in cwr if len(set(t))==r]) # comb: cwrs without dups
448 self.assertEqual(comb, list(filter(set(cwr).__contains__, perm))) # comb: perm that is a cwr
449 self.assertEqual(comb, list(filter(set(perm).__contains__, cwr))) # comb: cwr that is a perm
450 self.assertEqual(comb, sorted(set(cwr) & set(perm))) # comb: both a cwr and a perm
Raymond Hettingereb508ad2009-01-27 09:35:21 +0000451
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000452 def test_compress(self):
Raymond Hettinger15a49502009-02-19 02:17:09 +0000453 self.assertEqual(list(compress(data='ABCDEF', selectors=[1,0,1,0,1,1])), list('ACEF'))
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000454 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
455 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
456 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
457 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
458 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
459 n = 10000
460 data = chain.from_iterable(repeat(range(6), n))
461 selectors = chain.from_iterable(repeat((0, 1)))
462 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
463 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
464 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
465 self.assertRaises(TypeError, compress, range(6)) # too few args
466 self.assertRaises(TypeError, compress, range(6), None) # too many args
467
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000468 # check copy, deepcopy, pickle
469 for op in (lambda a:copy.copy(a), lambda a:copy.deepcopy(a), lambda a:pickle.loads(pickle.dumps(a))):
470 for data, selectors, result1, result2 in [
471 ('ABCDEF', [1,0,1,0,1,1], 'ACEF', 'CEF'),
472 ('ABCDEF', [0,0,0,0,0,0], '', ''),
473 ('ABCDEF', [1,1,1,1,1,1], 'ABCDEF', 'BCDEF'),
474 ('ABCDEF', [1,0,1], 'AC', 'C'),
475 ('ABC', [0,1,1,1,1,1], 'BC', 'C'),
476 ]:
477
478 self.assertEqual(list(op(compress(data=data, selectors=selectors))), list(result1))
479 self.assertEqual(list(op(compress(data, selectors))), list(result1))
480 testIntermediate = compress(data, selectors)
481 if result1:
482 next(testIntermediate)
483 self.assertEqual(list(op(testIntermediate)), list(result2))
484
485
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000486 def test_count(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000487 self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
488 self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
489 self.assertEqual(take(2, lzip('abc',count(3))), [('a', 3), ('b', 4)])
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000490 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
491 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000492 self.assertRaises(TypeError, count, 2, 3, 4)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000493 self.assertRaises(TypeError, count, 'a')
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000494 self.assertEqual(list(islice(count(maxsize-5), 10)),
495 list(range(maxsize-5, maxsize+5)))
496 self.assertEqual(list(islice(count(-maxsize-5), 10)),
497 list(range(-maxsize-5, -maxsize+5)))
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +0000498 self.assertEqual(list(islice(count(10, maxsize+5), 3)),
499 list(range(10, 10+3*(maxsize+5), maxsize+5)))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000500 c = count(3)
501 self.assertEqual(repr(c), 'count(3)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000502 next(c)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000503 self.assertEqual(repr(c), 'count(4)')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000504 c = count(-9)
505 self.assertEqual(repr(c), 'count(-9)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000506 next(c)
Raymond Hettinger30729212009-02-12 06:28:27 +0000507 self.assertEqual(repr(count(10.25)), 'count(10.25)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000508 self.assertEqual(next(c), -8)
Christian Heimesa37d4c62007-12-04 23:02:19 +0000509 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 +0300510 # Test repr
511 r1 = repr(count(i))
512 r2 = 'count(%r)'.__mod__(i)
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000513 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000514
Raymond Hettingera3e1ad22009-11-30 22:02:31 +0000515 # check copy, deepcopy, pickle
516 for value in -3, 3, maxsize-5, maxsize+5:
517 c = count(value)
518 self.assertEqual(next(copy.copy(c)), value)
519 self.assertEqual(next(copy.deepcopy(c)), value)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000520 self.pickletest(count(value))
Raymond Hettingera3e1ad22009-11-30 22:02:31 +0000521
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +0000522 #check proper internal error handling for large "step' sizes
523 count(1, maxsize+5); sys.exc_info()
524
Raymond Hettinger30729212009-02-12 06:28:27 +0000525 def test_count_with_stride(self):
526 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingereb13fdd2009-02-21 22:30:12 +0000527 self.assertEqual(lzip('abc',count(start=2,step=3)),
528 [('a', 2), ('b', 5), ('c', 8)])
529 self.assertEqual(lzip('abc',count(step=-1)),
530 [('a', 0), ('b', -1), ('c', -2)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000531 self.assertEqual(lzip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
532 self.assertEqual(lzip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000533 self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000534 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
535 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
536 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettinger8a882a72009-02-12 12:08:18 +0000537 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
538 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerde09acf2009-02-12 12:53:53 +0000539 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
540 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Raymond Hettinger30729212009-02-12 06:28:27 +0000541 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
542 c = count(3, 5)
543 self.assertEqual(repr(c), 'count(3, 5)')
544 next(c)
545 self.assertEqual(repr(c), 'count(8, 5)')
546 c = count(-9, 0)
547 self.assertEqual(repr(c), 'count(-9, 0)')
548 next(c)
549 self.assertEqual(repr(c), 'count(-9, 0)')
550 c = count(-9, -3)
551 self.assertEqual(repr(c), 'count(-9, -3)')
552 next(c)
553 self.assertEqual(repr(c), 'count(-12, -3)')
554 self.assertEqual(repr(c), 'count(-12, -3)')
555 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
556 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
557 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
558 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
559 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 +0300560 # Test repr
561 r1 = repr(count(i, j))
Raymond Hettinger30729212009-02-12 06:28:27 +0000562 if j == 1:
Serhiy Storchaka95949422013-08-27 19:40:23 +0300563 r2 = ('count(%r)' % i)
Raymond Hettinger30729212009-02-12 06:28:27 +0000564 else:
Serhiy Storchaka95949422013-08-27 19:40:23 +0300565 r2 = ('count(%r, %r)' % (i, j))
Raymond Hettinger30729212009-02-12 06:28:27 +0000566 self.assertEqual(r1, r2)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000567 self.pickletest(count(i, j))
Raymond Hettinger30729212009-02-12 06:28:27 +0000568
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000569 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000570 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000571 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000572 self.assertRaises(TypeError, cycle)
573 self.assertRaises(TypeError, cycle, 5)
574 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000575
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000576 # check copy, deepcopy, pickle
577 c = cycle('abc')
578 self.assertEqual(next(c), 'a')
579 #simple copy currently not supported, because __reduce__ returns
580 #an internal iterator
581 #self.assertEqual(take(10, copy.copy(c)), list('bcabcabcab'))
582 self.assertEqual(take(10, copy.deepcopy(c)), list('bcabcabcab'))
583 self.assertEqual(take(10, pickle.loads(pickle.dumps(c))), list('bcabcabcab'))
584 next(c)
585 self.assertEqual(take(10, pickle.loads(pickle.dumps(c))), list('cabcabcabc'))
586 self.pickletest(cycle('abc'))
587
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000588 def test_groupby(self):
589 # Check whether it accepts arguments correctly
590 self.assertEqual([], list(groupby([])))
591 self.assertEqual([], list(groupby([], key=id)))
592 self.assertRaises(TypeError, list, groupby('abc', []))
593 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000594 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000595
596 # Check normal input
597 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
598 (2,15,22), (3,16,23), (3,17,23)]
599 dup = []
600 for k, g in groupby(s, lambda r:r[0]):
601 for elem in g:
602 self.assertEqual(k, elem[0])
603 dup.append(elem)
604 self.assertEqual(s, dup)
605
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000606 # Check normal pickled
607 dup = []
608 for k, g in pickle.loads(pickle.dumps(groupby(s, testR))):
609 for elem in g:
610 self.assertEqual(k, elem[0])
611 dup.append(elem)
612 self.assertEqual(s, dup)
613
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000614 # Check nested case
615 dup = []
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000616 for k, g in groupby(s, testR):
617 for ik, ig in groupby(g, testR2):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000618 for elem in ig:
619 self.assertEqual(k, elem[0])
620 self.assertEqual(ik, elem[2])
621 dup.append(elem)
622 self.assertEqual(s, dup)
623
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000624 # Check nested and pickled
625 dup = []
626 for k, g in pickle.loads(pickle.dumps(groupby(s, testR))):
627 for ik, ig in pickle.loads(pickle.dumps(groupby(g, testR2))):
628 for elem in ig:
629 self.assertEqual(k, elem[0])
630 self.assertEqual(ik, elem[2])
631 dup.append(elem)
632 self.assertEqual(s, dup)
633
634
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000635 # Check case where inner iterator is not used
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000636 keys = [k for k, g in groupby(s, testR)]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000637 expectedkeys = set([r[0] for r in s])
638 self.assertEqual(set(keys), expectedkeys)
639 self.assertEqual(len(keys), len(expectedkeys))
640
641 # Exercise pipes and filters style
642 s = 'abracadabra'
643 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000644 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000645 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
646 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000647 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000648 self.assertEqual(r, ['a', 'b', 'r'])
649 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000650 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000651 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
652 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000653 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000654 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
655
Georg Brandla18af4e2007-04-21 15:47:16 +0000656 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000657 class ExpectedError(Exception):
658 pass
659 def delayed_raise(n=0):
660 for i in range(n):
661 yield 'yo'
662 raise ExpectedError
663 def gulp(iterable, keyp=None, func=list):
664 return [func(g) for k, g in groupby(iterable, keyp)]
665
Georg Brandla18af4e2007-04-21 15:47:16 +0000666 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000667 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000668 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000669 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
670
671 # __cmp__ failure
672 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000673 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000674 raise ExpectedError
675 s = [DummyCmp(), DummyCmp(), None]
676
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000677 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000678 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000679 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000680 self.assertRaises(ExpectedError, gulp, s)
681
682 # keyfunc failure
683 def keyfunc(obj):
684 if keyfunc.skip > 0:
685 keyfunc.skip -= 1
686 return obj
687 else:
688 raise ExpectedError
689
690 # keyfunc failure on outer object
691 keyfunc.skip = 0
692 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
693 keyfunc.skip = 1
694 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
695
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000696 def test_filter(self):
697 self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
698 self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])
699 self.assertEqual(list(filter(bool, [0,1,0,2,0])), [1,2])
700 self.assertEqual(take(4, filter(isEven, count())), [0,2,4,6])
701 self.assertRaises(TypeError, filter)
702 self.assertRaises(TypeError, filter, lambda x:x)
703 self.assertRaises(TypeError, filter, lambda x:x, range(6), 7)
704 self.assertRaises(TypeError, filter, isEven, 3)
705 self.assertRaises(TypeError, next, filter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000706
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000707 # check copy, deepcopy, pickle
708 ans = [0,2,4]
709
710 c = filter(isEven, range(6))
711 self.assertEqual(list(copy.copy(c)), ans)
712 c = filter(isEven, range(6))
713 self.assertEqual(list(copy.deepcopy(c)), ans)
714 c = filter(isEven, range(6))
715 self.assertEqual(list(pickle.loads(pickle.dumps(c))), ans)
716 next(c)
717 self.assertEqual(list(pickle.loads(pickle.dumps(c))), ans[1:])
718 c = filter(isEven, range(6))
719 self.pickletest(c)
720
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000721 def test_filterfalse(self):
722 self.assertEqual(list(filterfalse(isEven, range(6))), [1,3,5])
723 self.assertEqual(list(filterfalse(None, [0,1,0,2,0])), [0,0,0])
724 self.assertEqual(list(filterfalse(bool, [0,1,0,2,0])), [0,0,0])
725 self.assertEqual(take(4, filterfalse(isEven, count())), [1,3,5,7])
726 self.assertRaises(TypeError, filterfalse)
727 self.assertRaises(TypeError, filterfalse, lambda x:x)
728 self.assertRaises(TypeError, filterfalse, lambda x:x, range(6), 7)
729 self.assertRaises(TypeError, filterfalse, isEven, 3)
730 self.assertRaises(TypeError, next, filterfalse(range(6), range(6)))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000731 self.pickletest(filterfalse(isEven, range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000732
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000733 def test_zip(self):
734 # XXX This is rather silly now that builtin zip() calls zip()...
735 ans = [(x,y) for x, y in zip('abc',count())]
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000736 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000737 self.assertEqual(list(zip('abc', range(6))), lzip('abc', range(6)))
738 self.assertEqual(list(zip('abcdef', range(3))), lzip('abcdef', range(3)))
739 self.assertEqual(take(3,zip('abcdef', count())), lzip('abcdef', range(3)))
740 self.assertEqual(list(zip('abcdef')), lzip('abcdef'))
741 self.assertEqual(list(zip()), lzip())
742 self.assertRaises(TypeError, zip, 3)
743 self.assertRaises(TypeError, zip, range(3), 3)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000744 self.assertEqual([tuple(list(pair)) for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000745 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000746 self.assertEqual([pair for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000747 lzip('abc', 'def'))
Alex Gaynore151d212011-07-17 16:21:30 -0700748
749 @support.impl_detail("tuple reuse is specific to CPython")
750 def test_zip_tuple_reuse(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000751 ids = list(map(id, zip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000752 self.assertEqual(min(ids), max(ids))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000753 ids = list(map(id, list(zip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000754 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000755
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000756 # check copy, deepcopy, pickle
757 ans = [(x,y) for x, y in copy.copy(zip('abc',count()))]
758 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
759
760 ans = [(x,y) for x, y in copy.deepcopy(zip('abc',count()))]
761 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
762
763 ans = [(x,y) for x, y in pickle.loads(pickle.dumps(zip('abc',count())))]
764 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
765
766 testIntermediate = zip('abc',count())
767 next(testIntermediate)
768 ans = [(x,y) for x, y in pickle.loads(pickle.dumps(testIntermediate))]
769 self.assertEqual(ans, [('b', 1), ('c', 2)])
770
771 self.pickletest(zip('abc', count()))
772
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000773 def test_ziplongest(self):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000774 for args in [
775 ['abc', range(6)],
776 [range(6), 'abc'],
777 [range(1000), range(2000,2100), range(3000,3050)],
778 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
779 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
780 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000781 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
782 for i in range(max(map(len, args)))]
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000783 self.assertEqual(list(zip_longest(*args)), target)
784 self.assertEqual(list(zip_longest(*args, **{})), target)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000785 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 +0000786 self.assertEqual(list(zip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000787
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000788 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 +0000789
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000790 self.assertEqual(list(zip_longest()), list(zip()))
791 self.assertEqual(list(zip_longest([])), list(zip([])))
792 self.assertEqual(list(zip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000793
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000794 self.assertEqual(list(zip_longest('abc', 'defg', **{})),
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000795 list(zip(list('abc')+[None], 'defg'))) # empty keyword dict
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000796 self.assertRaises(TypeError, zip_longest, 3)
797 self.assertRaises(TypeError, zip_longest, range(3), 3)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000798
799 for stmt in [
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000800 "zip_longest('abc', fv=1)",
801 "zip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000802 ]:
803 try:
804 eval(stmt, globals(), locals())
805 except TypeError:
806 pass
807 else:
808 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000809
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000810 self.assertEqual([tuple(list(pair)) for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000811 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000812 self.assertEqual([pair for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000813 list(zip('abc', 'def')))
Alex Gaynore151d212011-07-17 16:21:30 -0700814
815 @support.impl_detail("tuple reuse is specific to CPython")
816 def test_zip_longest_tuple_reuse(self):
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000817 ids = list(map(id, zip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000818 self.assertEqual(min(ids), max(ids))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000819 ids = list(map(id, list(zip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000820 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
821
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000822 def test_zip_longest_pickling(self):
823 self.pickletest(zip_longest("abc", "def"))
824 self.pickletest(zip_longest("abc", "defgh"))
825 self.pickletest(zip_longest("abc", "defgh", fillvalue=1))
826 self.pickletest(zip_longest("", "defgh"))
827
Raymond Hettingerfc438512009-11-01 20:55:33 +0000828 def test_bug_7244(self):
829
830 class Repeater:
831 # this class is similar to itertools.repeat
832 def __init__(self, o, t, e):
833 self.o = o
834 self.t = int(t)
835 self.e = e
836 def __iter__(self): # its iterator is itself
837 return self
838 def __next__(self):
839 if self.t > 0:
840 self.t -= 1
841 return self.o
842 else:
843 raise self.e
844
845 # Formerly this code in would fail in debug mode
846 # with Undetected Error and Stop Iteration
847 r1 = Repeater(1, 3, StopIteration)
848 r2 = Repeater(2, 4, StopIteration)
849 def run(r1, r2):
850 result = []
851 for i, j in zip_longest(r1, r2, fillvalue=0):
852 with support.captured_output('stdout'):
853 print((i, j))
854 result.append((i, j))
855 return result
856 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
857
858 # Formerly, the RuntimeError would be lost
859 # and StopIteration would stop as expected
860 r1 = Repeater(1, 3, RuntimeError)
861 r2 = Repeater(2, 4, StopIteration)
862 it = zip_longest(r1, r2, fillvalue=0)
863 self.assertEqual(next(it), (1, 2))
864 self.assertEqual(next(it), (1, 2))
865 self.assertEqual(next(it), (1, 2))
866 self.assertRaises(RuntimeError, next, it)
867
Christian Heimesc3f30c42008-02-22 16:37:40 +0000868 def test_product(self):
869 for args, result in [
Christian Heimes78644762008-03-04 23:39:23 +0000870 ([], [()]), # zero iterables
Christian Heimesc3f30c42008-02-22 16:37:40 +0000871 (['ab'], [('a',), ('b',)]), # one iterable
872 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
873 ([range(0), range(2), range(3)], []), # first iterable with zero length
874 ([range(2), range(0), range(3)], []), # middle iterable with zero length
875 ([range(2), range(3), range(0)], []), # last iterable with zero length
876 ]:
877 self.assertEqual(list(product(*args)), result)
Christian Heimesf16baeb2008-02-29 14:57:44 +0000878 for r in range(4):
879 self.assertEqual(list(product(*(args*r))),
880 list(product(*args, **dict(repeat=r))))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000881 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
882 self.assertRaises(TypeError, product, range(6), None)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000883
884 def product1(*args, **kwds):
885 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
886 n = len(pools)
887 if n == 0:
888 yield ()
889 return
890 if any(len(pool) == 0 for pool in pools):
891 return
892 indices = [0] * n
893 yield tuple(pool[i] for pool, i in zip(pools, indices))
894 while 1:
895 for i in reversed(range(n)): # right to left
896 if indices[i] == len(pools[i]) - 1:
897 continue
898 indices[i] += 1
899 for j in range(i+1, n):
900 indices[j] = 0
901 yield tuple(pool[i] for pool, i in zip(pools, indices))
902 break
903 else:
904 return
905
906 def product2(*args, **kwds):
907 'Pure python version used in docs'
908 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
909 result = [[]]
910 for pool in pools:
911 result = [x+[y] for x in result for y in pool]
912 for prod in result:
913 yield tuple(prod)
914
Christian Heimesc3f30c42008-02-22 16:37:40 +0000915 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
916 set('abcdefg'), range(11), tuple(range(13))]
917 for i in range(100):
918 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Christian Heimes78644762008-03-04 23:39:23 +0000919 expected_len = prod(map(len, args))
920 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000921 self.assertEqual(list(product(*args)), list(product1(*args)))
922 self.assertEqual(list(product(*args)), list(product2(*args)))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000923 args = map(iter, args)
Christian Heimes78644762008-03-04 23:39:23 +0000924 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000925
Alex Gaynore151d212011-07-17 16:21:30 -0700926 @support.impl_detail("tuple reuse is specific to CPython")
927 def test_product_tuple_reuse(self):
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000928 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
929 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000930
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000931 def test_product_pickling(self):
932 # check copy, deepcopy, pickle
933 for args, result in [
934 ([], [()]), # zero iterables
935 (['ab'], [('a',), ('b',)]), # one iterable
936 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
937 ([range(0), range(2), range(3)], []), # first iterable with zero length
938 ([range(2), range(0), range(3)], []), # middle iterable with zero length
939 ([range(2), range(3), range(0)], []), # last iterable with zero length
940 ]:
941 self.assertEqual(list(copy.copy(product(*args))), result)
942 self.assertEqual(list(copy.deepcopy(product(*args))), result)
943 self.pickletest(product(*args))
944
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000945 def test_repeat(self):
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +0000946 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
Guido van Rossum805365e2007-05-07 22:24:25 +0000947 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000948 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000949 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000950 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000951 self.assertEqual(list(repeat('a', 0)), [])
952 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000953 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000954 self.assertRaises(TypeError, repeat, None, 3, 4)
955 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000956 r = repeat(1+0j)
957 self.assertEqual(repr(r), 'repeat((1+0j))')
958 r = repeat(1+0j, 5)
959 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
960 list(r)
961 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000962
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000963 # check copy, deepcopy, pickle
964 c = repeat(object='a', times=10)
965 self.assertEqual(next(c), 'a')
966 self.assertEqual(take(2, copy.copy(c)), list('a' * 2))
967 self.assertEqual(take(2, copy.deepcopy(c)), list('a' * 2))
968 self.pickletest(repeat(object='a', times=10))
969
Raymond Hettinger97d35552014-06-24 21:36:58 -0700970 def test_repeat_with_negative_times(self):
971 self.assertEqual(repr(repeat('a', -1)), "repeat('a', 0)")
972 self.assertEqual(repr(repeat('a', -2)), "repeat('a', 0)")
973 self.assertEqual(repr(repeat('a', times=-1)), "repeat('a', 0)")
974 self.assertEqual(repr(repeat('a', times=-2)), "repeat('a', 0)")
975
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000976 def test_map(self):
977 self.assertEqual(list(map(operator.pow, range(3), range(1,7))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000978 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000979 self.assertEqual(list(map(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000980 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000981 self.assertEqual(list(map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000982 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000983 self.assertEqual(take(2,map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000984 [('a',0),('b',1)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000985 self.assertEqual(list(map(operator.pow, [])), [])
986 self.assertRaises(TypeError, map)
987 self.assertRaises(TypeError, list, map(None, range(3), range(3)))
988 self.assertRaises(TypeError, map, operator.neg)
989 self.assertRaises(TypeError, next, map(10, range(5)))
990 self.assertRaises(ValueError, next, map(errfunc, [4], [5]))
991 self.assertRaises(TypeError, next, map(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000992
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000993 # check copy, deepcopy, pickle
994 ans = [('a',0),('b',1),('c',2)]
995
996 c = map(tupleize, 'abc', count())
997 self.assertEqual(list(copy.copy(c)), ans)
998
999 c = map(tupleize, 'abc', count())
1000 self.assertEqual(list(copy.deepcopy(c)), ans)
1001
1002 c = map(tupleize, 'abc', count())
1003 self.pickletest(c)
1004
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001005 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001006 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
1007 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001008 self.assertEqual(take(3, starmap(operator.pow, zip(count(), count(1)))),
Raymond Hettinger02420702003-06-29 20:36:23 +00001009 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001010 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +00001011 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
1012 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001013 self.assertRaises(TypeError, starmap)
1014 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001015 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
1016 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
1017 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001018
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001019 # check copy, deepcopy, pickle
1020 ans = [0**1, 1**2, 2**3]
1021
1022 c = starmap(operator.pow, zip(range(3), range(1,7)))
1023 self.assertEqual(list(copy.copy(c)), ans)
1024
1025 c = starmap(operator.pow, zip(range(3), range(1,7)))
1026 self.assertEqual(list(copy.deepcopy(c)), ans)
1027
1028 c = starmap(operator.pow, zip(range(3), range(1,7)))
1029 self.pickletest(c)
1030
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001031 def test_islice(self):
1032 for args in [ # islice(args) should agree with range(args)
1033 (10, 20, 3),
1034 (10, 3, 20),
1035 (10, 20),
1036 (10, 3),
1037 (20,)
1038 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +00001039 self.assertEqual(list(islice(range(100), *args)),
1040 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001041
1042 for args, tgtargs in [ # Stop when seqn is exhausted
1043 ((10, 110, 3), ((10, 100, 3))),
1044 ((10, 110), ((10, 100))),
1045 ((110,), (100,))
1046 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +00001047 self.assertEqual(list(islice(range(100), *args)),
1048 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001049
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001050 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +00001051 self.assertEqual(list(islice(range(10), None)), list(range(10)))
1052 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
1053 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
1054 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
1055 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001056
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +00001057 # Test number of items consumed SF #1171417
1058 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +00001059 self.assertEqual(list(islice(it, 3)), list(range(3)))
1060 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +00001061
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001062 # Test invalid arguments
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001063 ra = range(10)
1064 self.assertRaises(TypeError, islice, ra)
1065 self.assertRaises(TypeError, islice, ra, 1, 2, 3, 4)
1066 self.assertRaises(ValueError, islice, ra, -5, 10, 1)
1067 self.assertRaises(ValueError, islice, ra, 1, -5, -1)
1068 self.assertRaises(ValueError, islice, ra, 1, 10, -1)
1069 self.assertRaises(ValueError, islice, ra, 1, 10, 0)
1070 self.assertRaises(ValueError, islice, ra, 'a')
1071 self.assertRaises(ValueError, islice, ra, 'a', 1)
1072 self.assertRaises(ValueError, islice, ra, 1, 'a')
1073 self.assertRaises(ValueError, islice, ra, 'a', 1, 1)
1074 self.assertRaises(ValueError, islice, ra, 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +00001075 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001076
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001077 # Issue #10323: Less islice in a predictable state
1078 c = count()
1079 self.assertEqual(list(islice(c, 1, 3, 50)), [1])
1080 self.assertEqual(next(c), 3)
1081
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001082 # check copy, deepcopy, pickle
1083 for args in [ # islice(args) should agree with range(args)
1084 (10, 20, 3),
1085 (10, 3, 20),
1086 (10, 20),
1087 (10, 3),
1088 (20,)
1089 ]:
1090 self.assertEqual(list(copy.copy(islice(range(100), *args))),
1091 list(range(*args)))
1092 self.assertEqual(list(copy.deepcopy(islice(range(100), *args))),
1093 list(range(*args)))
1094 self.pickletest(islice(range(100), *args))
1095
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001096 # Issue #21321: check source iterator is not referenced
1097 # from islice() after the latter has been exhausted
1098 it = (x for x in (1, 2))
1099 wr = weakref.ref(it)
1100 it = islice(it, 1)
1101 self.assertIsNotNone(wr())
1102 list(it) # exhaust the iterator
1103 self.assertIsNone(wr())
1104
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001105 def test_takewhile(self):
1106 data = [1, 3, 5, 20, 2, 4, 6, 8]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001107 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001108 self.assertEqual(list(takewhile(underten, [])), [])
1109 self.assertRaises(TypeError, takewhile)
1110 self.assertRaises(TypeError, takewhile, operator.pow)
1111 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001112 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
1113 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001114 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
1115 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +00001116 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001117
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001118 # check copy, deepcopy, pickle
1119 self.assertEqual(list(copy.copy(takewhile(underten, data))), [1, 3, 5])
1120 self.assertEqual(list(copy.deepcopy(takewhile(underten, data))),
1121 [1, 3, 5])
1122 self.pickletest(takewhile(underten, data))
1123
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001124 def test_dropwhile(self):
1125 data = [1, 3, 5, 20, 2, 4, 6, 8]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001126 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001127 self.assertEqual(list(dropwhile(underten, [])), [])
1128 self.assertRaises(TypeError, dropwhile)
1129 self.assertRaises(TypeError, dropwhile, operator.pow)
1130 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +00001131 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
1132 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001133
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001134 # check copy, deepcopy, pickle
1135 self.assertEqual(list(copy.copy(dropwhile(underten, data))), [20, 2, 4, 6, 8])
1136 self.assertEqual(list(copy.deepcopy(dropwhile(underten, data))),
1137 [20, 2, 4, 6, 8])
1138 self.pickletest(dropwhile(underten, data))
1139
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001140 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +00001141 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001142
1143 a, b = tee([]) # test empty iterator
1144 self.assertEqual(list(a), [])
1145 self.assertEqual(list(b), [])
1146
1147 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +00001148 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001149
1150 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +00001151 self.assertEqual(list(a), list(range(n)))
1152 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001153
1154 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +00001155 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001156 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001157 del a
Guido van Rossum805365e2007-05-07 22:24:25 +00001158 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001159
1160 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +00001161 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001162 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001163 del b
Guido van Rossum805365e2007-05-07 22:24:25 +00001164 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001165
Guido van Rossum805365e2007-05-07 22:24:25 +00001166 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001167 order = [0]*n + [1]*n
1168 random.shuffle(order)
1169 lists = ([], [])
1170 its = tee(irange(n))
1171 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +00001172 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001173 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +00001174 self.assertEqual(lists[0], list(range(n)))
1175 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001176
Raymond Hettingerad983e72003-11-12 14:32:26 +00001177 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001178 self.assertRaises(TypeError, tee)
1179 self.assertRaises(TypeError, tee, 3)
1180 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +00001181 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001182
Raymond Hettingerad983e72003-11-12 14:32:26 +00001183 # tee object should be instantiable
1184 a, b = tee('abc')
1185 c = type(a)('def')
1186 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +00001187
Raymond Hettingerad983e72003-11-12 14:32:26 +00001188 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +00001189 a, b, c = tee(range(2000), 3)
1190 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +00001191 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +00001192 self.assertEqual(list(b), list(range(2000)))
1193 self.assertEqual([next(c), next(c)], list(range(2)))
1194 self.assertEqual(list(a), list(range(100,2000)))
1195 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +00001196
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001197 # test values of n
1198 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +00001199 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +00001200 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001201 result = tee('abc', n)
1202 self.assertEqual(type(result), tuple)
1203 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001204 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001205
Raymond Hettingerad983e72003-11-12 14:32:26 +00001206 # tee pass-through to copyable iterator
1207 a, b = tee('abc')
1208 c, d = tee(a)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +00001209 self.assertTrue(a is c)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001210
Raymond Hettinger4cda01e2004-09-28 04:45:28 +00001211 # test tee_new
1212 t1, t2 = tee('abc')
1213 tnew = type(t1)
1214 self.assertRaises(TypeError, tnew)
1215 self.assertRaises(TypeError, tnew, 10)
1216 t3 = tnew(t1)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +00001217 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +00001218
Raymond Hettingera9f60922004-10-17 16:40:14 +00001219 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +00001220 a, b = tee(range(10))
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001221 p = weakref.proxy(a)
Raymond Hettingera9f60922004-10-17 16:40:14 +00001222 self.assertEqual(getattr(p, '__class__'), type(b))
1223 del a
1224 self.assertRaises(ReferenceError, getattr, p, '__class__')
1225
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001226 ans = list('abc')
1227 long_ans = list(range(10000))
1228
1229 # check copy
1230 a, b = tee('abc')
1231 self.assertEqual(list(copy.copy(a)), ans)
1232 self.assertEqual(list(copy.copy(b)), ans)
1233 a, b = tee(list(range(10000)))
1234 self.assertEqual(list(copy.copy(a)), long_ans)
1235 self.assertEqual(list(copy.copy(b)), long_ans)
1236
1237 # check partially consumed copy
1238 a, b = tee('abc')
1239 take(2, a)
1240 take(1, b)
1241 self.assertEqual(list(copy.copy(a)), ans[2:])
1242 self.assertEqual(list(copy.copy(b)), ans[1:])
1243 self.assertEqual(list(a), ans[2:])
1244 self.assertEqual(list(b), ans[1:])
1245 a, b = tee(range(10000))
1246 take(100, a)
1247 take(60, b)
1248 self.assertEqual(list(copy.copy(a)), long_ans[100:])
1249 self.assertEqual(list(copy.copy(b)), long_ans[60:])
1250 self.assertEqual(list(a), long_ans[100:])
1251 self.assertEqual(list(b), long_ans[60:])
1252
1253 # check deepcopy
1254 a, b = tee('abc')
1255 self.assertEqual(list(copy.deepcopy(a)), ans)
1256 self.assertEqual(list(copy.deepcopy(b)), ans)
1257 self.assertEqual(list(a), ans)
1258 self.assertEqual(list(b), ans)
1259 a, b = tee(range(10000))
1260 self.assertEqual(list(copy.deepcopy(a)), long_ans)
1261 self.assertEqual(list(copy.deepcopy(b)), long_ans)
1262 self.assertEqual(list(a), long_ans)
1263 self.assertEqual(list(b), long_ans)
1264
1265 # check partially consumed deepcopy
1266 a, b = tee('abc')
1267 take(2, a)
1268 take(1, b)
1269 self.assertEqual(list(copy.deepcopy(a)), ans[2:])
1270 self.assertEqual(list(copy.deepcopy(b)), ans[1:])
1271 self.assertEqual(list(a), ans[2:])
1272 self.assertEqual(list(b), ans[1:])
1273 a, b = tee(range(10000))
1274 take(100, a)
1275 take(60, b)
1276 self.assertEqual(list(copy.deepcopy(a)), long_ans[100:])
1277 self.assertEqual(list(copy.deepcopy(b)), long_ans[60:])
1278 self.assertEqual(list(a), long_ans[100:])
1279 self.assertEqual(list(b), long_ans[60:])
1280
1281 # check pickle
1282 self.pickletest(iter(tee('abc')))
1283 a, b = tee('abc')
1284 self.pickletest(a, compare=ans)
1285 self.pickletest(b, compare=ans)
1286
Serhiy Storchakaa3e91282013-01-25 13:19:31 +02001287 # Issue 13454: Crash when deleting backward iterator from tee()
1288 def test_tee_del_backward(self):
Serhiy Storchaka5bb893c2013-01-26 11:52:06 +02001289 forward, backward = tee(repeat(None, 20000000))
1290 any(forward) # exhaust the iterator
Serhiy Storchakaa3e91282013-01-25 13:19:31 +02001291 del backward
1292
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001293 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001294 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +00001295
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001296 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +00001297 self.assertRaises(StopIteration, next, f([]))
1298 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001299
Georg Brandla18af4e2007-04-21 15:47:16 +00001300 self.assertRaises(StopIteration, next, islice([], None))
1301 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001302
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001303 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +00001304 self.assertRaises(StopIteration, next, p)
1305 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001306 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +00001307 self.assertRaises(StopIteration, next, p)
1308 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001309
Georg Brandla18af4e2007-04-21 15:47:16 +00001310 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001311
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001312 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +00001313 self.assertRaises(StopIteration, next, f(lambda x:x, []))
1314 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001315
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001316class TestExamples(unittest.TestCase):
1317
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001318 def test_accumulate(self):
Raymond Hettinger482ba772010-12-01 22:48:00 +00001319 self.assertEqual(list(accumulate([1,2,3,4,5])), [1, 3, 6, 10, 15])
1320
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001321 def test_accumulate_reducible(self):
1322 # check copy, deepcopy, pickle
1323 data = [1, 2, 3, 4, 5]
1324 accumulated = [1, 3, 6, 10, 15]
1325 it = accumulate(data)
1326
1327 self.assertEqual(list(pickle.loads(pickle.dumps(it))), accumulated[:])
1328 self.assertEqual(next(it), 1)
1329 self.assertEqual(list(pickle.loads(pickle.dumps(it))), accumulated[1:])
1330 self.assertEqual(list(copy.deepcopy(it)), accumulated[1:])
1331 self.assertEqual(list(copy.copy(it)), accumulated[1:])
1332
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001333 def test_chain(self):
1334 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
1335
1336 def test_chain_from_iterable(self):
1337 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
1338
1339 def test_combinations(self):
1340 self.assertEqual(list(combinations('ABCD', 2)),
1341 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
1342 self.assertEqual(list(combinations(range(4), 3)),
1343 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
1344
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001345 def test_combinations_with_replacement(self):
1346 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
1347 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
1348
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001349 def test_compress(self):
1350 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
1351
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001352 def test_count(self):
1353 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
1354
1355 def test_cycle(self):
1356 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
1357
1358 def test_dropwhile(self):
1359 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
1360
1361 def test_groupby(self):
1362 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
1363 list('ABCDAB'))
1364 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
1365 [list('AAAA'), list('BBB'), list('CC'), list('D')])
1366
1367 def test_filter(self):
1368 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
1369
1370 def test_filterfalse(self):
1371 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
1372
1373 def test_map(self):
1374 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
1375
1376 def test_islice(self):
1377 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
1378 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
1379 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
1380 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1381
1382 def test_zip(self):
1383 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
1384
1385 def test_zip_longest(self):
1386 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
1387 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
1388
1389 def test_permutations(self):
1390 self.assertEqual(list(permutations('ABCD', 2)),
1391 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
1392 self.assertEqual(list(permutations(range(3))),
1393 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
1394
1395 def test_product(self):
1396 self.assertEqual(list(product('ABCD', 'xy')),
1397 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
1398 self.assertEqual(list(product(range(2), repeat=3)),
1399 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
1400 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
1401
1402 def test_repeat(self):
1403 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
1404
1405 def test_stapmap(self):
1406 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
1407 [32, 9, 1000])
1408
1409 def test_takewhile(self):
1410 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
1411
1412
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001413class TestGC(unittest.TestCase):
1414
1415 def makecycle(self, iterator, container):
1416 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +00001417 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001418 del container, iterator
1419
Raymond Hettinger482ba772010-12-01 22:48:00 +00001420 def test_accumulate(self):
1421 a = []
1422 self.makecycle(accumulate([1,2,a,3]), a)
1423
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001424 def test_chain(self):
1425 a = []
1426 self.makecycle(chain(a), a)
1427
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001428 def test_chain_from_iterable(self):
1429 a = []
1430 self.makecycle(chain.from_iterable([a]), a)
1431
1432 def test_combinations(self):
1433 a = []
1434 self.makecycle(combinations([1,2,a,3], 3), a)
1435
Raymond Hettingerd07d9392009-01-27 04:20:44 +00001436 def test_combinations_with_replacement(self):
1437 a = []
1438 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1439
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001440 def test_compress(self):
1441 a = []
1442 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1443
Raymond Hettingerd6280f42009-02-16 20:50:56 +00001444 def test_count(self):
1445 a = []
1446 Int = type('Int', (int,), dict(x=a))
1447 self.makecycle(count(Int(0), Int(1)), a)
1448
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001449 def test_cycle(self):
1450 a = []
1451 self.makecycle(cycle([a]*2), a)
1452
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001453 def test_dropwhile(self):
1454 a = []
1455 self.makecycle(dropwhile(bool, [0, a, a]), a)
1456
1457 def test_groupby(self):
1458 a = []
1459 self.makecycle(groupby([a]*2, lambda x:x), a)
1460
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001461 def test_issue2246(self):
1462 # Issue 2246 -- the _grouper iterator was not included in GC
1463 n = 10
1464 keyfunc = lambda x: x
1465 for i, j in groupby(range(n), key=keyfunc):
1466 keyfunc.__dict__.setdefault('x',[]).append(j)
1467
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001468 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001469 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001470 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001471
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001472 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001473 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001474 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001475
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001476 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001477 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001478 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001479
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001480 def test_zip_longest(self):
1481 a = []
1482 self.makecycle(zip_longest([a]*2, [a]*3), a)
1483 b = [a, None]
1484 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
1485
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001486 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001487 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001488 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001489
1490 def test_islice(self):
1491 a = []
1492 self.makecycle(islice([a]*2, None), a)
1493
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001494 def test_permutations(self):
1495 a = []
1496 self.makecycle(permutations([1,2,a,3], 3), a)
1497
1498 def test_product(self):
1499 a = []
1500 self.makecycle(product([1,2,a,3], repeat=3), a)
1501
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001502 def test_repeat(self):
1503 a = []
1504 self.makecycle(repeat(a), a)
1505
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001506 def test_starmap(self):
1507 a = []
1508 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1509
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001510 def test_takewhile(self):
1511 a = []
1512 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1513
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001514def R(seqn):
1515 'Regular generator'
1516 for i in seqn:
1517 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001518
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001519class G:
1520 'Sequence using __getitem__'
1521 def __init__(self, seqn):
1522 self.seqn = seqn
1523 def __getitem__(self, i):
1524 return self.seqn[i]
1525
1526class I:
1527 'Sequence using iterator protocol'
1528 def __init__(self, seqn):
1529 self.seqn = seqn
1530 self.i = 0
1531 def __iter__(self):
1532 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001533 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001534 if self.i >= len(self.seqn): raise StopIteration
1535 v = self.seqn[self.i]
1536 self.i += 1
1537 return v
1538
1539class Ig:
1540 'Sequence using iterator protocol defined with a generator'
1541 def __init__(self, seqn):
1542 self.seqn = seqn
1543 self.i = 0
1544 def __iter__(self):
1545 for val in self.seqn:
1546 yield val
1547
1548class X:
1549 'Missing __getitem__ and __iter__'
1550 def __init__(self, seqn):
1551 self.seqn = seqn
1552 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +00001553 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001554 if self.i >= len(self.seqn): raise StopIteration
1555 v = self.seqn[self.i]
1556 self.i += 1
1557 return v
1558
1559class N:
Georg Brandla18af4e2007-04-21 15:47:16 +00001560 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001561 def __init__(self, seqn):
1562 self.seqn = seqn
1563 self.i = 0
1564 def __iter__(self):
1565 return self
1566
1567class E:
1568 'Test propagation of exceptions'
1569 def __init__(self, seqn):
1570 self.seqn = seqn
1571 self.i = 0
1572 def __iter__(self):
1573 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001574 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001575 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001576
1577class S:
1578 'Test immediate stop'
1579 def __init__(self, seqn):
1580 pass
1581 def __iter__(self):
1582 return self
Georg Brandla18af4e2007-04-21 15:47:16 +00001583 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001584 raise StopIteration
1585
1586def L(seqn):
1587 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001588 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001589
1590
1591class TestVariousIteratorArgs(unittest.TestCase):
1592
Raymond Hettinger482ba772010-12-01 22:48:00 +00001593 def test_accumulate(self):
1594 s = [1,2,3,4,5]
1595 r = [1,3,6,10,15]
1596 n = len(s)
1597 for g in (G, I, Ig, L, R):
1598 self.assertEqual(list(accumulate(g(s))), r)
1599 self.assertEqual(list(accumulate(S(s))), [])
1600 self.assertRaises(TypeError, accumulate, X(s))
1601 self.assertRaises(TypeError, accumulate, N(s))
1602 self.assertRaises(ZeroDivisionError, list, accumulate(E(s)))
1603
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001604 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001605 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001606 for g in (G, I, Ig, S, L, R):
1607 self.assertEqual(list(chain(g(s))), list(g(s)))
1608 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +00001609 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +00001610 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001611 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1612
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001613 def test_compress(self):
1614 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1615 n = len(s)
1616 for g in (G, I, Ig, S, L, R):
1617 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1618 self.assertRaises(TypeError, compress, X(s), repeat(1))
1619 self.assertRaises(TypeError, compress, N(s), repeat(1))
1620 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1621
Christian Heimesc3f30c42008-02-22 16:37:40 +00001622 def test_product(self):
1623 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
1624 self.assertRaises(TypeError, product, X(s))
1625 self.assertRaises(TypeError, product, N(s))
1626 self.assertRaises(ZeroDivisionError, product, E(s))
1627
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001628 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001629 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001630 for g in (G, I, Ig, S, L, R):
1631 tgtlen = len(s) * 3
1632 expected = list(g(s))*3
1633 actual = list(islice(cycle(g(s)), tgtlen))
1634 self.assertEqual(actual, expected)
1635 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001636 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001637 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1638
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001639 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001640 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001641 for g in (G, I, Ig, S, L, R):
1642 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1643 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001644 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001645 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1646
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001647 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001648 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001649 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001650 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001651 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001652 self.assertRaises(TypeError, filter, isEven, X(s))
1653 self.assertRaises(TypeError, filter, isEven, N(s))
1654 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001655
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001656 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001657 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001658 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001659 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001660 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001661 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1662 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1663 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001664
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001665 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001666 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001667 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001668 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1669 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1670 self.assertRaises(TypeError, zip, X(s))
1671 self.assertRaises(TypeError, zip, N(s))
1672 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001673
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001674 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001675 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001676 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001677 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1678 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1679 self.assertRaises(TypeError, zip_longest, X(s))
1680 self.assertRaises(TypeError, zip_longest, N(s))
1681 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001682
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001683 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001684 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001685 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001686 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001687 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001688 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001689 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001690 self.assertRaises(TypeError, map, onearg, X(s))
1691 self.assertRaises(TypeError, map, onearg, N(s))
1692 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001693
1694 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001695 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001696 for g in (G, I, Ig, S, L, R):
1697 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1698 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001699 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001700 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1701
1702 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001703 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001704 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001705 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001706 self.assertEqual(list(starmap(operator.pow, g(ss))),
1707 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001708 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001709 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001710 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1711
1712 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001713 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001714 for g in (G, I, Ig, S, L, R):
1715 tgt = []
1716 for elem in g(s):
1717 if not isEven(elem): break
1718 tgt.append(elem)
1719 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1720 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001721 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001722 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1723
1724 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001725 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001726 for g in (G, I, Ig, S, L, R):
1727 tgt = []
1728 for elem in g(s):
1729 if not tgt and isOdd(elem): continue
1730 tgt.append(elem)
1731 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1732 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001733 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001734 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1735
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001736 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001737 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001738 for g in (G, I, Ig, S, L, R):
1739 it1, it2 = tee(g(s))
1740 self.assertEqual(list(it1), list(g(s)))
1741 self.assertEqual(list(it2), list(g(s)))
1742 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001743 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001744 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1745
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001746class LengthTransparency(unittest.TestCase):
1747
1748 def test_repeat(self):
Armin Ronacheraa9a79d2012-10-06 14:03:24 +02001749 self.assertEqual(operator.length_hint(repeat(None, 50)), 50)
Raymond Hettinger97d35552014-06-24 21:36:58 -07001750 self.assertEqual(operator.length_hint(repeat(None, 0)), 0)
Armin Ronacheraa9a79d2012-10-06 14:03:24 +02001751 self.assertEqual(operator.length_hint(repeat(None), 12), 12)
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001752
Raymond Hettinger97d35552014-06-24 21:36:58 -07001753 def test_repeat_with_negative_times(self):
1754 self.assertEqual(operator.length_hint(repeat(None, -1)), 0)
1755 self.assertEqual(operator.length_hint(repeat(None, -2)), 0)
1756 self.assertEqual(operator.length_hint(repeat(None, times=-1)), 0)
1757 self.assertEqual(operator.length_hint(repeat(None, times=-2)), 0)
1758
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001759class RegressionTests(unittest.TestCase):
1760
1761 def test_sf_793826(self):
1762 # Fix Armin Rigo's successful efforts to wreak havoc
1763
1764 def mutatingtuple(tuple1, f, tuple2):
1765 # this builds a tuple t which is a copy of tuple1,
1766 # then calls f(t), then mutates t to be equal to tuple2
1767 # (needs len(tuple1) == len(tuple2)).
1768 def g(value, first=[1]):
1769 if first:
1770 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00001771 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001772 return value
1773 items = list(tuple2)
1774 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001775 gen = map(g, items)
1776 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00001777 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001778
1779 def f(t):
1780 global T
1781 T = t
1782 first[:] = list(T)
1783
1784 first = []
1785 mutatingtuple((1,2,3), f, (4,5,6))
1786 second = list(T)
1787 self.assertEqual(first, second)
1788
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001789
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001790 def test_sf_950057(self):
1791 # Make sure that chain() and cycle() catch exceptions immediately
1792 # rather than when shifting between input sources
1793
1794 def gen1():
1795 hist.append(0)
1796 yield 1
1797 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001798 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001799 hist.append(2)
1800
1801 def gen2(x):
1802 hist.append(3)
1803 yield 2
1804 hist.append(4)
1805 if x:
1806 raise StopIteration
1807
1808 hist = []
1809 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1810 self.assertEqual(hist, [0,1])
1811
1812 hist = []
1813 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1814 self.assertEqual(hist, [0,1])
1815
1816 hist = []
1817 self.assertRaises(AssertionError, list, cycle(gen1()))
1818 self.assertEqual(hist, [0,1])
1819
Thomas Woutersb2137042007-02-01 18:02:27 +00001820class SubclassWithKwargsTest(unittest.TestCase):
1821 def test_keywords_in_subclass(self):
1822 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001823 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001824 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00001825 class Subclass(cls):
1826 def __init__(self, newarg=None, *args):
1827 cls.__init__(self, *args)
1828 try:
1829 Subclass(newarg=1)
1830 except TypeError as err:
1831 # we expect type errors because of wrong argument count
Ezio Melottib58e0bd2010-01-23 15:40:09 +00001832 self.assertNotIn("does not take keyword arguments", err.args[0])
Thomas Woutersb2137042007-02-01 18:02:27 +00001833
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02001834@support.cpython_only
1835class SizeofTest(unittest.TestCase):
1836 def setUp(self):
1837 self.ssize_t = struct.calcsize('n')
1838
1839 check_sizeof = support.check_sizeof
1840
1841 def test_product_sizeof(self):
1842 basesize = support.calcobjsize('3Pi')
1843 check = self.check_sizeof
1844 check(product('ab', '12'), basesize + 2 * self.ssize_t)
1845 check(product(*(('abc',) * 10)), basesize + 10 * self.ssize_t)
1846
1847 def test_combinations_sizeof(self):
1848 basesize = support.calcobjsize('3Pni')
1849 check = self.check_sizeof
1850 check(combinations('abcd', 3), basesize + 3 * self.ssize_t)
1851 check(combinations(range(10), 4), basesize + 4 * self.ssize_t)
1852
1853 def test_combinations_with_replacement_sizeof(self):
1854 cwr = combinations_with_replacement
1855 basesize = support.calcobjsize('3Pni')
1856 check = self.check_sizeof
1857 check(cwr('abcd', 3), basesize + 3 * self.ssize_t)
1858 check(cwr(range(10), 4), basesize + 4 * self.ssize_t)
1859
1860 def test_permutations_sizeof(self):
1861 basesize = support.calcobjsize('4Pni')
1862 check = self.check_sizeof
1863 check(permutations('abcd'),
1864 basesize + 4 * self.ssize_t + 4 * self.ssize_t)
1865 check(permutations('abcd', 3),
1866 basesize + 4 * self.ssize_t + 3 * self.ssize_t)
1867 check(permutations('abcde', 3),
1868 basesize + 5 * self.ssize_t + 3 * self.ssize_t)
1869 check(permutations(range(10), 4),
1870 basesize + 10 * self.ssize_t + 4 * self.ssize_t)
1871
Thomas Woutersb2137042007-02-01 18:02:27 +00001872
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001873libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001874
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001875
1876>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001877>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00001878... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001879...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001880Check 1200 is for $120.15
1881Check 1201 is for $764.05
1882Check 1202 is for $823.14
1883
1884>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001885>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00001886... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00001887...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000018881
18898
189027
1891
1892>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001893>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00001894... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00001895...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001896Alex
1897Laura
1898Martin
1899Walter
1900Samuele
1901
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001902>>> from operator import itemgetter
1903>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00001904>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001905>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001906... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001907...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000019081 ['a', 'c', 'e']
19092 ['b', 'd', 'f']
19103 ['g']
1911
Raymond Hettinger734fb572004-01-20 20:04:40 +00001912# Find runs of consecutive numbers using groupby. The key to the solution
1913# is differencing with a range so that consecutive numbers all appear in
1914# same group.
1915>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00001916>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001917... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001918...
Raymond Hettinger734fb572004-01-20 20:04:40 +00001919[1]
1920[4, 5, 6]
1921[10]
1922[15, 16, 17, 18]
1923[22]
1924[25, 26, 27, 28]
1925
Georg Brandl3dbca812008-07-23 16:10:53 +00001926>>> def take(n, iterable):
1927... "Return first n items of the iterable as a list"
1928... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001929
Georg Brandl3dbca812008-07-23 16:10:53 +00001930>>> def enumerate(iterable, start=0):
1931... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001932
Georg Brandl3dbca812008-07-23 16:10:53 +00001933>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001934... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00001935... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001936
Raymond Hettingercdf8ba32009-02-19 04:45:07 +00001937>>> def nth(iterable, n, default=None):
1938... "Returns the nth item or a default value"
1939... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001940
Georg Brandl3dbca812008-07-23 16:10:53 +00001941>>> def quantify(iterable, pred=bool):
1942... "Count how many times the predicate is true"
1943... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001944
Georg Brandl3dbca812008-07-23 16:10:53 +00001945>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001946... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00001947... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001948
Georg Brandl3dbca812008-07-23 16:10:53 +00001949>>> def ncycles(iterable, n):
Ezio Melotti13925002011-03-16 11:05:33 +02001950... "Returns the sequence elements n times"
Georg Brandl3dbca812008-07-23 16:10:53 +00001951... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001952
1953>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001954... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001955
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001956>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001957... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001958
1959>>> def repeatfunc(func, times=None, *args):
1960... "Repeat calls to func with specified arguments."
1961... " Example: repeatfunc(random.random)"
1962... if times is None:
1963... return starmap(func, repeat(args))
1964... else:
1965... return starmap(func, repeat(args, times))
1966
Raymond Hettingerd591f662003-10-26 15:34:50 +00001967>>> def pairwise(iterable):
1968... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1969... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001970... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00001971... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001972... except StopIteration:
1973... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001974... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001975
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001976>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001977... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001978... args = [iter(iterable)] * n
Raymond Hettinger883d2762009-01-27 04:57:51 +00001979... return zip_longest(*args, fillvalue=fillvalue)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001980
1981>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001982... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001983... # Recipe credited to George Sakkis
1984... pending = len(iterables)
1985... nexts = cycle(iter(it).__next__ for it in iterables)
1986... while pending:
1987... try:
1988... for next in nexts:
1989... yield next()
1990... except StopIteration:
1991... pending -= 1
1992... nexts = cycle(islice(nexts, pending))
1993
1994>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00001995... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1996... s = list(iterable)
1997... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001998
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001999>>> def unique_everseen(iterable, key=None):
2000... "List unique elements, preserving order. Remember all elements ever seen."
2001... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
2002... # unique_everseen('ABBCcAD', str.lower) --> A B C D
2003... seen = set()
2004... seen_add = seen.add
2005... if key is None:
2006... for element in iterable:
2007... if element not in seen:
2008... seen_add(element)
2009... yield element
2010... else:
2011... for element in iterable:
2012... k = key(element)
2013... if k not in seen:
2014... seen_add(k)
2015... yield element
2016
2017>>> def unique_justseen(iterable, key=None):
2018... "List unique elements, preserving order. Remember only the element just seen."
2019... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
2020... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
2021... return map(next, map(itemgetter(1), groupby(iterable, key)))
2022
Raymond Hettinger31b26f62014-04-02 03:16:42 -07002023>>> def first_true(iterable, default=False, pred=None):
2024... '''Returns the first true value in the iterable.
2025...
2026... If no true value is found, returns *default*
2027...
2028... If *pred* is not None, returns the first item
2029... for which pred(item) is true.
2030...
2031... '''
2032... # first_true([a,b,c], x) --> a or b or c or x
2033... # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x
2034... return next(filter(pred, iterable), default)
2035
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002036This is not part of the examples but it tests to make sure the definitions
2037perform as purported.
2038
Raymond Hettingera098b332003-09-08 23:58:40 +00002039>>> take(10, count())
2040[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
2041
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002042>>> list(enumerate('abc'))
2043[(0, 'a'), (1, 'b'), (2, 'c')]
2044
2045>>> list(islice(tabulate(lambda x: 2*x), 4))
2046[0, 2, 4, 6]
2047
2048>>> nth('abcde', 3)
Raymond Hettingerd04fa312009-02-04 19:45:13 +00002049'd'
2050
2051>>> nth('abcde', 9) is None
2052True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002053
Guido van Rossum805365e2007-05-07 22:24:25 +00002054>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000205550
2056
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002057>>> a = [[1, 2, 3], [4, 5, 6]]
2058>>> flatten(a)
2059[1, 2, 3, 4, 5, 6]
2060
2061>>> list(repeatfunc(pow, 5, 2, 3))
2062[8, 8, 8, 8, 8]
2063
2064>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00002065>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002066[0, 0, 0, 0, 0]
2067
Raymond Hettingerd591f662003-10-26 15:34:50 +00002068>>> list(pairwise('abcd'))
2069[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002070
Raymond Hettingerd591f662003-10-26 15:34:50 +00002071>>> list(pairwise([]))
2072[]
2073
2074>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00002075[]
2076
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002077>>> list(islice(padnone('abc'), 0, 6))
2078['a', 'b', 'c', None, None, None]
2079
2080>>> list(ncycles('abc', 3))
2081['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
2082
2083>>> dotproduct([1,2,3], [4,5,6])
208432
2085
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002086>>> list(grouper(3, 'abcdefg', 'x'))
2087[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
2088
2089>>> list(roundrobin('abc', 'd', 'ef'))
2090['a', 'd', 'e', 'b', 'f', 'c']
2091
Raymond Hettingerace67332009-01-26 02:23:50 +00002092>>> list(powerset([1,2,3]))
2093[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002094
Raymond Hettinger191e8502009-01-27 13:29:43 +00002095>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
2096True
2097
2098>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
2099True
2100
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00002101>>> list(unique_everseen('AAAABBBCCDAABBB'))
2102['A', 'B', 'C', 'D']
2103
2104>>> list(unique_everseen('ABBCcAD', str.lower))
2105['A', 'B', 'C', 'D']
2106
2107>>> list(unique_justseen('AAAABBBCCDAABBB'))
2108['A', 'B', 'C', 'D', 'A', 'B']
2109
2110>>> list(unique_justseen('ABBCcAD', str.lower))
2111['A', 'B', 'C', 'A', 'D']
2112
Raymond Hettinger31b26f62014-04-02 03:16:42 -07002113>>> first_true('ABC0DEF1', '9', str.isdigit)
2114'0'
2115
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002116"""
2117
2118__test__ = {'libreftest' : libreftest}
2119
2120def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002121 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00002122 RegressionTests, LengthTransparency,
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002123 SubclassWithKwargsTest, TestExamples,
2124 SizeofTest)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002125 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002126
2127 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002128 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00002129 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00002130 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00002131 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002132 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00002133 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00002134 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00002135 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002136
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002137 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002138 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00002139
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002140if __name__ == "__main__":
2141 test_main(verbose=True)