blob: 04279792065edfae8c69605f1779b45339960baa [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001import unittest
2from test import test_support
3from itertools import *
Antoine Pitrou3ec903f2014-04-29 12:13:46 +02004import weakref
Raymond Hettingeraa044612009-02-12 12:04:26 +00005from decimal import Decimal
Raymond Hettingerdbe3bfb2009-02-12 12:43:01 +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 Hettingere09f45a2009-11-30 19:44:40 +000010import copy
11import pickle
Ezio Melottidde5b942010-02-03 05:37:26 +000012from functools import reduce
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +000013maxsize = test_support.MAX_Py_ssize_t
14minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000015
16def onearg(x):
17 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000018 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000019
20def errfunc(*args):
21 'Test function that raises an error'
22 raise ValueError
23
24def gen3():
25 'Non-restartable source sequence'
26 for i in (0, 1, 2):
27 yield i
28
29def isEven(x):
30 'Test predicate'
31 return x%2==0
32
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000033def isOdd(x):
34 'Test predicate'
35 return x%2==1
36
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000037class StopNow:
38 'Class emulating an empty iterable.'
39 def __iter__(self):
40 return self
41 def next(self):
42 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000043
Raymond Hettinger02420702003-06-29 20:36:23 +000044def take(n, seq):
45 'Convenience function for partially consuming a long of infinite iterable'
46 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000047
Raymond Hettingerd553d852008-03-04 04:17:08 +000048def prod(iterable):
49 return reduce(operator.mul, iterable, 1)
50
Raymond Hettinger93e804d2008-02-26 23:40:50 +000051def fact(n):
52 'Factorial'
Raymond Hettingerd553d852008-03-04 04:17:08 +000053 return prod(range(1, n+1))
54
Raymond Hettinger96ef8112003-02-01 00:10:11 +000055class TestBasicOps(unittest.TestCase):
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000056 def test_chain(self):
Raymond Hettingerad47fa12008-03-06 20:52:01 +000057
58 def chain2(*iterables):
59 'Pure python version in the docs'
60 for it in iterables:
61 for element in it:
62 yield element
63
64 for c in (chain, chain2):
65 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
66 self.assertEqual(list(c('abc')), list('abc'))
67 self.assertEqual(list(c('')), [])
68 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
69 self.assertRaises(TypeError, list,c(2, 3))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000070
Raymond Hettingerb4cbc982008-02-28 22:46:41 +000071 def test_chain_from_iterable(self):
72 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
73 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
74 self.assertEqual(list(chain.from_iterable([''])), [])
75 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
76 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
77
Raymond Hettinger93e804d2008-02-26 23:40:50 +000078 def test_combinations(self):
Raymond Hettinger5b913e32009-01-08 06:39:04 +000079 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Raymond Hettinger93e804d2008-02-26 23:40:50 +000080 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Raymond Hettingerd553d852008-03-04 04:17:08 +000081 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Raymond Hettinger93e804d2008-02-26 23:40:50 +000082 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Raymond Hettinger5b913e32009-01-08 06:39:04 +000083 self.assertEqual(list(combinations('abc', 32)), []) # r > n
Raymond Hettinger93e804d2008-02-26 23:40:50 +000084 self.assertEqual(list(combinations(range(4), 3)),
85 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
Raymond Hettingerd553d852008-03-04 04:17:08 +000086
87 def combinations1(iterable, r):
88 'Pure python version shown in the docs'
89 pool = tuple(iterable)
90 n = len(pool)
Raymond Hettinger5b913e32009-01-08 06:39:04 +000091 if r > n:
92 return
Raymond Hettingerd553d852008-03-04 04:17:08 +000093 indices = range(r)
94 yield tuple(pool[i] for i in indices)
95 while 1:
96 for i in reversed(range(r)):
97 if indices[i] != i + n - r:
98 break
99 else:
100 return
101 indices[i] += 1
102 for j in range(i+1, r):
103 indices[j] = indices[j-1] + 1
104 yield tuple(pool[i] for i in indices)
105
106 def combinations2(iterable, r):
107 'Pure python version shown in the docs'
108 pool = tuple(iterable)
109 n = len(pool)
110 for indices in permutations(range(n), r):
111 if sorted(indices) == list(indices):
112 yield tuple(pool[i] for i in indices)
113
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000114 def combinations3(iterable, r):
115 'Pure python version from cwr()'
116 pool = tuple(iterable)
117 n = len(pool)
118 for indices in combinations_with_replacement(range(n), r):
119 if len(set(indices)) == r:
120 yield tuple(pool[i] for i in indices)
121
Raymond Hettingerd553d852008-03-04 04:17:08 +0000122 for n in range(7):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000123 values = [5*x-12 for x in range(n)]
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000124 for r in range(n+2):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000125 result = list(combinations(values, r))
Ezio Melottidde5b942010-02-03 05:37:26 +0000126 self.assertEqual(len(result), 0 if r>n else fact(n) // fact(r) // fact(n-r)) # right number of combs
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000127 self.assertEqual(len(result), len(set(result))) # no repeats
128 self.assertEqual(result, sorted(result)) # lexicographic order
129 for c in result:
130 self.assertEqual(len(c), r) # r-length combinations
131 self.assertEqual(len(set(c)), r) # no duplicate elements
132 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Peterson5c8da862009-06-30 22:57:08 +0000133 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000134 self.assertEqual(list(c),
135 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Raymond Hettingerd553d852008-03-04 04:17:08 +0000136 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000137 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000138 self.assertEqual(result, list(combinations3(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000139
Benjamin Peterson021dec12015-02-01 20:59:00 -0500140 @test_support.bigaddrspacetest
141 def test_combinations_overflow(self):
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +0200142 with self.assertRaises((OverflowError, MemoryError)):
Benjamin Peterson021dec12015-02-01 20:59:00 -0500143 combinations("AA", 2**29)
144
Alex Gaynor97376482011-07-17 16:44:31 -0700145 @test_support.impl_detail("tuple reuse is specific to CPython")
146 def test_combinations_tuple_reuse(self):
Raymond Hettingerd553d852008-03-04 04:17:08 +0000147 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
148 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
149
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000150 def test_combinations_with_replacement(self):
151 cwr = combinations_with_replacement
152 self.assertRaises(TypeError, cwr, 'abc') # missing r argument
153 self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
154 self.assertRaises(TypeError, cwr, None) # pool is not iterable
155 self.assertRaises(ValueError, cwr, 'abc', -2) # r is negative
156 self.assertEqual(list(cwr('ABC', 2)),
157 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
158
159 def cwr1(iterable, r):
160 'Pure python version shown in the docs'
161 # number items returned: (n+r-1)! / r! / (n-1)! when n>0
162 pool = tuple(iterable)
163 n = len(pool)
164 if not n and r:
165 return
166 indices = [0] * r
167 yield tuple(pool[i] for i in indices)
168 while 1:
169 for i in reversed(range(r)):
170 if indices[i] != n - 1:
171 break
172 else:
173 return
174 indices[i:] = [indices[i] + 1] * (r - i)
175 yield tuple(pool[i] for i in indices)
176
177 def cwr2(iterable, r):
178 'Pure python version shown in the docs'
179 pool = tuple(iterable)
180 n = len(pool)
181 for indices in product(range(n), repeat=r):
182 if sorted(indices) == list(indices):
183 yield tuple(pool[i] for i in indices)
184
185 def numcombs(n, r):
186 if not n:
187 return 0 if r else 1
Ezio Melottidde5b942010-02-03 05:37:26 +0000188 return fact(n+r-1) // fact(r) // fact(n-1)
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000189
190 for n in range(7):
191 values = [5*x-12 for x in range(n)]
192 for r in range(n+2):
193 result = list(cwr(values, r))
194
195 self.assertEqual(len(result), numcombs(n, r)) # right number of combs
196 self.assertEqual(len(result), len(set(result))) # no repeats
197 self.assertEqual(result, sorted(result)) # lexicographic order
198
199 regular_combs = list(combinations(values, r)) # compare to combs without replacement
200 if n == 0 or r <= 1:
Ezio Melotti2623a372010-11-21 13:34:58 +0000201 self.assertEqual(result, regular_combs) # cases that should be identical
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000202 else:
Benjamin Peterson5c8da862009-06-30 22:57:08 +0000203 self.assertTrue(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000204
205 for c in result:
206 self.assertEqual(len(c), r) # r-length combinations
207 noruns = [k for k,v in groupby(c)] # combo without consecutive repeats
208 self.assertEqual(len(noruns), len(set(noruns))) # no repeats other than consecutive
209 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Peterson5c8da862009-06-30 22:57:08 +0000210 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000211 self.assertEqual(noruns,
212 [e for e in values if e in c]) # comb is a subsequence of the input iterable
213 self.assertEqual(result, list(cwr1(values, r))) # matches first pure python version
214 self.assertEqual(result, list(cwr2(values, r))) # matches second pure python version
215
Benjamin Peterson17845c12015-02-01 21:10:47 -0500216 @test_support.bigaddrspacetest
217 def test_combinations_with_replacement_overflow(self):
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +0200218 with self.assertRaises((OverflowError, MemoryError)):
Benjamin Peterson17845c12015-02-01 21:10:47 -0500219 combinations_with_replacement("AA", 2**30)
220
Alex Gaynor97376482011-07-17 16:44:31 -0700221 @test_support.impl_detail("tuple reuse is specific to CPython")
222 def test_combinations_with_replacement_tuple_reuse(self):
223 cwr = combinations_with_replacement
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000224 self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
225 self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
226
Raymond Hettingerd553d852008-03-04 04:17:08 +0000227 def test_permutations(self):
228 self.assertRaises(TypeError, permutations) # too few arguments
229 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000230 self.assertRaises(TypeError, permutations, None) # pool is not iterable
231 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000232 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000233 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Raymond Hettingerd553d852008-03-04 04:17:08 +0000234 self.assertEqual(list(permutations(range(3), 2)),
235 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
236
237 def permutations1(iterable, r=None):
238 'Pure python version shown in the docs'
239 pool = tuple(iterable)
240 n = len(pool)
241 r = n if r is None else r
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000242 if r > n:
243 return
Raymond Hettingerd553d852008-03-04 04:17:08 +0000244 indices = range(n)
Raymond Hettingere70bb8d2008-03-23 00:55:46 +0000245 cycles = range(n, n-r, -1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000246 yield tuple(pool[i] for i in indices[:r])
247 while n:
248 for i in reversed(range(r)):
249 cycles[i] -= 1
250 if cycles[i] == 0:
251 indices[i:] = indices[i+1:] + indices[i:i+1]
252 cycles[i] = n - i
253 else:
254 j = cycles[i]
255 indices[i], indices[-j] = indices[-j], indices[i]
256 yield tuple(pool[i] for i in indices[:r])
257 break
258 else:
259 return
260
261 def permutations2(iterable, r=None):
262 'Pure python version shown in the docs'
263 pool = tuple(iterable)
264 n = len(pool)
265 r = n if r is None else r
266 for indices in product(range(n), repeat=r):
267 if len(set(indices)) == r:
268 yield tuple(pool[i] for i in indices)
269
270 for n in range(7):
271 values = [5*x-12 for x in range(n)]
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000272 for r in range(n+2):
Raymond Hettingerd553d852008-03-04 04:17:08 +0000273 result = list(permutations(values, r))
Ezio Melottidde5b942010-02-03 05:37:26 +0000274 self.assertEqual(len(result), 0 if r>n else fact(n) // fact(n-r)) # right number of perms
Raymond Hettingerd553d852008-03-04 04:17:08 +0000275 self.assertEqual(len(result), len(set(result))) # no repeats
276 self.assertEqual(result, sorted(result)) # lexicographic order
277 for p in result:
278 self.assertEqual(len(p), r) # r-length permutations
279 self.assertEqual(len(set(p)), r) # no duplicate elements
Benjamin Peterson5c8da862009-06-30 22:57:08 +0000280 self.assertTrue(all(e in values for e in p)) # elements taken from input iterable
Raymond Hettingerd553d852008-03-04 04:17:08 +0000281 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000282 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000283 if r == n:
284 self.assertEqual(result, list(permutations(values, None))) # test r as None
285 self.assertEqual(result, list(permutations(values))) # test default r
286
Benjamin Petersondda91212015-02-01 21:34:07 -0500287 @test_support.bigaddrspacetest
288 def test_permutations_overflow(self):
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +0200289 with self.assertRaises((OverflowError, MemoryError)):
Benjamin Petersondda91212015-02-01 21:34:07 -0500290 permutations("A", 2**30)
Benjamin Petersondda91212015-02-01 21:34:07 -0500291
Zachary Ware4e0df172014-04-24 13:20:27 -0500292 @test_support.impl_detail("tuple reuse is specific to CPython")
Alex Gaynor97376482011-07-17 16:44:31 -0700293 def test_permutations_tuple_reuse(self):
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000294 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000295 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000296
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000297 def test_combinatorics(self):
298 # Test relationships between product(), permutations(),
299 # combinations() and combinations_with_replacement().
300
Raymond Hettinger2f6c2e02009-01-27 10:36:14 +0000301 for n in range(6):
302 s = 'ABCDEFG'[:n]
303 for r in range(8):
304 prod = list(product(s, repeat=r))
305 cwr = list(combinations_with_replacement(s, r))
306 perm = list(permutations(s, r))
307 comb = list(combinations(s, r))
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000308
Raymond Hettinger2f6c2e02009-01-27 10:36:14 +0000309 # Check size
Ezio Melotti2623a372010-11-21 13:34:58 +0000310 self.assertEqual(len(prod), n**r)
311 self.assertEqual(len(cwr), (fact(n+r-1) // fact(r) // fact(n-1)) if n else (not r))
312 self.assertEqual(len(perm), 0 if r>n else fact(n) // fact(n-r))
313 self.assertEqual(len(comb), 0 if r>n else fact(n) // fact(r) // fact(n-r))
Raymond Hettinger2f6c2e02009-01-27 10:36:14 +0000314
315 # Check lexicographic order without repeated tuples
Ezio Melotti2623a372010-11-21 13:34:58 +0000316 self.assertEqual(prod, sorted(set(prod)))
317 self.assertEqual(cwr, sorted(set(cwr)))
318 self.assertEqual(perm, sorted(set(perm)))
319 self.assertEqual(comb, sorted(set(comb)))
Raymond Hettinger2f6c2e02009-01-27 10:36:14 +0000320
321 # Check interrelationships
Ezio Melotti2623a372010-11-21 13:34:58 +0000322 self.assertEqual(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
323 self.assertEqual(perm, [t for t in prod if len(set(t))==r]) # perm: prods with no dups
Raymond Hettinger2f6c2e02009-01-27 10:36:14 +0000324 self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
325 self.assertEqual(comb, [t for t in cwr if len(set(t))==r]) # comb: cwrs without dups
326 self.assertEqual(comb, filter(set(cwr).__contains__, perm)) # comb: perm that is a cwr
327 self.assertEqual(comb, filter(set(perm).__contains__, cwr)) # comb: cwr that is a perm
328 self.assertEqual(comb, sorted(set(cwr) & set(perm))) # comb: both a cwr and a perm
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000329
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000330 def test_compress(self):
Raymond Hettinger2e2909f2009-02-19 02:15:14 +0000331 self.assertEqual(list(compress(data='ABCDEF', selectors=[1,0,1,0,1,1])), list('ACEF'))
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000332 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
333 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
334 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
335 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
336 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
337 n = 10000
338 data = chain.from_iterable(repeat(range(6), n))
339 selectors = chain.from_iterable(repeat((0, 1)))
340 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
341 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
342 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
343 self.assertRaises(TypeError, compress, range(6)) # too few args
344 self.assertRaises(TypeError, compress, range(6), None) # too many args
345
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000346 def test_count(self):
347 self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
348 self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000349 self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000350 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
351 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000352 self.assertRaises(TypeError, count, 2, 3, 4)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000353 self.assertRaises(TypeError, count, 'a')
Serhiy Storchaka7f8ce852016-09-10 09:53:29 +0300354 self.assertEqual(take(10, count(maxsize-5)), range(maxsize-5, maxsize+5))
355 self.assertEqual(take(10, count(-maxsize-5)), range(-maxsize-5, -maxsize+5))
356 self.assertEqual(take(3, count(3.25)), [3.25, 4.25, 5.25])
357 self.assertEqual(take(3, count(3.25-4j)), [3.25-4j, 4.25-4j, 5.25-4j])
358 self.assertEqual(take(3, count(Decimal('1.1'))),
359 [Decimal('1.1'), Decimal('2.1'), Decimal('3.1')])
360 self.assertEqual(take(3, count(Fraction(2, 3))),
361 [Fraction(2, 3), Fraction(5, 3), Fraction(8, 3)])
362 BIGINT = 1<<1000
363 self.assertEqual(take(3, count(BIGINT)), [BIGINT, BIGINT+1, BIGINT+2])
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000364 c = count(3)
365 self.assertEqual(repr(c), 'count(3)')
366 c.next()
367 self.assertEqual(repr(c), 'count(4)')
Jack Diederich36234e82006-09-21 17:50:26 +0000368 c = count(-9)
369 self.assertEqual(repr(c), 'count(-9)')
370 c.next()
Serhiy Storchaka7f8ce852016-09-10 09:53:29 +0300371 self.assertEqual(next(c), -8)
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000372 self.assertEqual(repr(count(10.25)), 'count(10.25)')
Serhiy Storchaka7f8ce852016-09-10 09:53:29 +0300373 self.assertEqual(repr(count(10.0)), 'count(10.0)')
374 self.assertEqual(type(next(count(10.0))), float)
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000375 for i in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 10, sys.maxint-5, sys.maxint+5):
Raymond Hettinger3a0de082007-10-12 17:53:11 +0000376 # Test repr (ignoring the L in longs)
377 r1 = repr(count(i)).replace('L', '')
378 r2 = 'count(%r)'.__mod__(i).replace('L', '')
379 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000380
Raymond Hettingere09f45a2009-11-30 19:44:40 +0000381 # check copy, deepcopy, pickle
382 for value in -3, 3, sys.maxint-5, sys.maxint+5:
383 c = count(value)
384 self.assertEqual(next(copy.copy(c)), value)
385 self.assertEqual(next(copy.deepcopy(c)), value)
Serhiy Storchaka655720e2014-12-15 14:02:43 +0200386 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
387 self.assertEqual(next(pickle.loads(pickle.dumps(c, proto))), value)
Raymond Hettingere09f45a2009-11-30 19:44:40 +0000388
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000389 def test_count_with_stride(self):
390 self.assertEqual(zip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingera4038032009-02-14 00:25:51 +0000391 self.assertEqual(zip('abc',count(start=2,step=3)),
392 [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingeraa681c72009-02-21 07:17:22 +0000393 self.assertEqual(zip('abc',count(step=-1)),
394 [('a', 0), ('b', -1), ('c', -2)])
Serhiy Storchaka7f8ce852016-09-10 09:53:29 +0300395 self.assertRaises(TypeError, count, 'a', 'b')
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000396 self.assertEqual(zip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
397 self.assertEqual(zip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
Serhiy Storchaka7f8ce852016-09-10 09:53:29 +0300398 self.assertEqual(zip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
399 self.assertEqual(zip('abc',count(2,1L)), [('a', 2), ('b', 3), ('c', 4)])
400 self.assertEqual(zip('abc',count(2,3L)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000401 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
402 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
Serhiy Storchaka7f8ce852016-09-10 09:53:29 +0300403 self.assertEqual(take(3, count(10, maxsize+5)),
404 range(10, 10+3*(maxsize+5), maxsize+5))
405 self.assertEqual(take(3, count(2, 1.25)), [2, 3.25, 4.5])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000406 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettingeraa044612009-02-12 12:04:26 +0000407 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
408 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerdbe3bfb2009-02-12 12:43:01 +0000409 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
410 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Serhiy Storchaka7f8ce852016-09-10 09:53:29 +0300411 BIGINT = 1<<1000
412 self.assertEqual(take(3, count(step=BIGINT)), [0, BIGINT, 2*BIGINT])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000413 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
414 c = count(3, 5)
415 self.assertEqual(repr(c), 'count(3, 5)')
416 c.next()
417 self.assertEqual(repr(c), 'count(8, 5)')
418 c = count(-9, 0)
419 self.assertEqual(repr(c), 'count(-9, 0)')
420 c.next()
421 self.assertEqual(repr(c), 'count(-9, 0)')
422 c = count(-9, -3)
423 self.assertEqual(repr(c), 'count(-9, -3)')
424 c.next()
425 self.assertEqual(repr(c), 'count(-12, -3)')
426 self.assertEqual(repr(c), 'count(-12, -3)')
427 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
428 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
429 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
Serhiy Storchaka7f8ce852016-09-10 09:53:29 +0300430 self.assertEqual(repr(count(10, 1.00)), 'count(10, 1.0)')
431 c = count(10, 1.0)
432 self.assertEqual(type(next(c)), int)
433 self.assertEqual(type(next(c)), float)
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000434 for i in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 10, sys.maxint-5, sys.maxint+5):
435 for j in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 1, 10, sys.maxint-5, sys.maxint+5):
436 # Test repr (ignoring the L in longs)
437 r1 = repr(count(i, j)).replace('L', '')
438 if j == 1:
439 r2 = ('count(%r)' % i).replace('L', '')
440 else:
441 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
442 self.assertEqual(r1, r2)
443
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000444 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000445 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000446 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000447 self.assertRaises(TypeError, cycle)
448 self.assertRaises(TypeError, cycle, 5)
449 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000450
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000451 def test_groupby(self):
452 # Check whether it accepts arguments correctly
453 self.assertEqual([], list(groupby([])))
454 self.assertEqual([], list(groupby([], key=id)))
455 self.assertRaises(TypeError, list, groupby('abc', []))
456 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000457 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000458
459 # Check normal input
460 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
461 (2,15,22), (3,16,23), (3,17,23)]
462 dup = []
463 for k, g in groupby(s, lambda r:r[0]):
464 for elem in g:
465 self.assertEqual(k, elem[0])
466 dup.append(elem)
467 self.assertEqual(s, dup)
468
469 # Check nested case
470 dup = []
471 for k, g in groupby(s, lambda r:r[0]):
472 for ik, ig in groupby(g, lambda r:r[2]):
473 for elem in ig:
474 self.assertEqual(k, elem[0])
475 self.assertEqual(ik, elem[2])
476 dup.append(elem)
477 self.assertEqual(s, dup)
478
479 # Check case where inner iterator is not used
480 keys = [k for k, g in groupby(s, lambda r:r[0])]
481 expectedkeys = set([r[0] for r in s])
482 self.assertEqual(set(keys), expectedkeys)
483 self.assertEqual(len(keys), len(expectedkeys))
484
485 # Exercise pipes and filters style
486 s = 'abracadabra'
487 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000488 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000489 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
490 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000491 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000492 self.assertEqual(r, ['a', 'b', 'r'])
493 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000494 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000495 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
496 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000497 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000498 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
499
500 # iter.next failure
501 class ExpectedError(Exception):
502 pass
503 def delayed_raise(n=0):
504 for i in range(n):
505 yield 'yo'
506 raise ExpectedError
507 def gulp(iterable, keyp=None, func=list):
508 return [func(g) for k, g in groupby(iterable, keyp)]
509
510 # iter.next failure on outer object
511 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
512 # iter.next failure on inner object
513 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
514
515 # __cmp__ failure
516 class DummyCmp:
517 def __cmp__(self, dst):
518 raise ExpectedError
519 s = [DummyCmp(), DummyCmp(), None]
520
521 # __cmp__ failure on outer object
522 self.assertRaises(ExpectedError, gulp, s, func=id)
523 # __cmp__ failure on inner object
524 self.assertRaises(ExpectedError, gulp, s)
525
526 # keyfunc failure
527 def keyfunc(obj):
528 if keyfunc.skip > 0:
529 keyfunc.skip -= 1
530 return obj
531 else:
532 raise ExpectedError
533
534 # keyfunc failure on outer object
535 keyfunc.skip = 0
536 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
537 keyfunc.skip = 1
538 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
539
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000540 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000541 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000542 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000543 self.assertEqual(list(ifilter(bool, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000544 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000545 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000546 self.assertRaises(TypeError, ifilter, lambda x:x)
547 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000548 self.assertRaises(TypeError, ifilter, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000549 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000550
551 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000552 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
553 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000554 self.assertEqual(list(ifilterfalse(bool, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000555 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000556 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000557 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
558 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000559 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000560 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000561
562 def test_izip(self):
563 ans = [(x,y) for x, y in izip('abc',count())]
564 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000565 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
566 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
Raymond Hettinger02420702003-06-29 20:36:23 +0000567 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000568 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000569 self.assertEqual(list(izip()), zip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000570 self.assertRaises(TypeError, izip, 3)
571 self.assertRaises(TypeError, izip, range(3), 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000572 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
573 zip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000574 self.assertEqual([pair for pair in izip('abc', 'def')],
575 zip('abc', 'def'))
Alex Gaynor97376482011-07-17 16:44:31 -0700576
577 @test_support.impl_detail("tuple reuse is specific to CPython")
Zachary Ware4e0df172014-04-24 13:20:27 -0500578 def test_izip_tuple_reuse(self):
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000579 ids = map(id, izip('abc', 'def'))
580 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000581 ids = map(id, list(izip('abc', 'def')))
582 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000583
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000584 def test_iziplongest(self):
585 for args in [
586 ['abc', range(6)],
587 [range(6), 'abc'],
588 [range(1000), range(2000,2100), range(3000,3050)],
589 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
590 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
591 ]:
Ezio Melottidde5b942010-02-03 05:37:26 +0000592 # target = map(None, *args) <- this raises a py3k warning
593 # this is the replacement:
594 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
595 for i in range(max(map(len, args)))]
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000596 self.assertEqual(list(izip_longest(*args)), target)
597 self.assertEqual(list(izip_longest(*args, **{})), target)
598 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
599 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Tim Petersea5962f2007-03-12 18:07:52 +0000600
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000601 self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
602
603 self.assertEqual(list(izip_longest()), zip())
604 self.assertEqual(list(izip_longest([])), zip([]))
605 self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
Tim Petersea5962f2007-03-12 18:07:52 +0000606
Ezio Melottidde5b942010-02-03 05:37:26 +0000607 self.assertEqual(list(izip_longest('abc', 'defg', **{})),
608 zip(list('abc') + [None], 'defg')) # empty keyword dict
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000609 self.assertRaises(TypeError, izip_longest, 3)
610 self.assertRaises(TypeError, izip_longest, range(3), 3)
611
612 for stmt in [
613 "izip_longest('abc', fv=1)",
Tim Petersea5962f2007-03-12 18:07:52 +0000614 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000615 ]:
616 try:
617 eval(stmt, globals(), locals())
618 except TypeError:
619 pass
620 else:
621 self.fail('Did not raise Type in: ' + stmt)
Tim Petersea5962f2007-03-12 18:07:52 +0000622
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000623 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
624 zip('abc', 'def'))
625 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
626 zip('abc', 'def'))
Alex Gaynor97376482011-07-17 16:44:31 -0700627
628 @test_support.impl_detail("tuple reuse is specific to CPython")
629 def test_izip_longest_tuple_reuse(self):
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000630 ids = map(id, izip_longest('abc', 'def'))
631 self.assertEqual(min(ids), max(ids))
632 ids = map(id, list(izip_longest('abc', 'def')))
633 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
634
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +0000635 def test_bug_7244(self):
636
637 class Repeater(object):
638 # this class is similar to itertools.repeat
639 def __init__(self, o, t, e):
640 self.o = o
641 self.t = int(t)
642 self.e = e
643 def __iter__(self): # its iterator is itself
644 return self
645 def next(self):
646 if self.t > 0:
647 self.t -= 1
648 return self.o
649 else:
650 raise self.e
651
652 # Formerly this code in would fail in debug mode
653 # with Undetected Error and Stop Iteration
654 r1 = Repeater(1, 3, StopIteration)
655 r2 = Repeater(2, 4, StopIteration)
656 def run(r1, r2):
657 result = []
658 for i, j in izip_longest(r1, r2, fillvalue=0):
659 with test_support.captured_output('stdout'):
660 print (i, j)
661 result.append((i, j))
662 return result
663 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
664
665 # Formerly, the RuntimeError would be lost
666 # and StopIteration would stop as expected
667 r1 = Repeater(1, 3, RuntimeError)
668 r2 = Repeater(2, 4, StopIteration)
669 it = izip_longest(r1, r2, fillvalue=0)
670 self.assertEqual(next(it), (1, 2))
671 self.assertEqual(next(it), (1, 2))
672 self.assertEqual(next(it), (1, 2))
673 self.assertRaises(RuntimeError, next, it)
674
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000675 def test_product(self):
676 for args, result in [
Raymond Hettingerd553d852008-03-04 04:17:08 +0000677 ([], [()]), # zero iterables
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000678 (['ab'], [('a',), ('b',)]), # one iterable
679 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
680 ([range(0), range(2), range(3)], []), # first iterable with zero length
681 ([range(2), range(0), range(3)], []), # middle iterable with zero length
682 ([range(2), range(3), range(0)], []), # last iterable with zero length
683 ]:
684 self.assertEqual(list(product(*args)), result)
Raymond Hettinger08ff6822008-02-29 02:21:48 +0000685 for r in range(4):
686 self.assertEqual(list(product(*(args*r))),
687 list(product(*args, **dict(repeat=r))))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000688 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
689 self.assertRaises(TypeError, product, range(6), None)
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000690
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000691 def product1(*args, **kwds):
692 pools = map(tuple, args) * kwds.get('repeat', 1)
693 n = len(pools)
694 if n == 0:
695 yield ()
696 return
697 if any(len(pool) == 0 for pool in pools):
698 return
699 indices = [0] * n
700 yield tuple(pool[i] for pool, i in zip(pools, indices))
701 while 1:
702 for i in reversed(range(n)): # right to left
703 if indices[i] == len(pools[i]) - 1:
704 continue
705 indices[i] += 1
706 for j in range(i+1, n):
707 indices[j] = 0
708 yield tuple(pool[i] for pool, i in zip(pools, indices))
709 break
710 else:
711 return
712
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000713 def product2(*args, **kwds):
714 'Pure python version used in docs'
715 pools = map(tuple, args) * kwds.get('repeat', 1)
716 result = [[]]
717 for pool in pools:
718 result = [x+[y] for x in result for y in pool]
719 for prod in result:
720 yield tuple(prod)
721
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000722 argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
723 set('abcdefg'), range(11), tuple(range(13))]
724 for i in range(100):
725 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Raymond Hettingerd553d852008-03-04 04:17:08 +0000726 expected_len = prod(map(len, args))
727 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000728 self.assertEqual(list(product(*args)), list(product1(*args)))
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000729 self.assertEqual(list(product(*args)), list(product2(*args)))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000730 args = map(iter, args)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000731 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000732
Benjamin Petersondda91212015-02-01 21:34:07 -0500733 @test_support.bigaddrspacetest
734 def test_product_overflow(self):
Serhiy Storchaka42aa9c02015-02-03 01:34:09 +0200735 with self.assertRaises((OverflowError, MemoryError)):
736 product(*(['ab']*2**5), repeat=2**25)
Benjamin Petersondda91212015-02-01 21:34:07 -0500737
Alex Gaynor97376482011-07-17 16:44:31 -0700738 @test_support.impl_detail("tuple reuse is specific to CPython")
739 def test_product_tuple_reuse(self):
Raymond Hettinger73d79632008-02-23 02:20:41 +0000740 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
741 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000742
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000743 def test_repeat(self):
Raymond Hettinger182edae2009-02-19 02:38:25 +0000744 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
Raymond Hettinger58ad2452014-06-24 21:53:45 -0700745 self.assertEqual(list(repeat(object='a', times=0)), [])
746 self.assertEqual(list(repeat(object='a', times=-1)), [])
747 self.assertEqual(list(repeat(object='a', times=-2)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000748 self.assertEqual(zip(xrange(3),repeat('a')),
749 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000750 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000751 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000752 self.assertEqual(list(repeat('a', 0)), [])
753 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000754 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000755 self.assertRaises(TypeError, repeat, None, 3, 4)
756 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000757 r = repeat(1+0j)
758 self.assertEqual(repr(r), 'repeat((1+0j))')
759 r = repeat(1+0j, 5)
760 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
761 list(r)
762 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000763
Raymond Hettinger58ad2452014-06-24 21:53:45 -0700764 def test_repeat_with_negative_times(self):
765 self.assertEqual(repr(repeat('a', -1)), "repeat('a', 0)")
766 self.assertEqual(repr(repeat('a', -2)), "repeat('a', 0)")
767 self.assertEqual(repr(repeat('a', times=-1)), "repeat('a', 0)")
768 self.assertEqual(repr(repeat('a', times=-2)), "repeat('a', 0)")
769
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000770 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000771 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
772 [0**1, 1**2, 2**3])
773 self.assertEqual(list(imap(None, 'abc', range(5))),
774 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000775 self.assertEqual(list(imap(None, 'abc', count())),
776 [('a',0),('b',1),('c',2)])
777 self.assertEqual(take(2,imap(None, 'abc', count())),
778 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000779 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000780 self.assertRaises(TypeError, imap)
781 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000782 self.assertRaises(TypeError, imap(10, range(5)).next)
783 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
784 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000785
786 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000787 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
788 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000789 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
790 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000791 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger47317092008-01-17 03:02:14 +0000792 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
793 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000794 self.assertRaises(TypeError, starmap)
795 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
796 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
797 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
798 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000799
800 def test_islice(self):
801 for args in [ # islice(args) should agree with range(args)
802 (10, 20, 3),
803 (10, 3, 20),
804 (10, 20),
Cheryl Sabella325191b2018-04-02 01:29:01 -0400805 (10, 10),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000806 (10, 3),
807 (20,)
808 ]:
809 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
810
811 for args, tgtargs in [ # Stop when seqn is exhausted
812 ((10, 110, 3), ((10, 100, 3))),
813 ((10, 110), ((10, 100))),
814 ((110,), (100,))
815 ]:
816 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
817
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000818 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000819 self.assertEqual(list(islice(xrange(10), None)), range(10))
Raymond Hettingerb2594052004-12-05 09:25:51 +0000820 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
821 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000822 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
823 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
824
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000825 # Test number of items consumed SF #1171417
826 it = iter(range(10))
827 self.assertEqual(list(islice(it, 3)), range(3))
828 self.assertEqual(list(it), range(3, 10))
829
Cheryl Sabella325191b2018-04-02 01:29:01 -0400830 it = iter(range(10))
831 self.assertEqual(list(islice(it, 3, 3)), [])
832 self.assertEqual(list(it), range(3, 10))
833
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000834 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000835 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000836 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
837 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
838 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
839 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
840 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000841 self.assertRaises(ValueError, islice, xrange(10), 'a')
842 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
843 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
844 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
845 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +0000846 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000847
Raymond Hettinger061bf7a2010-11-30 03:15:35 +0000848 # Issue #10323: Less islice in a predictable state
849 c = count()
850 self.assertEqual(list(islice(c, 1, 3, 50)), [1])
851 self.assertEqual(next(c), 3)
852
Antoine Pitrou3ec903f2014-04-29 12:13:46 +0200853 # Issue #21321: check source iterator is not referenced
854 # from islice() after the latter has been exhausted
855 it = (x for x in (1, 2))
856 wr = weakref.ref(it)
857 it = islice(it, 1)
858 self.assertIsNotNone(wr())
859 list(it) # exhaust the iterator
Benjamin Petersonec9d5472014-08-24 18:07:28 -0500860 test_support.gc_collect()
Antoine Pitrou3ec903f2014-04-29 12:13:46 +0200861 self.assertIsNone(wr())
862
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000863 def test_takewhile(self):
864 data = [1, 3, 5, 20, 2, 4, 6, 8]
865 underten = lambda x: x<10
866 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000867 self.assertEqual(list(takewhile(underten, [])), [])
868 self.assertRaises(TypeError, takewhile)
869 self.assertRaises(TypeError, takewhile, operator.pow)
870 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
871 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
872 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000873 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
874 self.assertEqual(list(t), [1, 1, 1])
875 self.assertRaises(StopIteration, t.next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000876
877 def test_dropwhile(self):
878 data = [1, 3, 5, 20, 2, 4, 6, 8]
879 underten = lambda x: x<10
880 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000881 self.assertEqual(list(dropwhile(underten, [])), [])
882 self.assertRaises(TypeError, dropwhile)
883 self.assertRaises(TypeError, dropwhile, operator.pow)
884 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
885 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
886 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000887
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000888 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000889 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000890 def irange(n):
891 for i in xrange(n):
892 yield i
893
894 a, b = tee([]) # test empty iterator
895 self.assertEqual(list(a), [])
896 self.assertEqual(list(b), [])
897
898 a, b = tee(irange(n)) # test 100% interleaved
899 self.assertEqual(zip(a,b), zip(range(n),range(n)))
900
901 a, b = tee(irange(n)) # test 0% interleaved
902 self.assertEqual(list(a), range(n))
903 self.assertEqual(list(b), range(n))
904
905 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000906 for i in xrange(100):
907 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000908 del a
909 self.assertEqual(list(b), range(n))
910
911 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000912 for i in xrange(100):
913 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000914 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000915 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000916
917 for j in xrange(5): # test randomly interleaved
918 order = [0]*n + [1]*n
919 random.shuffle(order)
920 lists = ([], [])
921 its = tee(irange(n))
922 for i in order:
923 value = its[i].next()
924 lists[i].append(value)
925 self.assertEqual(lists[0], range(n))
926 self.assertEqual(lists[1], range(n))
927
Raymond Hettingerad983e72003-11-12 14:32:26 +0000928 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000929 self.assertRaises(TypeError, tee)
930 self.assertRaises(TypeError, tee, 3)
931 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000932 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000933
Raymond Hettingerad983e72003-11-12 14:32:26 +0000934 # tee object should be instantiable
935 a, b = tee('abc')
936 c = type(a)('def')
937 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000938
Raymond Hettingerad983e72003-11-12 14:32:26 +0000939 # test long-lagged and multi-way split
940 a, b, c = tee(xrange(2000), 3)
941 for i in xrange(100):
942 self.assertEqual(a.next(), i)
943 self.assertEqual(list(b), range(2000))
944 self.assertEqual([c.next(), c.next()], range(2))
945 self.assertEqual(list(a), range(100,2000))
946 self.assertEqual(list(c), range(2,2000))
947
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000948 # test values of n
949 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Neal Norwitz69e88972006-09-02 02:58:13 +0000950 self.assertRaises(ValueError, tee, [], -1)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000951 for n in xrange(5):
952 result = tee('abc', n)
953 self.assertEqual(type(result), tuple)
954 self.assertEqual(len(result), n)
955 self.assertEqual(map(list, result), [list('abc')]*n)
956
Raymond Hettingerad983e72003-11-12 14:32:26 +0000957 # tee pass-through to copyable iterator
958 a, b = tee('abc')
959 c, d = tee(a)
Benjamin Peterson5c8da862009-06-30 22:57:08 +0000960 self.assertTrue(a is c)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000961
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000962 # test tee_new
963 t1, t2 = tee('abc')
964 tnew = type(t1)
965 self.assertRaises(TypeError, tnew)
966 self.assertRaises(TypeError, tnew, 10)
967 t3 = tnew(t1)
Benjamin Peterson5c8da862009-06-30 22:57:08 +0000968 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000969
Raymond Hettingera9f60922004-10-17 16:40:14 +0000970 # test that tee objects are weak referencable
971 a, b = tee(xrange(10))
Antoine Pitrou3ec903f2014-04-29 12:13:46 +0200972 p = weakref.proxy(a)
Raymond Hettingera9f60922004-10-17 16:40:14 +0000973 self.assertEqual(getattr(p, '__class__'), type(b))
974 del a
975 self.assertRaises(ReferenceError, getattr, p, '__class__')
976
Serhiy Storchakab09ec9b2013-01-25 13:31:05 +0200977 # Issue 13454: Crash when deleting backward iterator from tee()
978 def test_tee_del_backward(self):
Serhiy Storchaka6fef14d2013-01-26 11:51:42 +0200979 forward, backward = tee(repeat(None, 20000000))
Serhiy Storchaka53ea1622015-03-28 20:38:48 +0200980 try:
981 any(forward) # exhaust the iterator
982 del backward
983 except:
984 del forward, backward
985 raise
Serhiy Storchakab09ec9b2013-01-25 13:31:05 +0200986
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000987 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000988 self.assertRaises(StopIteration, izip().next)
989
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000990 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000991 self.assertRaises(StopIteration, f([]).next)
992 self.assertRaises(StopIteration, f(StopNow()).next)
993
994 self.assertRaises(StopIteration, islice([], None).next)
995 self.assertRaises(StopIteration, islice(StopNow(), None).next)
996
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000997 p, q = tee([])
998 self.assertRaises(StopIteration, p.next)
999 self.assertRaises(StopIteration, q.next)
1000 p, q = tee(StopNow())
1001 self.assertRaises(StopIteration, p.next)
1002 self.assertRaises(StopIteration, q.next)
1003
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001004 self.assertRaises(StopIteration, repeat(None, 0).next)
1005
1006 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
1007 self.assertRaises(StopIteration, f(lambda x:x, []).next)
1008 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
1009
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001010class TestExamples(unittest.TestCase):
1011
1012 def test_chain(self):
1013 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
1014
1015 def test_chain_from_iterable(self):
1016 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
1017
1018 def test_combinations(self):
1019 self.assertEqual(list(combinations('ABCD', 2)),
1020 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
1021 self.assertEqual(list(combinations(range(4), 3)),
1022 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
1023
Raymond Hettingerd081abc2009-01-27 02:58:49 +00001024 def test_combinations_with_replacement(self):
1025 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
1026 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
1027
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001028 def test_compress(self):
1029 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
1030
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001031 def test_count(self):
1032 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
1033
1034 def test_cycle(self):
1035 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
1036
1037 def test_dropwhile(self):
1038 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
1039
1040 def test_groupby(self):
1041 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
1042 list('ABCDAB'))
1043 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
1044 [list('AAAA'), list('BBB'), list('CC'), list('D')])
1045
1046 def test_ifilter(self):
1047 self.assertEqual(list(ifilter(lambda x: x%2, range(10))), [1,3,5,7,9])
1048
1049 def test_ifilterfalse(self):
1050 self.assertEqual(list(ifilterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
1051
1052 def test_imap(self):
1053 self.assertEqual(list(imap(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
1054
1055 def test_islice(self):
1056 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
1057 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
1058 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
1059 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1060
1061 def test_izip(self):
1062 self.assertEqual(list(izip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
1063
1064 def test_izip_longest(self):
1065 self.assertEqual(list(izip_longest('ABCD', 'xy', fillvalue='-')),
1066 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
1067
1068 def test_permutations(self):
1069 self.assertEqual(list(permutations('ABCD', 2)),
1070 map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split()))
1071 self.assertEqual(list(permutations(range(3))),
1072 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
1073
1074 def test_product(self):
1075 self.assertEqual(list(product('ABCD', 'xy')),
1076 map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split()))
1077 self.assertEqual(list(product(range(2), repeat=3)),
1078 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
1079 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
1080
1081 def test_repeat(self):
1082 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
1083
1084 def test_stapmap(self):
1085 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
1086 [32, 9, 1000])
1087
1088 def test_takewhile(self):
1089 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
1090
1091
Cheryl Sabella325191b2018-04-02 01:29:01 -04001092class TestPurePythonRoughEquivalents(unittest.TestCase):
1093
1094 @staticmethod
1095 def islice(iterable, *args):
1096 s = slice(*args)
1097 start, stop, step = s.start or 0, s.stop or sys.maxint, s.step or 1
1098 it = iter(xrange(start, stop, step))
1099 try:
1100 nexti = next(it)
1101 except StopIteration:
1102 # Consume *iterable* up to the *start* position.
1103 for i, element in izip(xrange(start), iterable):
1104 pass
1105 return
1106 try:
1107 for i, element in enumerate(iterable):
1108 if i == nexti:
1109 yield element
1110 nexti = next(it)
1111 except StopIteration:
1112 # Consume to *stop*.
1113 for i, element in izip(xrange(i + 1, stop), iterable):
1114 pass
1115
1116 def test_islice_recipe(self):
1117 self.assertEqual(list(self.islice('ABCDEFG', 2)), list('AB'))
1118 self.assertEqual(list(self.islice('ABCDEFG', 2, 4)), list('CD'))
1119 self.assertEqual(list(self.islice('ABCDEFG', 2, None)), list('CDEFG'))
1120 self.assertEqual(list(self.islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1121 # Test items consumed.
1122 it = iter(xrange(10))
1123 self.assertEqual(list(self.islice(it, 3)), range(3))
1124 self.assertEqual(list(it), range(3, 10))
1125 it = iter(xrange(10))
1126 self.assertEqual(list(self.islice(it, 3, 3)), [])
1127 self.assertEqual(list(it), range(3, 10))
1128 # Test that slice finishes in predictable state.
1129 c = count()
1130 self.assertEqual(list(self.islice(c, 1, 3, 50)), [1])
1131 self.assertEqual(next(c), 3)
1132
1133
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001134class TestGC(unittest.TestCase):
1135
1136 def makecycle(self, iterator, container):
1137 container.append(iterator)
1138 iterator.next()
1139 del container, iterator
1140
1141 def test_chain(self):
1142 a = []
1143 self.makecycle(chain(a), a)
1144
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001145 def test_chain_from_iterable(self):
1146 a = []
1147 self.makecycle(chain.from_iterable([a]), a)
1148
1149 def test_combinations(self):
1150 a = []
1151 self.makecycle(combinations([1,2,a,3], 3), a)
1152
Raymond Hettingerd081abc2009-01-27 02:58:49 +00001153 def test_combinations_with_replacement(self):
1154 a = []
1155 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1156
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001157 def test_compress(self):
1158 a = []
1159 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1160
Raymond Hettingerb21d8102009-02-16 20:39:12 +00001161 def test_count(self):
1162 a = []
1163 Int = type('Int', (int,), dict(x=a))
1164 self.makecycle(count(Int(0), Int(1)), a)
1165
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001166 def test_cycle(self):
1167 a = []
1168 self.makecycle(cycle([a]*2), a)
1169
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001170 def test_dropwhile(self):
1171 a = []
1172 self.makecycle(dropwhile(bool, [0, a, a]), a)
1173
1174 def test_groupby(self):
1175 a = []
1176 self.makecycle(groupby([a]*2, lambda x:x), a)
1177
Raymond Hettingera1ca94a2008-03-06 22:51:36 +00001178 def test_issue2246(self):
1179 # Issue 2246 -- the _grouper iterator was not included in GC
1180 n = 10
1181 keyfunc = lambda x: x
1182 for i, j in groupby(xrange(n), key=keyfunc):
1183 keyfunc.__dict__.setdefault('x',[]).append(j)
1184
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001185 def test_ifilter(self):
1186 a = []
1187 self.makecycle(ifilter(lambda x:True, [a]*2), a)
1188
1189 def test_ifilterfalse(self):
1190 a = []
1191 self.makecycle(ifilterfalse(lambda x:False, a), a)
1192
1193 def test_izip(self):
1194 a = []
1195 self.makecycle(izip([a]*2, [a]*3), a)
1196
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001197 def test_izip_longest(self):
1198 a = []
1199 self.makecycle(izip_longest([a]*2, [a]*3), a)
1200 b = [a, None]
1201 self.makecycle(izip_longest([a]*2, [a]*3, fillvalue=b), a)
1202
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001203 def test_imap(self):
1204 a = []
1205 self.makecycle(imap(lambda x:x, [a]*2), a)
1206
1207 def test_islice(self):
1208 a = []
1209 self.makecycle(islice([a]*2, None), a)
1210
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001211 def test_permutations(self):
1212 a = []
1213 self.makecycle(permutations([1,2,a,3], 3), a)
1214
1215 def test_product(self):
1216 a = []
1217 self.makecycle(product([1,2,a,3], repeat=3), a)
1218
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001219 def test_repeat(self):
1220 a = []
1221 self.makecycle(repeat(a), a)
1222
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001223 def test_starmap(self):
1224 a = []
1225 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1226
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001227 def test_takewhile(self):
1228 a = []
1229 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1230
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001231def R(seqn):
1232 'Regular generator'
1233 for i in seqn:
1234 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001235
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001236class G:
1237 'Sequence using __getitem__'
1238 def __init__(self, seqn):
1239 self.seqn = seqn
1240 def __getitem__(self, i):
1241 return self.seqn[i]
1242
1243class I:
1244 'Sequence using iterator protocol'
1245 def __init__(self, seqn):
1246 self.seqn = seqn
1247 self.i = 0
1248 def __iter__(self):
1249 return self
1250 def next(self):
1251 if self.i >= len(self.seqn): raise StopIteration
1252 v = self.seqn[self.i]
1253 self.i += 1
1254 return v
1255
1256class Ig:
1257 'Sequence using iterator protocol defined with a generator'
1258 def __init__(self, seqn):
1259 self.seqn = seqn
1260 self.i = 0
1261 def __iter__(self):
1262 for val in self.seqn:
1263 yield val
1264
1265class X:
1266 'Missing __getitem__ and __iter__'
1267 def __init__(self, seqn):
1268 self.seqn = seqn
1269 self.i = 0
1270 def next(self):
1271 if self.i >= len(self.seqn): raise StopIteration
1272 v = self.seqn[self.i]
1273 self.i += 1
1274 return v
1275
1276class N:
1277 'Iterator missing next()'
1278 def __init__(self, seqn):
1279 self.seqn = seqn
1280 self.i = 0
1281 def __iter__(self):
1282 return self
1283
1284class E:
1285 'Test propagation of exceptions'
1286 def __init__(self, seqn):
1287 self.seqn = seqn
1288 self.i = 0
1289 def __iter__(self):
1290 return self
1291 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001292 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001293
1294class S:
1295 'Test immediate stop'
1296 def __init__(self, seqn):
1297 pass
1298 def __iter__(self):
1299 return self
1300 def next(self):
1301 raise StopIteration
1302
1303def L(seqn):
1304 'Test multiple tiers of iterators'
1305 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
1306
1307
1308class TestVariousIteratorArgs(unittest.TestCase):
1309
1310 def test_chain(self):
1311 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1312 for g in (G, I, Ig, S, L, R):
1313 self.assertEqual(list(chain(g(s))), list(g(s)))
1314 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001315 self.assertRaises(TypeError, list, chain(X(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001316 self.assertRaises(TypeError, list, chain(N(s)))
1317 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1318
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001319 def test_compress(self):
1320 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1321 n = len(s)
1322 for g in (G, I, Ig, S, L, R):
1323 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1324 self.assertRaises(TypeError, compress, X(s), repeat(1))
1325 self.assertRaises(TypeError, list, compress(N(s), repeat(1)))
1326 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1327
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001328 def test_product(self):
1329 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1330 self.assertRaises(TypeError, product, X(s))
1331 self.assertRaises(TypeError, product, N(s))
1332 self.assertRaises(ZeroDivisionError, product, E(s))
1333
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001334 def test_cycle(self):
1335 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1336 for g in (G, I, Ig, S, L, R):
1337 tgtlen = len(s) * 3
1338 expected = list(g(s))*3
1339 actual = list(islice(cycle(g(s)), tgtlen))
1340 self.assertEqual(actual, expected)
1341 self.assertRaises(TypeError, cycle, X(s))
1342 self.assertRaises(TypeError, list, cycle(N(s)))
1343 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1344
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001345 def test_groupby(self):
1346 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1347 for g in (G, I, Ig, S, L, R):
1348 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1349 self.assertRaises(TypeError, groupby, X(s))
1350 self.assertRaises(TypeError, list, groupby(N(s)))
1351 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1352
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001353 def test_ifilter(self):
1354 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1355 for g in (G, I, Ig, S, L, R):
1356 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
1357 self.assertRaises(TypeError, ifilter, isEven, X(s))
1358 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
1359 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
1360
1361 def test_ifilterfalse(self):
1362 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1363 for g in (G, I, Ig, S, L, R):
1364 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
1365 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
1366 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
1367 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
1368
1369 def test_izip(self):
1370 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1371 for g in (G, I, Ig, S, L, R):
1372 self.assertEqual(list(izip(g(s))), zip(g(s)))
1373 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
1374 self.assertRaises(TypeError, izip, X(s))
1375 self.assertRaises(TypeError, list, izip(N(s)))
1376 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
1377
Raymond Hettingerd36862c2007-02-21 05:20:38 +00001378 def test_iziplongest(self):
1379 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1380 for g in (G, I, Ig, S, L, R):
1381 self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
1382 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
1383 self.assertRaises(TypeError, izip_longest, X(s))
1384 self.assertRaises(TypeError, list, izip_longest(N(s)))
1385 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
1386
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001387 def test_imap(self):
1388 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1389 for g in (G, I, Ig, S, L, R):
1390 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
1391 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
1392 self.assertRaises(TypeError, imap, onearg, X(s))
1393 self.assertRaises(TypeError, list, imap(onearg, N(s)))
1394 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
1395
1396 def test_islice(self):
1397 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1398 for g in (G, I, Ig, S, L, R):
1399 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1400 self.assertRaises(TypeError, islice, X(s), 10)
1401 self.assertRaises(TypeError, list, islice(N(s), 10))
1402 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1403
1404 def test_starmap(self):
1405 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1406 for g in (G, I, Ig, S, L, R):
1407 ss = zip(s, s)
1408 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
1409 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
1410 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
1411 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1412
1413 def test_takewhile(self):
1414 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1415 for g in (G, I, Ig, S, L, R):
1416 tgt = []
1417 for elem in g(s):
1418 if not isEven(elem): break
1419 tgt.append(elem)
1420 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1421 self.assertRaises(TypeError, takewhile, isEven, X(s))
1422 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
1423 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1424
1425 def test_dropwhile(self):
1426 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1427 for g in (G, I, Ig, S, L, R):
1428 tgt = []
1429 for elem in g(s):
1430 if not tgt and isOdd(elem): continue
1431 tgt.append(elem)
1432 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1433 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
1434 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
1435 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1436
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001437 def test_tee(self):
1438 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1439 for g in (G, I, Ig, S, L, R):
1440 it1, it2 = tee(g(s))
1441 self.assertEqual(list(it1), list(g(s)))
1442 self.assertEqual(list(it2), list(g(s)))
1443 self.assertRaises(TypeError, tee, X(s))
1444 self.assertRaises(TypeError, list, tee(N(s))[0])
1445 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1446
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001447class LengthTransparency(unittest.TestCase):
1448
1449 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001450 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001451 self.assertEqual(len(repeat(None, 50)), 50)
1452 self.assertRaises(TypeError, len, repeat(None))
1453
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001454class RegressionTests(unittest.TestCase):
1455
1456 def test_sf_793826(self):
1457 # Fix Armin Rigo's successful efforts to wreak havoc
1458
1459 def mutatingtuple(tuple1, f, tuple2):
1460 # this builds a tuple t which is a copy of tuple1,
1461 # then calls f(t), then mutates t to be equal to tuple2
1462 # (needs len(tuple1) == len(tuple2)).
1463 def g(value, first=[1]):
1464 if first:
1465 del first[:]
1466 f(z.next())
1467 return value
1468 items = list(tuple2)
1469 items[1:1] = list(tuple1)
1470 gen = imap(g, items)
1471 z = izip(*[gen]*len(tuple1))
1472 z.next()
1473
1474 def f(t):
1475 global T
1476 T = t
1477 first[:] = list(T)
1478
1479 first = []
1480 mutatingtuple((1,2,3), f, (4,5,6))
1481 second = list(T)
1482 self.assertEqual(first, second)
1483
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001484
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001485 def test_sf_950057(self):
1486 # Make sure that chain() and cycle() catch exceptions immediately
1487 # rather than when shifting between input sources
1488
1489 def gen1():
1490 hist.append(0)
1491 yield 1
1492 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001493 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001494 hist.append(2)
1495
1496 def gen2(x):
1497 hist.append(3)
1498 yield 2
1499 hist.append(4)
1500 if x:
1501 raise StopIteration
1502
1503 hist = []
1504 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1505 self.assertEqual(hist, [0,1])
1506
1507 hist = []
1508 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1509 self.assertEqual(hist, [0,1])
1510
1511 hist = []
1512 self.assertRaises(AssertionError, list, cycle(gen1()))
1513 self.assertEqual(hist, [0,1])
1514
T. Woutersd694a062017-03-30 12:49:22 -07001515 def test_long_chain_of_empty_iterables(self):
1516 # Make sure itertools.chain doesn't run into recursion limits when
1517 # dealing with long chains of empty iterables. Even with a high
1518 # number this would probably only fail in Py_DEBUG mode.
1519 it = chain.from_iterable(() for unused in xrange(10000000))
1520 with self.assertRaises(StopIteration):
1521 next(it)
1522
Serhiy Storchakad0b9dc32017-09-26 23:15:36 +03001523 def test_issue30347_1(self):
1524 def f(n):
1525 if n == 5:
1526 list(b)
1527 return n != 6
1528 for (k, b) in groupby(range(10), f):
1529 list(b) # shouldn't crash
1530
1531 def test_issue30347_2(self):
1532 class K(object):
1533 i = 0
1534 def __init__(self, v):
1535 pass
1536 def __eq__(self, other):
1537 K.i += 1
1538 if K.i == 1:
1539 next(g, None)
1540 return True
Victor Stinner2ce1ef52017-11-08 14:45:55 -08001541 def __hash__(self):
1542 return 1
Serhiy Storchakad0b9dc32017-09-26 23:15:36 +03001543 g = next(groupby(range(10), K))[1]
1544 for j in range(2):
1545 next(g, None) # shouldn't crash
1546
1547
Georg Brandlb84c1372007-01-21 10:28:43 +00001548class SubclassWithKwargsTest(unittest.TestCase):
1549 def test_keywords_in_subclass(self):
1550 # count is not subclassable...
1551 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001552 starmap, islice, takewhile, dropwhile, cycle, compress):
Georg Brandlb84c1372007-01-21 10:28:43 +00001553 class Subclass(cls):
1554 def __init__(self, newarg=None, *args):
1555 cls.__init__(self, *args)
1556 try:
1557 Subclass(newarg=1)
1558 except TypeError, err:
1559 # we expect type errors because of wrong argument count
Ezio Melottiaa980582010-01-23 23:04:36 +00001560 self.assertNotIn("does not take keyword arguments", err.args[0])
Georg Brandlb84c1372007-01-21 10:28:43 +00001561
1562
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001563libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001564
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001565
1566>>> amounts = [120.15, 764.05, 823.14]
1567>>> for checknum, amount in izip(count(1200), amounts):
1568... print 'Check %d is for $%.2f' % (checknum, amount)
1569...
1570Check 1200 is for $120.15
1571Check 1201 is for $764.05
1572Check 1202 is for $823.14
1573
1574>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001575>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1576... print cube
1577...
15781
15798
158027
1581
1582>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001583>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001584... print name.title()
1585...
1586Alex
1587Laura
1588Martin
1589Walter
1590Samuele
1591
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001592>>> from operator import itemgetter
1593>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Armin Rigoa3f09272006-05-28 19:13:17 +00001594>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001595>>> for k, g in groupby(di, itemgetter(1)):
1596... print k, map(itemgetter(0), g)
1597...
15981 ['a', 'c', 'e']
15992 ['b', 'd', 'f']
16003 ['g']
1601
Raymond Hettinger734fb572004-01-20 20:04:40 +00001602# Find runs of consecutive numbers using groupby. The key to the solution
1603# is differencing with a range so that consecutive numbers all appear in
1604# same group.
1605>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Ezio Melottidde5b942010-02-03 05:37:26 +00001606>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Raymond Hettinger734fb572004-01-20 20:04:40 +00001607... print map(operator.itemgetter(1), g)
1608...
1609[1]
1610[4, 5, 6]
1611[10]
1612[15, 16, 17, 18]
1613[22]
1614[25, 26, 27, 28]
1615
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001616>>> def take(n, iterable):
1617... "Return first n items of the iterable as a list"
1618... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001619
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001620>>> def enumerate(iterable, start=0):
1621... return izip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001622
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001623>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001624... "Return function(0), function(1), ..."
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001625... return imap(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001626
Cheryl Sabella325191b2018-04-02 01:29:01 -04001627>>> import collections
1628>>> def consume(iterator, n=None):
1629... "Advance the iterator n-steps ahead. If n is None, consume entirely."
1630... # Use functions that consume iterators at C speed.
1631... if n is None:
1632... # feed the entire iterator into a zero-length deque
1633... collections.deque(iterator, maxlen=0)
1634... else:
1635... # advance to the empty slice starting at position n
1636... next(islice(iterator, n, n), None)
1637
Raymond Hettingerf9bce832009-02-19 05:34:35 +00001638>>> def nth(iterable, n, default=None):
1639... "Returns the nth item or a default value"
1640... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001641
Raymond Hettinger2ec5fd52016-03-06 18:06:29 -08001642>>> def all_equal(iterable):
1643... "Returns True if all the elements are equal to each other"
1644... g = groupby(iterable)
1645... return next(g, True) and not next(g, False)
1646
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001647>>> def quantify(iterable, pred=bool):
1648... "Count how many times the predicate is true"
1649... return sum(imap(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001650
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001651>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001652... "Returns the sequence elements and then returns None indefinitely"
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001653... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001654
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001655>>> def ncycles(iterable, n):
Ezio Melottic2077b02011-03-16 12:34:31 +02001656... "Returns the sequence elements n times"
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001657... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001658
1659>>> def dotproduct(vec1, vec2):
1660... return sum(imap(operator.mul, vec1, vec2))
1661
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001662>>> def flatten(listOfLists):
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001663... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001664
1665>>> def repeatfunc(func, times=None, *args):
1666... "Repeat calls to func with specified arguments."
1667... " Example: repeatfunc(random.random)"
1668... if times is None:
1669... return starmap(func, repeat(args))
1670... else:
1671... return starmap(func, repeat(args, times))
1672
Raymond Hettingerd591f662003-10-26 15:34:50 +00001673>>> def pairwise(iterable):
1674... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1675... a, b = tee(iterable)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001676... for elem in b:
1677... break
Raymond Hettingerad983e72003-11-12 14:32:26 +00001678... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001679
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001680>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001681... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001682... args = [iter(iterable)] * n
Raymond Hettingerf080e6d2008-07-31 01:19:50 +00001683... return izip_longest(fillvalue=fillvalue, *args)
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001684
1685>>> def roundrobin(*iterables):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001686... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001687... # Recipe credited to George Sakkis
1688... pending = len(iterables)
1689... nexts = cycle(iter(it).next for it in iterables)
1690... while pending:
1691... try:
1692... for next in nexts:
1693... yield next()
1694... except StopIteration:
1695... pending -= 1
1696... nexts = cycle(islice(nexts, pending))
1697
1698>>> def powerset(iterable):
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001699... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1700... s = list(iterable)
1701... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001702
Raymond Hettinger44e15812009-01-02 21:26:45 +00001703>>> def unique_everseen(iterable, key=None):
1704... "List unique elements, preserving order. Remember all elements ever seen."
1705... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1706... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1707... seen = set()
1708... seen_add = seen.add
1709... if key is None:
1710... for element in iterable:
1711... if element not in seen:
1712... seen_add(element)
1713... yield element
1714... else:
1715... for element in iterable:
1716... k = key(element)
1717... if k not in seen:
1718... seen_add(k)
1719... yield element
1720
1721>>> def unique_justseen(iterable, key=None):
1722... "List unique elements, preserving order. Remember only the element just seen."
1723... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1724... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1725... return imap(next, imap(itemgetter(1), groupby(iterable, key)))
1726
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001727This is not part of the examples but it tests to make sure the definitions
1728perform as purported.
1729
Raymond Hettingera098b332003-09-08 23:58:40 +00001730>>> take(10, count())
1731[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1732
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001733>>> list(enumerate('abc'))
1734[(0, 'a'), (1, 'b'), (2, 'c')]
1735
1736>>> list(islice(tabulate(lambda x: 2*x), 4))
1737[0, 2, 4, 6]
1738
Cheryl Sabella325191b2018-04-02 01:29:01 -04001739>>> it = iter(xrange(10))
1740>>> consume(it, 3)
1741>>> next(it)
17423
1743>>> consume(it)
1744>>> next(it, 'Done')
1745'Done'
1746
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001747>>> nth('abcde', 3)
Raymond Hettingerd507afd2009-02-04 10:52:32 +00001748'd'
1749
1750>>> nth('abcde', 9) is None
1751True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001752
Raymond Hettinger2ec5fd52016-03-06 18:06:29 -08001753>>> [all_equal(s) for s in ('', 'A', 'AAAA', 'AAAB', 'AAABA')]
1754[True, True, True, False, False]
1755
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001756>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000175750
1758
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001759>>> a = [[1, 2, 3], [4, 5, 6]]
1760>>> flatten(a)
1761[1, 2, 3, 4, 5, 6]
1762
1763>>> list(repeatfunc(pow, 5, 2, 3))
1764[8, 8, 8, 8, 8]
1765
1766>>> import random
1767>>> take(5, imap(int, repeatfunc(random.random)))
1768[0, 0, 0, 0, 0]
1769
Raymond Hettingerd591f662003-10-26 15:34:50 +00001770>>> list(pairwise('abcd'))
1771[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001772
Raymond Hettingerd591f662003-10-26 15:34:50 +00001773>>> list(pairwise([]))
1774[]
1775
1776>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001777[]
1778
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001779>>> list(islice(padnone('abc'), 0, 6))
1780['a', 'b', 'c', None, None, None]
1781
1782>>> list(ncycles('abc', 3))
1783['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1784
1785>>> dotproduct([1,2,3], [4,5,6])
178632
1787
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001788>>> list(grouper(3, 'abcdefg', 'x'))
1789[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1790
1791>>> list(roundrobin('abc', 'd', 'ef'))
1792['a', 'd', 'e', 'b', 'f', 'c']
1793
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001794>>> list(powerset([1,2,3]))
1795[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001796
Raymond Hettinger560f9a82009-01-27 13:26:35 +00001797>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1798True
1799
1800>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1801True
1802
Raymond Hettinger44e15812009-01-02 21:26:45 +00001803>>> list(unique_everseen('AAAABBBCCDAABBB'))
1804['A', 'B', 'C', 'D']
1805
1806>>> list(unique_everseen('ABBCcAD', str.lower))
1807['A', 'B', 'C', 'D']
1808
1809>>> list(unique_justseen('AAAABBBCCDAABBB'))
1810['A', 'B', 'C', 'D', 'A', 'B']
1811
1812>>> list(unique_justseen('ABBCcAD', str.lower))
1813['A', 'B', 'C', 'A', 'D']
1814
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001815"""
1816
1817__test__ = {'libreftest' : libreftest}
1818
1819def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001820 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Georg Brandlb84c1372007-01-21 10:28:43 +00001821 RegressionTests, LengthTransparency,
Cheryl Sabella325191b2018-04-02 01:29:01 -04001822 SubclassWithKwargsTest, TestExamples,
1823 TestPurePythonRoughEquivalents)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001824 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001825
1826 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001827 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001828 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001829 counts = [None] * 5
1830 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001831 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001832 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001833 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001834 print counts
1835
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001836 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001837 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001838
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001839if __name__ == "__main__":
1840 test_main(verbose=True)