blob: 31d10dc5c84086f39c4fcf6dfe77603df4b88ed8 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001import unittest
2from test import test_support
3from itertools import *
Raymond Hettingera9f60922004-10-17 16:40:14 +00004from weakref import proxy
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
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +000010maxsize = test_support.MAX_Py_ssize_t
11minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000012
13def onearg(x):
14 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000015 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000016
17def errfunc(*args):
18 'Test function that raises an error'
19 raise ValueError
20
21def gen3():
22 'Non-restartable source sequence'
23 for i in (0, 1, 2):
24 yield i
25
26def isEven(x):
27 'Test predicate'
28 return x%2==0
29
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000030def isOdd(x):
31 'Test predicate'
32 return x%2==1
33
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000034class StopNow:
35 'Class emulating an empty iterable.'
36 def __iter__(self):
37 return self
38 def next(self):
39 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000040
Raymond Hettinger02420702003-06-29 20:36:23 +000041def take(n, seq):
42 'Convenience function for partially consuming a long of infinite iterable'
43 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000044
Raymond Hettingerd553d852008-03-04 04:17:08 +000045def prod(iterable):
46 return reduce(operator.mul, iterable, 1)
47
Raymond Hettinger93e804d2008-02-26 23:40:50 +000048def fact(n):
49 'Factorial'
Raymond Hettingerd553d852008-03-04 04:17:08 +000050 return prod(range(1, n+1))
51
Raymond Hettinger96ef8112003-02-01 00:10:11 +000052class TestBasicOps(unittest.TestCase):
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000053 def test_chain(self):
Raymond Hettingerad47fa12008-03-06 20:52:01 +000054
55 def chain2(*iterables):
56 'Pure python version in the docs'
57 for it in iterables:
58 for element in it:
59 yield element
60
61 for c in (chain, chain2):
62 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
63 self.assertEqual(list(c('abc')), list('abc'))
64 self.assertEqual(list(c('')), [])
65 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
66 self.assertRaises(TypeError, list,c(2, 3))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000067
Raymond Hettingerb4cbc982008-02-28 22:46:41 +000068 def test_chain_from_iterable(self):
69 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
70 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
71 self.assertEqual(list(chain.from_iterable([''])), [])
72 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
73 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
74
Raymond Hettinger93e804d2008-02-26 23:40:50 +000075 def test_combinations(self):
Raymond Hettinger5b913e32009-01-08 06:39:04 +000076 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Raymond Hettinger93e804d2008-02-26 23:40:50 +000077 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Raymond Hettingerd553d852008-03-04 04:17:08 +000078 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Raymond Hettinger93e804d2008-02-26 23:40:50 +000079 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Raymond Hettinger5b913e32009-01-08 06:39:04 +000080 self.assertEqual(list(combinations('abc', 32)), []) # r > n
Raymond Hettinger93e804d2008-02-26 23:40:50 +000081 self.assertEqual(list(combinations(range(4), 3)),
82 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
Raymond Hettingerd553d852008-03-04 04:17:08 +000083
84 def combinations1(iterable, r):
85 'Pure python version shown in the docs'
86 pool = tuple(iterable)
87 n = len(pool)
Raymond Hettinger5b913e32009-01-08 06:39:04 +000088 if r > n:
89 return
Raymond Hettingerd553d852008-03-04 04:17:08 +000090 indices = range(r)
91 yield tuple(pool[i] for i in indices)
92 while 1:
93 for i in reversed(range(r)):
94 if indices[i] != i + n - r:
95 break
96 else:
97 return
98 indices[i] += 1
99 for j in range(i+1, r):
100 indices[j] = indices[j-1] + 1
101 yield tuple(pool[i] for i in indices)
102
103 def combinations2(iterable, r):
104 'Pure python version shown in the docs'
105 pool = tuple(iterable)
106 n = len(pool)
107 for indices in permutations(range(n), r):
108 if sorted(indices) == list(indices):
109 yield tuple(pool[i] for i in indices)
110
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000111 def combinations3(iterable, r):
112 'Pure python version from cwr()'
113 pool = tuple(iterable)
114 n = len(pool)
115 for indices in combinations_with_replacement(range(n), r):
116 if len(set(indices)) == r:
117 yield tuple(pool[i] for i in indices)
118
Raymond Hettingerd553d852008-03-04 04:17:08 +0000119 for n in range(7):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000120 values = [5*x-12 for x in range(n)]
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000121 for r in range(n+2):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000122 result = list(combinations(values, r))
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000123 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 +0000124 self.assertEqual(len(result), len(set(result))) # no repeats
125 self.assertEqual(result, sorted(result)) # lexicographic order
126 for c in result:
127 self.assertEqual(len(c), r) # r-length combinations
128 self.assertEqual(len(set(c)), r) # no duplicate elements
129 self.assertEqual(list(c), sorted(c)) # keep original ordering
130 self.assert_(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000131 self.assertEqual(list(c),
132 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Raymond Hettingerd553d852008-03-04 04:17:08 +0000133 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000134 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000135 self.assertEqual(result, list(combinations3(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000136
137 # Test implementation detail: tuple re-use
138 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
139 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
140
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000141 def test_combinations_with_replacement(self):
142 cwr = combinations_with_replacement
143 self.assertRaises(TypeError, cwr, 'abc') # missing r argument
144 self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
145 self.assertRaises(TypeError, cwr, None) # pool is not iterable
146 self.assertRaises(ValueError, cwr, 'abc', -2) # r is negative
147 self.assertEqual(list(cwr('ABC', 2)),
148 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
149
150 def cwr1(iterable, r):
151 'Pure python version shown in the docs'
152 # number items returned: (n+r-1)! / r! / (n-1)! when n>0
153 pool = tuple(iterable)
154 n = len(pool)
155 if not n and r:
156 return
157 indices = [0] * r
158 yield tuple(pool[i] for i in indices)
159 while 1:
160 for i in reversed(range(r)):
161 if indices[i] != n - 1:
162 break
163 else:
164 return
165 indices[i:] = [indices[i] + 1] * (r - i)
166 yield tuple(pool[i] for i in indices)
167
168 def cwr2(iterable, r):
169 'Pure python version shown in the docs'
170 pool = tuple(iterable)
171 n = len(pool)
172 for indices in product(range(n), repeat=r):
173 if sorted(indices) == list(indices):
174 yield tuple(pool[i] for i in indices)
175
176 def numcombs(n, r):
177 if not n:
178 return 0 if r else 1
179 return fact(n+r-1) / fact(r)/ fact(n-1)
180
181 for n in range(7):
182 values = [5*x-12 for x in range(n)]
183 for r in range(n+2):
184 result = list(cwr(values, r))
185
186 self.assertEqual(len(result), numcombs(n, r)) # right number of combs
187 self.assertEqual(len(result), len(set(result))) # no repeats
188 self.assertEqual(result, sorted(result)) # lexicographic order
189
190 regular_combs = list(combinations(values, r)) # compare to combs without replacement
191 if n == 0 or r <= 1:
192 self.assertEquals(result, regular_combs) # cases that should be identical
193 else:
194 self.assert_(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
195
196 for c in result:
197 self.assertEqual(len(c), r) # r-length combinations
198 noruns = [k for k,v in groupby(c)] # combo without consecutive repeats
199 self.assertEqual(len(noruns), len(set(noruns))) # no repeats other than consecutive
200 self.assertEqual(list(c), sorted(c)) # keep original ordering
201 self.assert_(all(e in values for e in c)) # elements taken from input iterable
202 self.assertEqual(noruns,
203 [e for e in values if e in c]) # comb is a subsequence of the input iterable
204 self.assertEqual(result, list(cwr1(values, r))) # matches first pure python version
205 self.assertEqual(result, list(cwr2(values, r))) # matches second pure python version
206
207 # Test implementation detail: tuple re-use
208 self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
209 self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
210
Raymond Hettingerd553d852008-03-04 04:17:08 +0000211 def test_permutations(self):
212 self.assertRaises(TypeError, permutations) # too few arguments
213 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000214 self.assertRaises(TypeError, permutations, None) # pool is not iterable
215 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000216 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000217 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Raymond Hettingerd553d852008-03-04 04:17:08 +0000218 self.assertEqual(list(permutations(range(3), 2)),
219 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
220
221 def permutations1(iterable, r=None):
222 'Pure python version shown in the docs'
223 pool = tuple(iterable)
224 n = len(pool)
225 r = n if r is None else r
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000226 if r > n:
227 return
Raymond Hettingerd553d852008-03-04 04:17:08 +0000228 indices = range(n)
Raymond Hettingere70bb8d2008-03-23 00:55:46 +0000229 cycles = range(n, n-r, -1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000230 yield tuple(pool[i] for i in indices[:r])
231 while n:
232 for i in reversed(range(r)):
233 cycles[i] -= 1
234 if cycles[i] == 0:
235 indices[i:] = indices[i+1:] + indices[i:i+1]
236 cycles[i] = n - i
237 else:
238 j = cycles[i]
239 indices[i], indices[-j] = indices[-j], indices[i]
240 yield tuple(pool[i] for i in indices[:r])
241 break
242 else:
243 return
244
245 def permutations2(iterable, r=None):
246 'Pure python version shown in the docs'
247 pool = tuple(iterable)
248 n = len(pool)
249 r = n if r is None else r
250 for indices in product(range(n), repeat=r):
251 if len(set(indices)) == r:
252 yield tuple(pool[i] for i in indices)
253
254 for n in range(7):
255 values = [5*x-12 for x in range(n)]
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000256 for r in range(n+2):
Raymond Hettingerd553d852008-03-04 04:17:08 +0000257 result = list(permutations(values, r))
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000258 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 +0000259 self.assertEqual(len(result), len(set(result))) # no repeats
260 self.assertEqual(result, sorted(result)) # lexicographic order
261 for p in result:
262 self.assertEqual(len(p), r) # r-length permutations
263 self.assertEqual(len(set(p)), r) # no duplicate elements
264 self.assert_(all(e in values for e in p)) # elements taken from input iterable
265 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000266 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000267 if r == n:
268 self.assertEqual(result, list(permutations(values, None))) # test r as None
269 self.assertEqual(result, list(permutations(values))) # test default r
270
271 # Test implementation detail: tuple re-use
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000272 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000273 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000274
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000275 def test_combinatorics(self):
276 # Test relationships between product(), permutations(),
277 # combinations() and combinations_with_replacement().
278
Raymond Hettinger2f6c2e02009-01-27 10:36:14 +0000279 for n in range(6):
280 s = 'ABCDEFG'[:n]
281 for r in range(8):
282 prod = list(product(s, repeat=r))
283 cwr = list(combinations_with_replacement(s, r))
284 perm = list(permutations(s, r))
285 comb = list(combinations(s, r))
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000286
Raymond Hettinger2f6c2e02009-01-27 10:36:14 +0000287 # Check size
288 self.assertEquals(len(prod), n**r)
289 self.assertEquals(len(cwr), (fact(n+r-1) / fact(r)/ fact(n-1)) if n else (not r))
290 self.assertEquals(len(perm), 0 if r>n else fact(n) / fact(n-r))
291 self.assertEquals(len(comb), 0 if r>n else fact(n) / fact(r) / fact(n-r))
292
293 # Check lexicographic order without repeated tuples
294 self.assertEquals(prod, sorted(set(prod)))
295 self.assertEquals(cwr, sorted(set(cwr)))
296 self.assertEquals(perm, sorted(set(perm)))
297 self.assertEquals(comb, sorted(set(comb)))
298
299 # Check interrelationships
300 self.assertEquals(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
301 self.assertEquals(perm, [t for t in prod if len(set(t))==r]) # perm: prods with no dups
302 self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
303 self.assertEqual(comb, [t for t in cwr if len(set(t))==r]) # comb: cwrs without dups
304 self.assertEqual(comb, filter(set(cwr).__contains__, perm)) # comb: perm that is a cwr
305 self.assertEqual(comb, filter(set(perm).__contains__, cwr)) # comb: cwr that is a perm
306 self.assertEqual(comb, sorted(set(cwr) & set(perm))) # comb: both a cwr and a perm
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000307
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000308 def test_compress(self):
Raymond Hettinger2e2909f2009-02-19 02:15:14 +0000309 self.assertEqual(list(compress(data='ABCDEF', selectors=[1,0,1,0,1,1])), list('ACEF'))
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000310 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
311 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
312 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
313 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
314 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
315 n = 10000
316 data = chain.from_iterable(repeat(range(6), n))
317 selectors = chain.from_iterable(repeat((0, 1)))
318 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
319 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
320 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
321 self.assertRaises(TypeError, compress, range(6)) # too few args
322 self.assertRaises(TypeError, compress, range(6), None) # too many args
323
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000324 def test_count(self):
325 self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
326 self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000327 self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000328 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
329 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000330 self.assertRaises(TypeError, count, 2, 3, 4)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000331 self.assertRaises(TypeError, count, 'a')
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000332 self.assertEqual(list(islice(count(maxsize-5), 10)), range(maxsize-5, maxsize+5))
333 self.assertEqual(list(islice(count(-maxsize-5), 10)), range(-maxsize-5, -maxsize+5))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000334 c = count(3)
335 self.assertEqual(repr(c), 'count(3)')
336 c.next()
337 self.assertEqual(repr(c), 'count(4)')
Jack Diederich36234e82006-09-21 17:50:26 +0000338 c = count(-9)
339 self.assertEqual(repr(c), 'count(-9)')
340 c.next()
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000341 self.assertEqual(repr(count(10.25)), 'count(10.25)')
Jack Diederich36234e82006-09-21 17:50:26 +0000342 self.assertEqual(c.next(), -8)
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000343 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 +0000344 # Test repr (ignoring the L in longs)
345 r1 = repr(count(i)).replace('L', '')
346 r2 = 'count(%r)'.__mod__(i).replace('L', '')
347 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000348
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000349 def test_count_with_stride(self):
350 self.assertEqual(zip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingera4038032009-02-14 00:25:51 +0000351 self.assertEqual(zip('abc',count(start=2,step=3)),
352 [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingeraa681c72009-02-21 07:17:22 +0000353 self.assertEqual(zip('abc',count(step=-1)),
354 [('a', 0), ('b', -1), ('c', -2)])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000355 self.assertEqual(zip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
356 self.assertEqual(zip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
357 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
358 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
359 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettingeraa044612009-02-12 12:04:26 +0000360 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
361 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerdbe3bfb2009-02-12 12:43:01 +0000362 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
363 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000364 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
365 c = count(3, 5)
366 self.assertEqual(repr(c), 'count(3, 5)')
367 c.next()
368 self.assertEqual(repr(c), 'count(8, 5)')
369 c = count(-9, 0)
370 self.assertEqual(repr(c), 'count(-9, 0)')
371 c.next()
372 self.assertEqual(repr(c), 'count(-9, 0)')
373 c = count(-9, -3)
374 self.assertEqual(repr(c), 'count(-9, -3)')
375 c.next()
376 self.assertEqual(repr(c), 'count(-12, -3)')
377 self.assertEqual(repr(c), 'count(-12, -3)')
378 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
379 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
380 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
381 for i in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 10, sys.maxint-5, sys.maxint+5):
382 for j in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 1, 10, sys.maxint-5, sys.maxint+5):
383 # Test repr (ignoring the L in longs)
384 r1 = repr(count(i, j)).replace('L', '')
385 if j == 1:
386 r2 = ('count(%r)' % i).replace('L', '')
387 else:
388 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
389 self.assertEqual(r1, r2)
390
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000391 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000392 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000393 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000394 self.assertRaises(TypeError, cycle)
395 self.assertRaises(TypeError, cycle, 5)
396 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000397
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000398 def test_groupby(self):
399 # Check whether it accepts arguments correctly
400 self.assertEqual([], list(groupby([])))
401 self.assertEqual([], list(groupby([], key=id)))
402 self.assertRaises(TypeError, list, groupby('abc', []))
403 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000404 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000405
406 # Check normal input
407 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
408 (2,15,22), (3,16,23), (3,17,23)]
409 dup = []
410 for k, g in groupby(s, lambda r:r[0]):
411 for elem in g:
412 self.assertEqual(k, elem[0])
413 dup.append(elem)
414 self.assertEqual(s, dup)
415
416 # Check nested case
417 dup = []
418 for k, g in groupby(s, lambda r:r[0]):
419 for ik, ig in groupby(g, lambda r:r[2]):
420 for elem in ig:
421 self.assertEqual(k, elem[0])
422 self.assertEqual(ik, elem[2])
423 dup.append(elem)
424 self.assertEqual(s, dup)
425
426 # Check case where inner iterator is not used
427 keys = [k for k, g in groupby(s, lambda r:r[0])]
428 expectedkeys = set([r[0] for r in s])
429 self.assertEqual(set(keys), expectedkeys)
430 self.assertEqual(len(keys), len(expectedkeys))
431
432 # Exercise pipes and filters style
433 s = 'abracadabra'
434 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000435 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000436 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
437 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000438 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000439 self.assertEqual(r, ['a', 'b', 'r'])
440 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000441 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000442 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
443 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000444 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000445 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
446
447 # iter.next failure
448 class ExpectedError(Exception):
449 pass
450 def delayed_raise(n=0):
451 for i in range(n):
452 yield 'yo'
453 raise ExpectedError
454 def gulp(iterable, keyp=None, func=list):
455 return [func(g) for k, g in groupby(iterable, keyp)]
456
457 # iter.next failure on outer object
458 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
459 # iter.next failure on inner object
460 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
461
462 # __cmp__ failure
463 class DummyCmp:
464 def __cmp__(self, dst):
465 raise ExpectedError
466 s = [DummyCmp(), DummyCmp(), None]
467
468 # __cmp__ failure on outer object
469 self.assertRaises(ExpectedError, gulp, s, func=id)
470 # __cmp__ failure on inner object
471 self.assertRaises(ExpectedError, gulp, s)
472
473 # keyfunc failure
474 def keyfunc(obj):
475 if keyfunc.skip > 0:
476 keyfunc.skip -= 1
477 return obj
478 else:
479 raise ExpectedError
480
481 # keyfunc failure on outer object
482 keyfunc.skip = 0
483 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
484 keyfunc.skip = 1
485 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
486
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000487 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000488 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000489 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000490 self.assertEqual(list(ifilter(bool, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000491 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000492 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000493 self.assertRaises(TypeError, ifilter, lambda x:x)
494 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000495 self.assertRaises(TypeError, ifilter, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000496 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000497
498 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000499 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
500 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000501 self.assertEqual(list(ifilterfalse(bool, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000502 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000503 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000504 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
505 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000506 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000507 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000508
509 def test_izip(self):
510 ans = [(x,y) for x, y in izip('abc',count())]
511 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000512 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
513 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
Raymond Hettinger02420702003-06-29 20:36:23 +0000514 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000515 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000516 self.assertEqual(list(izip()), zip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000517 self.assertRaises(TypeError, izip, 3)
518 self.assertRaises(TypeError, izip, range(3), 3)
519 # Check tuple re-use (implementation detail)
520 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
521 zip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000522 self.assertEqual([pair for pair in izip('abc', 'def')],
523 zip('abc', 'def'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000524 ids = map(id, izip('abc', 'def'))
525 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000526 ids = map(id, list(izip('abc', 'def')))
527 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000528
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000529 def test_iziplongest(self):
530 for args in [
531 ['abc', range(6)],
532 [range(6), 'abc'],
533 [range(1000), range(2000,2100), range(3000,3050)],
534 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
535 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
536 ]:
537 target = map(None, *args)
538 self.assertEqual(list(izip_longest(*args)), target)
539 self.assertEqual(list(izip_longest(*args, **{})), target)
540 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
541 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Tim Petersea5962f2007-03-12 18:07:52 +0000542
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000543 self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
544
545 self.assertEqual(list(izip_longest()), zip())
546 self.assertEqual(list(izip_longest([])), zip([]))
547 self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
Tim Petersea5962f2007-03-12 18:07:52 +0000548
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000549 self.assertEqual(list(izip_longest('abc', 'defg', **{})), map(None, 'abc', 'defg')) # empty keyword dict
550 self.assertRaises(TypeError, izip_longest, 3)
551 self.assertRaises(TypeError, izip_longest, range(3), 3)
552
553 for stmt in [
554 "izip_longest('abc', fv=1)",
Tim Petersea5962f2007-03-12 18:07:52 +0000555 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000556 ]:
557 try:
558 eval(stmt, globals(), locals())
559 except TypeError:
560 pass
561 else:
562 self.fail('Did not raise Type in: ' + stmt)
Tim Petersea5962f2007-03-12 18:07:52 +0000563
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000564 # Check tuple re-use (implementation detail)
565 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
566 zip('abc', 'def'))
567 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
568 zip('abc', 'def'))
569 ids = map(id, izip_longest('abc', 'def'))
570 self.assertEqual(min(ids), max(ids))
571 ids = map(id, list(izip_longest('abc', 'def')))
572 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
573
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000574 def test_product(self):
575 for args, result in [
Raymond Hettingerd553d852008-03-04 04:17:08 +0000576 ([], [()]), # zero iterables
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000577 (['ab'], [('a',), ('b',)]), # one iterable
578 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
579 ([range(0), range(2), range(3)], []), # first iterable with zero length
580 ([range(2), range(0), range(3)], []), # middle iterable with zero length
581 ([range(2), range(3), range(0)], []), # last iterable with zero length
582 ]:
583 self.assertEqual(list(product(*args)), result)
Raymond Hettinger08ff6822008-02-29 02:21:48 +0000584 for r in range(4):
585 self.assertEqual(list(product(*(args*r))),
586 list(product(*args, **dict(repeat=r))))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000587 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
588 self.assertRaises(TypeError, product, range(6), None)
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000589
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000590 def product1(*args, **kwds):
591 pools = map(tuple, args) * kwds.get('repeat', 1)
592 n = len(pools)
593 if n == 0:
594 yield ()
595 return
596 if any(len(pool) == 0 for pool in pools):
597 return
598 indices = [0] * n
599 yield tuple(pool[i] for pool, i in zip(pools, indices))
600 while 1:
601 for i in reversed(range(n)): # right to left
602 if indices[i] == len(pools[i]) - 1:
603 continue
604 indices[i] += 1
605 for j in range(i+1, n):
606 indices[j] = 0
607 yield tuple(pool[i] for pool, i in zip(pools, indices))
608 break
609 else:
610 return
611
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000612 def product2(*args, **kwds):
613 'Pure python version used in docs'
614 pools = map(tuple, args) * kwds.get('repeat', 1)
615 result = [[]]
616 for pool in pools:
617 result = [x+[y] for x in result for y in pool]
618 for prod in result:
619 yield tuple(prod)
620
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000621 argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
622 set('abcdefg'), range(11), tuple(range(13))]
623 for i in range(100):
624 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Raymond Hettingerd553d852008-03-04 04:17:08 +0000625 expected_len = prod(map(len, args))
626 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000627 self.assertEqual(list(product(*args)), list(product1(*args)))
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000628 self.assertEqual(list(product(*args)), list(product2(*args)))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000629 args = map(iter, args)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000630 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000631
Raymond Hettinger73d79632008-02-23 02:20:41 +0000632 # Test implementation detail: tuple re-use
633 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
634 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000635
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000636 def test_repeat(self):
Raymond Hettinger182edae2009-02-19 02:38:25 +0000637 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000638 self.assertEqual(zip(xrange(3),repeat('a')),
639 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000640 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000641 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000642 self.assertEqual(list(repeat('a', 0)), [])
643 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000644 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000645 self.assertRaises(TypeError, repeat, None, 3, 4)
646 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000647 r = repeat(1+0j)
648 self.assertEqual(repr(r), 'repeat((1+0j))')
649 r = repeat(1+0j, 5)
650 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
651 list(r)
652 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000653
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000654 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000655 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
656 [0**1, 1**2, 2**3])
657 self.assertEqual(list(imap(None, 'abc', range(5))),
658 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000659 self.assertEqual(list(imap(None, 'abc', count())),
660 [('a',0),('b',1),('c',2)])
661 self.assertEqual(take(2,imap(None, 'abc', count())),
662 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000663 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000664 self.assertRaises(TypeError, imap)
665 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000666 self.assertRaises(TypeError, imap(10, range(5)).next)
667 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
668 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000669
670 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000671 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
672 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000673 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
674 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000675 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger47317092008-01-17 03:02:14 +0000676 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
677 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000678 self.assertRaises(TypeError, starmap)
679 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
680 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
681 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
682 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000683
684 def test_islice(self):
685 for args in [ # islice(args) should agree with range(args)
686 (10, 20, 3),
687 (10, 3, 20),
688 (10, 20),
689 (10, 3),
690 (20,)
691 ]:
692 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
693
694 for args, tgtargs in [ # Stop when seqn is exhausted
695 ((10, 110, 3), ((10, 100, 3))),
696 ((10, 110), ((10, 100))),
697 ((110,), (100,))
698 ]:
699 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
700
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000701 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000702 self.assertEqual(list(islice(xrange(10), None)), range(10))
Raymond Hettingerb2594052004-12-05 09:25:51 +0000703 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
704 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000705 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
706 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
707
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000708 # Test number of items consumed SF #1171417
709 it = iter(range(10))
710 self.assertEqual(list(islice(it, 3)), range(3))
711 self.assertEqual(list(it), range(3, 10))
712
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000713 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000714 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000715 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
716 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
717 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
718 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
719 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000720 self.assertRaises(ValueError, islice, xrange(10), 'a')
721 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
722 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
723 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
724 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +0000725 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000726
727 def test_takewhile(self):
728 data = [1, 3, 5, 20, 2, 4, 6, 8]
729 underten = lambda x: x<10
730 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000731 self.assertEqual(list(takewhile(underten, [])), [])
732 self.assertRaises(TypeError, takewhile)
733 self.assertRaises(TypeError, takewhile, operator.pow)
734 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
735 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
736 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000737 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
738 self.assertEqual(list(t), [1, 1, 1])
739 self.assertRaises(StopIteration, t.next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000740
741 def test_dropwhile(self):
742 data = [1, 3, 5, 20, 2, 4, 6, 8]
743 underten = lambda x: x<10
744 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000745 self.assertEqual(list(dropwhile(underten, [])), [])
746 self.assertRaises(TypeError, dropwhile)
747 self.assertRaises(TypeError, dropwhile, operator.pow)
748 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
749 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
750 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000751
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000752 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000753 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000754 def irange(n):
755 for i in xrange(n):
756 yield i
757
758 a, b = tee([]) # test empty iterator
759 self.assertEqual(list(a), [])
760 self.assertEqual(list(b), [])
761
762 a, b = tee(irange(n)) # test 100% interleaved
763 self.assertEqual(zip(a,b), zip(range(n),range(n)))
764
765 a, b = tee(irange(n)) # test 0% interleaved
766 self.assertEqual(list(a), range(n))
767 self.assertEqual(list(b), range(n))
768
769 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000770 for i in xrange(100):
771 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000772 del a
773 self.assertEqual(list(b), range(n))
774
775 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000776 for i in xrange(100):
777 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000778 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000779 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000780
781 for j in xrange(5): # test randomly interleaved
782 order = [0]*n + [1]*n
783 random.shuffle(order)
784 lists = ([], [])
785 its = tee(irange(n))
786 for i in order:
787 value = its[i].next()
788 lists[i].append(value)
789 self.assertEqual(lists[0], range(n))
790 self.assertEqual(lists[1], range(n))
791
Raymond Hettingerad983e72003-11-12 14:32:26 +0000792 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000793 self.assertRaises(TypeError, tee)
794 self.assertRaises(TypeError, tee, 3)
795 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000796 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000797
Raymond Hettingerad983e72003-11-12 14:32:26 +0000798 # tee object should be instantiable
799 a, b = tee('abc')
800 c = type(a)('def')
801 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000802
Raymond Hettingerad983e72003-11-12 14:32:26 +0000803 # test long-lagged and multi-way split
804 a, b, c = tee(xrange(2000), 3)
805 for i in xrange(100):
806 self.assertEqual(a.next(), i)
807 self.assertEqual(list(b), range(2000))
808 self.assertEqual([c.next(), c.next()], range(2))
809 self.assertEqual(list(a), range(100,2000))
810 self.assertEqual(list(c), range(2,2000))
811
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000812 # test values of n
813 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Neal Norwitz69e88972006-09-02 02:58:13 +0000814 self.assertRaises(ValueError, tee, [], -1)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000815 for n in xrange(5):
816 result = tee('abc', n)
817 self.assertEqual(type(result), tuple)
818 self.assertEqual(len(result), n)
819 self.assertEqual(map(list, result), [list('abc')]*n)
820
Raymond Hettingerad983e72003-11-12 14:32:26 +0000821 # tee pass-through to copyable iterator
822 a, b = tee('abc')
823 c, d = tee(a)
824 self.assert_(a is c)
825
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000826 # test tee_new
827 t1, t2 = tee('abc')
828 tnew = type(t1)
829 self.assertRaises(TypeError, tnew)
830 self.assertRaises(TypeError, tnew, 10)
831 t3 = tnew(t1)
832 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000833
Raymond Hettingera9f60922004-10-17 16:40:14 +0000834 # test that tee objects are weak referencable
835 a, b = tee(xrange(10))
836 p = proxy(a)
837 self.assertEqual(getattr(p, '__class__'), type(b))
838 del a
839 self.assertRaises(ReferenceError, getattr, p, '__class__')
840
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000841 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000842 self.assertRaises(StopIteration, izip().next)
843
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000844 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000845 self.assertRaises(StopIteration, f([]).next)
846 self.assertRaises(StopIteration, f(StopNow()).next)
847
848 self.assertRaises(StopIteration, islice([], None).next)
849 self.assertRaises(StopIteration, islice(StopNow(), None).next)
850
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000851 p, q = tee([])
852 self.assertRaises(StopIteration, p.next)
853 self.assertRaises(StopIteration, q.next)
854 p, q = tee(StopNow())
855 self.assertRaises(StopIteration, p.next)
856 self.assertRaises(StopIteration, q.next)
857
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000858 self.assertRaises(StopIteration, repeat(None, 0).next)
859
860 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
861 self.assertRaises(StopIteration, f(lambda x:x, []).next)
862 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
863
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000864class TestExamples(unittest.TestCase):
865
866 def test_chain(self):
867 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
868
869 def test_chain_from_iterable(self):
870 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
871
872 def test_combinations(self):
873 self.assertEqual(list(combinations('ABCD', 2)),
874 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
875 self.assertEqual(list(combinations(range(4), 3)),
876 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
877
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000878 def test_combinations_with_replacement(self):
879 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
880 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
881
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000882 def test_compress(self):
883 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
884
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000885 def test_count(self):
886 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
887
888 def test_cycle(self):
889 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
890
891 def test_dropwhile(self):
892 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
893
894 def test_groupby(self):
895 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
896 list('ABCDAB'))
897 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
898 [list('AAAA'), list('BBB'), list('CC'), list('D')])
899
900 def test_ifilter(self):
901 self.assertEqual(list(ifilter(lambda x: x%2, range(10))), [1,3,5,7,9])
902
903 def test_ifilterfalse(self):
904 self.assertEqual(list(ifilterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
905
906 def test_imap(self):
907 self.assertEqual(list(imap(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
908
909 def test_islice(self):
910 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
911 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
912 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
913 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
914
915 def test_izip(self):
916 self.assertEqual(list(izip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
917
918 def test_izip_longest(self):
919 self.assertEqual(list(izip_longest('ABCD', 'xy', fillvalue='-')),
920 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
921
922 def test_permutations(self):
923 self.assertEqual(list(permutations('ABCD', 2)),
924 map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split()))
925 self.assertEqual(list(permutations(range(3))),
926 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
927
928 def test_product(self):
929 self.assertEqual(list(product('ABCD', 'xy')),
930 map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split()))
931 self.assertEqual(list(product(range(2), repeat=3)),
932 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
933 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
934
935 def test_repeat(self):
936 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
937
938 def test_stapmap(self):
939 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
940 [32, 9, 1000])
941
942 def test_takewhile(self):
943 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
944
945
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000946class TestGC(unittest.TestCase):
947
948 def makecycle(self, iterator, container):
949 container.append(iterator)
950 iterator.next()
951 del container, iterator
952
953 def test_chain(self):
954 a = []
955 self.makecycle(chain(a), a)
956
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000957 def test_chain_from_iterable(self):
958 a = []
959 self.makecycle(chain.from_iterable([a]), a)
960
961 def test_combinations(self):
962 a = []
963 self.makecycle(combinations([1,2,a,3], 3), a)
964
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000965 def test_combinations_with_replacement(self):
966 a = []
967 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
968
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000969 def test_compress(self):
970 a = []
971 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
972
Raymond Hettingerb21d8102009-02-16 20:39:12 +0000973 def test_count(self):
974 a = []
975 Int = type('Int', (int,), dict(x=a))
976 self.makecycle(count(Int(0), Int(1)), a)
977
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000978 def test_cycle(self):
979 a = []
980 self.makecycle(cycle([a]*2), a)
981
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000982 def test_dropwhile(self):
983 a = []
984 self.makecycle(dropwhile(bool, [0, a, a]), a)
985
986 def test_groupby(self):
987 a = []
988 self.makecycle(groupby([a]*2, lambda x:x), a)
989
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000990 def test_issue2246(self):
991 # Issue 2246 -- the _grouper iterator was not included in GC
992 n = 10
993 keyfunc = lambda x: x
994 for i, j in groupby(xrange(n), key=keyfunc):
995 keyfunc.__dict__.setdefault('x',[]).append(j)
996
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000997 def test_ifilter(self):
998 a = []
999 self.makecycle(ifilter(lambda x:True, [a]*2), a)
1000
1001 def test_ifilterfalse(self):
1002 a = []
1003 self.makecycle(ifilterfalse(lambda x:False, a), a)
1004
1005 def test_izip(self):
1006 a = []
1007 self.makecycle(izip([a]*2, [a]*3), a)
1008
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001009 def test_izip_longest(self):
1010 a = []
1011 self.makecycle(izip_longest([a]*2, [a]*3), a)
1012 b = [a, None]
1013 self.makecycle(izip_longest([a]*2, [a]*3, fillvalue=b), a)
1014
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001015 def test_imap(self):
1016 a = []
1017 self.makecycle(imap(lambda x:x, [a]*2), a)
1018
1019 def test_islice(self):
1020 a = []
1021 self.makecycle(islice([a]*2, None), a)
1022
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001023 def test_permutations(self):
1024 a = []
1025 self.makecycle(permutations([1,2,a,3], 3), a)
1026
1027 def test_product(self):
1028 a = []
1029 self.makecycle(product([1,2,a,3], repeat=3), a)
1030
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001031 def test_repeat(self):
1032 a = []
1033 self.makecycle(repeat(a), a)
1034
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001035 def test_starmap(self):
1036 a = []
1037 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1038
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001039 def test_takewhile(self):
1040 a = []
1041 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1042
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001043def R(seqn):
1044 'Regular generator'
1045 for i in seqn:
1046 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001047
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001048class G:
1049 'Sequence using __getitem__'
1050 def __init__(self, seqn):
1051 self.seqn = seqn
1052 def __getitem__(self, i):
1053 return self.seqn[i]
1054
1055class I:
1056 'Sequence using iterator protocol'
1057 def __init__(self, seqn):
1058 self.seqn = seqn
1059 self.i = 0
1060 def __iter__(self):
1061 return self
1062 def next(self):
1063 if self.i >= len(self.seqn): raise StopIteration
1064 v = self.seqn[self.i]
1065 self.i += 1
1066 return v
1067
1068class Ig:
1069 'Sequence using iterator protocol defined with a generator'
1070 def __init__(self, seqn):
1071 self.seqn = seqn
1072 self.i = 0
1073 def __iter__(self):
1074 for val in self.seqn:
1075 yield val
1076
1077class X:
1078 'Missing __getitem__ and __iter__'
1079 def __init__(self, seqn):
1080 self.seqn = seqn
1081 self.i = 0
1082 def next(self):
1083 if self.i >= len(self.seqn): raise StopIteration
1084 v = self.seqn[self.i]
1085 self.i += 1
1086 return v
1087
1088class N:
1089 'Iterator missing next()'
1090 def __init__(self, seqn):
1091 self.seqn = seqn
1092 self.i = 0
1093 def __iter__(self):
1094 return self
1095
1096class E:
1097 'Test propagation of exceptions'
1098 def __init__(self, seqn):
1099 self.seqn = seqn
1100 self.i = 0
1101 def __iter__(self):
1102 return self
1103 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001104 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001105
1106class S:
1107 'Test immediate stop'
1108 def __init__(self, seqn):
1109 pass
1110 def __iter__(self):
1111 return self
1112 def next(self):
1113 raise StopIteration
1114
1115def L(seqn):
1116 'Test multiple tiers of iterators'
1117 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
1118
1119
1120class TestVariousIteratorArgs(unittest.TestCase):
1121
1122 def test_chain(self):
1123 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1124 for g in (G, I, Ig, S, L, R):
1125 self.assertEqual(list(chain(g(s))), list(g(s)))
1126 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001127 self.assertRaises(TypeError, list, chain(X(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001128 self.assertRaises(TypeError, list, chain(N(s)))
1129 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1130
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001131 def test_compress(self):
1132 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1133 n = len(s)
1134 for g in (G, I, Ig, S, L, R):
1135 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1136 self.assertRaises(TypeError, compress, X(s), repeat(1))
1137 self.assertRaises(TypeError, list, compress(N(s), repeat(1)))
1138 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1139
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001140 def test_product(self):
1141 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1142 self.assertRaises(TypeError, product, X(s))
1143 self.assertRaises(TypeError, product, N(s))
1144 self.assertRaises(ZeroDivisionError, product, E(s))
1145
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001146 def test_cycle(self):
1147 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1148 for g in (G, I, Ig, S, L, R):
1149 tgtlen = len(s) * 3
1150 expected = list(g(s))*3
1151 actual = list(islice(cycle(g(s)), tgtlen))
1152 self.assertEqual(actual, expected)
1153 self.assertRaises(TypeError, cycle, X(s))
1154 self.assertRaises(TypeError, list, cycle(N(s)))
1155 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1156
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001157 def test_groupby(self):
1158 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1159 for g in (G, I, Ig, S, L, R):
1160 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1161 self.assertRaises(TypeError, groupby, X(s))
1162 self.assertRaises(TypeError, list, groupby(N(s)))
1163 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1164
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001165 def test_ifilter(self):
1166 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1167 for g in (G, I, Ig, S, L, R):
1168 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
1169 self.assertRaises(TypeError, ifilter, isEven, X(s))
1170 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
1171 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
1172
1173 def test_ifilterfalse(self):
1174 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1175 for g in (G, I, Ig, S, L, R):
1176 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
1177 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
1178 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
1179 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
1180
1181 def test_izip(self):
1182 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1183 for g in (G, I, Ig, S, L, R):
1184 self.assertEqual(list(izip(g(s))), zip(g(s)))
1185 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
1186 self.assertRaises(TypeError, izip, X(s))
1187 self.assertRaises(TypeError, list, izip(N(s)))
1188 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
1189
Raymond Hettingerd36862c2007-02-21 05:20:38 +00001190 def test_iziplongest(self):
1191 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1192 for g in (G, I, Ig, S, L, R):
1193 self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
1194 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
1195 self.assertRaises(TypeError, izip_longest, X(s))
1196 self.assertRaises(TypeError, list, izip_longest(N(s)))
1197 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
1198
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001199 def test_imap(self):
1200 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1201 for g in (G, I, Ig, S, L, R):
1202 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
1203 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
1204 self.assertRaises(TypeError, imap, onearg, X(s))
1205 self.assertRaises(TypeError, list, imap(onearg, N(s)))
1206 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
1207
1208 def test_islice(self):
1209 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1210 for g in (G, I, Ig, S, L, R):
1211 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1212 self.assertRaises(TypeError, islice, X(s), 10)
1213 self.assertRaises(TypeError, list, islice(N(s), 10))
1214 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1215
1216 def test_starmap(self):
1217 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1218 for g in (G, I, Ig, S, L, R):
1219 ss = zip(s, s)
1220 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
1221 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
1222 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
1223 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1224
1225 def test_takewhile(self):
1226 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1227 for g in (G, I, Ig, S, L, R):
1228 tgt = []
1229 for elem in g(s):
1230 if not isEven(elem): break
1231 tgt.append(elem)
1232 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1233 self.assertRaises(TypeError, takewhile, isEven, X(s))
1234 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
1235 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1236
1237 def test_dropwhile(self):
1238 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1239 for g in (G, I, Ig, S, L, R):
1240 tgt = []
1241 for elem in g(s):
1242 if not tgt and isOdd(elem): continue
1243 tgt.append(elem)
1244 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1245 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
1246 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
1247 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1248
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001249 def test_tee(self):
1250 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1251 for g in (G, I, Ig, S, L, R):
1252 it1, it2 = tee(g(s))
1253 self.assertEqual(list(it1), list(g(s)))
1254 self.assertEqual(list(it2), list(g(s)))
1255 self.assertRaises(TypeError, tee, X(s))
1256 self.assertRaises(TypeError, list, tee(N(s))[0])
1257 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1258
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001259class LengthTransparency(unittest.TestCase):
1260
1261 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001262 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001263 self.assertEqual(len(repeat(None, 50)), 50)
1264 self.assertRaises(TypeError, len, repeat(None))
1265
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001266class RegressionTests(unittest.TestCase):
1267
1268 def test_sf_793826(self):
1269 # Fix Armin Rigo's successful efforts to wreak havoc
1270
1271 def mutatingtuple(tuple1, f, tuple2):
1272 # this builds a tuple t which is a copy of tuple1,
1273 # then calls f(t), then mutates t to be equal to tuple2
1274 # (needs len(tuple1) == len(tuple2)).
1275 def g(value, first=[1]):
1276 if first:
1277 del first[:]
1278 f(z.next())
1279 return value
1280 items = list(tuple2)
1281 items[1:1] = list(tuple1)
1282 gen = imap(g, items)
1283 z = izip(*[gen]*len(tuple1))
1284 z.next()
1285
1286 def f(t):
1287 global T
1288 T = t
1289 first[:] = list(T)
1290
1291 first = []
1292 mutatingtuple((1,2,3), f, (4,5,6))
1293 second = list(T)
1294 self.assertEqual(first, second)
1295
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001296
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001297 def test_sf_950057(self):
1298 # Make sure that chain() and cycle() catch exceptions immediately
1299 # rather than when shifting between input sources
1300
1301 def gen1():
1302 hist.append(0)
1303 yield 1
1304 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001305 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001306 hist.append(2)
1307
1308 def gen2(x):
1309 hist.append(3)
1310 yield 2
1311 hist.append(4)
1312 if x:
1313 raise StopIteration
1314
1315 hist = []
1316 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1317 self.assertEqual(hist, [0,1])
1318
1319 hist = []
1320 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1321 self.assertEqual(hist, [0,1])
1322
1323 hist = []
1324 self.assertRaises(AssertionError, list, cycle(gen1()))
1325 self.assertEqual(hist, [0,1])
1326
Georg Brandlb84c1372007-01-21 10:28:43 +00001327class SubclassWithKwargsTest(unittest.TestCase):
1328 def test_keywords_in_subclass(self):
1329 # count is not subclassable...
1330 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001331 starmap, islice, takewhile, dropwhile, cycle, compress):
Georg Brandlb84c1372007-01-21 10:28:43 +00001332 class Subclass(cls):
1333 def __init__(self, newarg=None, *args):
1334 cls.__init__(self, *args)
1335 try:
1336 Subclass(newarg=1)
1337 except TypeError, err:
1338 # we expect type errors because of wrong argument count
1339 self.failIf("does not take keyword arguments" in err.args[0])
1340
1341
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001342libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001343
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001344
1345>>> amounts = [120.15, 764.05, 823.14]
1346>>> for checknum, amount in izip(count(1200), amounts):
1347... print 'Check %d is for $%.2f' % (checknum, amount)
1348...
1349Check 1200 is for $120.15
1350Check 1201 is for $764.05
1351Check 1202 is for $823.14
1352
1353>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001354>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1355... print cube
1356...
13571
13588
135927
1360
1361>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001362>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001363... print name.title()
1364...
1365Alex
1366Laura
1367Martin
1368Walter
1369Samuele
1370
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001371>>> from operator import itemgetter
1372>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Armin Rigoa3f09272006-05-28 19:13:17 +00001373>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001374>>> for k, g in groupby(di, itemgetter(1)):
1375... print k, map(itemgetter(0), g)
1376...
13771 ['a', 'c', 'e']
13782 ['b', 'd', 'f']
13793 ['g']
1380
Raymond Hettinger734fb572004-01-20 20:04:40 +00001381# Find runs of consecutive numbers using groupby. The key to the solution
1382# is differencing with a range so that consecutive numbers all appear in
1383# same group.
1384>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
1385>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
1386... print map(operator.itemgetter(1), g)
1387...
1388[1]
1389[4, 5, 6]
1390[10]
1391[15, 16, 17, 18]
1392[22]
1393[25, 26, 27, 28]
1394
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001395>>> def take(n, iterable):
1396... "Return first n items of the iterable as a list"
1397... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001398
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001399>>> def enumerate(iterable, start=0):
1400... return izip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001401
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001402>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001403... "Return function(0), function(1), ..."
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001404... return imap(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001405
Raymond Hettingerf9bce832009-02-19 05:34:35 +00001406>>> def nth(iterable, n, default=None):
1407... "Returns the nth item or a default value"
1408... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001409
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001410>>> def quantify(iterable, pred=bool):
1411... "Count how many times the predicate is true"
1412... return sum(imap(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001413
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001414>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001415... "Returns the sequence elements and then returns None indefinitely"
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001416... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001417
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001418>>> def ncycles(iterable, n):
1419... "Returns the seqeuence elements n times"
1420... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001421
1422>>> def dotproduct(vec1, vec2):
1423... return sum(imap(operator.mul, vec1, vec2))
1424
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001425>>> def flatten(listOfLists):
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001426... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001427
1428>>> def repeatfunc(func, times=None, *args):
1429... "Repeat calls to func with specified arguments."
1430... " Example: repeatfunc(random.random)"
1431... if times is None:
1432... return starmap(func, repeat(args))
1433... else:
1434... return starmap(func, repeat(args, times))
1435
Raymond Hettingerd591f662003-10-26 15:34:50 +00001436>>> def pairwise(iterable):
1437... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1438... a, b = tee(iterable)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001439... for elem in b:
1440... break
Raymond Hettingerad983e72003-11-12 14:32:26 +00001441... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001442
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001443>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001444... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001445... args = [iter(iterable)] * n
Raymond Hettingerf080e6d2008-07-31 01:19:50 +00001446... return izip_longest(fillvalue=fillvalue, *args)
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001447
1448>>> def roundrobin(*iterables):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001449... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001450... # Recipe credited to George Sakkis
1451... pending = len(iterables)
1452... nexts = cycle(iter(it).next for it in iterables)
1453... while pending:
1454... try:
1455... for next in nexts:
1456... yield next()
1457... except StopIteration:
1458... pending -= 1
1459... nexts = cycle(islice(nexts, pending))
1460
1461>>> def powerset(iterable):
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001462... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1463... s = list(iterable)
1464... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001465
Raymond Hettinger44e15812009-01-02 21:26:45 +00001466>>> def unique_everseen(iterable, key=None):
1467... "List unique elements, preserving order. Remember all elements ever seen."
1468... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1469... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1470... seen = set()
1471... seen_add = seen.add
1472... if key is None:
1473... for element in iterable:
1474... if element not in seen:
1475... seen_add(element)
1476... yield element
1477... else:
1478... for element in iterable:
1479... k = key(element)
1480... if k not in seen:
1481... seen_add(k)
1482... yield element
1483
1484>>> def unique_justseen(iterable, key=None):
1485... "List unique elements, preserving order. Remember only the element just seen."
1486... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1487... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1488... return imap(next, imap(itemgetter(1), groupby(iterable, key)))
1489
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001490This is not part of the examples but it tests to make sure the definitions
1491perform as purported.
1492
Raymond Hettingera098b332003-09-08 23:58:40 +00001493>>> take(10, count())
1494[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1495
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001496>>> list(enumerate('abc'))
1497[(0, 'a'), (1, 'b'), (2, 'c')]
1498
1499>>> list(islice(tabulate(lambda x: 2*x), 4))
1500[0, 2, 4, 6]
1501
1502>>> nth('abcde', 3)
Raymond Hettingerd507afd2009-02-04 10:52:32 +00001503'd'
1504
1505>>> nth('abcde', 9) is None
1506True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001507
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001508>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000150950
1510
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001511>>> a = [[1, 2, 3], [4, 5, 6]]
1512>>> flatten(a)
1513[1, 2, 3, 4, 5, 6]
1514
1515>>> list(repeatfunc(pow, 5, 2, 3))
1516[8, 8, 8, 8, 8]
1517
1518>>> import random
1519>>> take(5, imap(int, repeatfunc(random.random)))
1520[0, 0, 0, 0, 0]
1521
Raymond Hettingerd591f662003-10-26 15:34:50 +00001522>>> list(pairwise('abcd'))
1523[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001524
Raymond Hettingerd591f662003-10-26 15:34:50 +00001525>>> list(pairwise([]))
1526[]
1527
1528>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001529[]
1530
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001531>>> list(islice(padnone('abc'), 0, 6))
1532['a', 'b', 'c', None, None, None]
1533
1534>>> list(ncycles('abc', 3))
1535['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1536
1537>>> dotproduct([1,2,3], [4,5,6])
153832
1539
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001540>>> list(grouper(3, 'abcdefg', 'x'))
1541[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1542
1543>>> list(roundrobin('abc', 'd', 'ef'))
1544['a', 'd', 'e', 'b', 'f', 'c']
1545
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001546>>> list(powerset([1,2,3]))
1547[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001548
Raymond Hettinger560f9a82009-01-27 13:26:35 +00001549>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1550True
1551
1552>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1553True
1554
Raymond Hettinger44e15812009-01-02 21:26:45 +00001555>>> list(unique_everseen('AAAABBBCCDAABBB'))
1556['A', 'B', 'C', 'D']
1557
1558>>> list(unique_everseen('ABBCcAD', str.lower))
1559['A', 'B', 'C', 'D']
1560
1561>>> list(unique_justseen('AAAABBBCCDAABBB'))
1562['A', 'B', 'C', 'D', 'A', 'B']
1563
1564>>> list(unique_justseen('ABBCcAD', str.lower))
1565['A', 'B', 'C', 'A', 'D']
1566
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001567"""
1568
1569__test__ = {'libreftest' : libreftest}
1570
1571def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001572 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Georg Brandlb84c1372007-01-21 10:28:43 +00001573 RegressionTests, LengthTransparency,
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001574 SubclassWithKwargsTest, TestExamples)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001575 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001576
1577 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001578 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001579 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001580 counts = [None] * 5
1581 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001582 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001583 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001584 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001585 print counts
1586
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001587 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001588 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001589
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001590if __name__ == "__main__":
1591 test_main(verbose=True)