blob: 9eb138954d2572cc10721cea38dc484515944fb0 [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 Hettinger31c769c2009-02-12 05:39:46 +0000353 self.assertEqual(zip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
354 self.assertEqual(zip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
355 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
356 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
357 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettingeraa044612009-02-12 12:04:26 +0000358 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
359 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerdbe3bfb2009-02-12 12:43:01 +0000360 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
361 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000362 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
363 c = count(3, 5)
364 self.assertEqual(repr(c), 'count(3, 5)')
365 c.next()
366 self.assertEqual(repr(c), 'count(8, 5)')
367 c = count(-9, 0)
368 self.assertEqual(repr(c), 'count(-9, 0)')
369 c.next()
370 self.assertEqual(repr(c), 'count(-9, 0)')
371 c = count(-9, -3)
372 self.assertEqual(repr(c), 'count(-9, -3)')
373 c.next()
374 self.assertEqual(repr(c), 'count(-12, -3)')
375 self.assertEqual(repr(c), 'count(-12, -3)')
376 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
377 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
378 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
379 for i in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 10, sys.maxint-5, sys.maxint+5):
380 for j in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 1, 10, sys.maxint-5, sys.maxint+5):
381 # Test repr (ignoring the L in longs)
382 r1 = repr(count(i, j)).replace('L', '')
383 if j == 1:
384 r2 = ('count(%r)' % i).replace('L', '')
385 else:
386 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
387 self.assertEqual(r1, r2)
388
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000389 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000390 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000391 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000392 self.assertRaises(TypeError, cycle)
393 self.assertRaises(TypeError, cycle, 5)
394 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000395
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000396 def test_groupby(self):
397 # Check whether it accepts arguments correctly
398 self.assertEqual([], list(groupby([])))
399 self.assertEqual([], list(groupby([], key=id)))
400 self.assertRaises(TypeError, list, groupby('abc', []))
401 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000402 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000403
404 # Check normal input
405 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
406 (2,15,22), (3,16,23), (3,17,23)]
407 dup = []
408 for k, g in groupby(s, lambda r:r[0]):
409 for elem in g:
410 self.assertEqual(k, elem[0])
411 dup.append(elem)
412 self.assertEqual(s, dup)
413
414 # Check nested case
415 dup = []
416 for k, g in groupby(s, lambda r:r[0]):
417 for ik, ig in groupby(g, lambda r:r[2]):
418 for elem in ig:
419 self.assertEqual(k, elem[0])
420 self.assertEqual(ik, elem[2])
421 dup.append(elem)
422 self.assertEqual(s, dup)
423
424 # Check case where inner iterator is not used
425 keys = [k for k, g in groupby(s, lambda r:r[0])]
426 expectedkeys = set([r[0] for r in s])
427 self.assertEqual(set(keys), expectedkeys)
428 self.assertEqual(len(keys), len(expectedkeys))
429
430 # Exercise pipes and filters style
431 s = 'abracadabra'
432 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000433 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000434 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
435 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000436 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000437 self.assertEqual(r, ['a', 'b', 'r'])
438 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000439 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000440 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
441 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000442 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000443 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
444
445 # iter.next failure
446 class ExpectedError(Exception):
447 pass
448 def delayed_raise(n=0):
449 for i in range(n):
450 yield 'yo'
451 raise ExpectedError
452 def gulp(iterable, keyp=None, func=list):
453 return [func(g) for k, g in groupby(iterable, keyp)]
454
455 # iter.next failure on outer object
456 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
457 # iter.next failure on inner object
458 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
459
460 # __cmp__ failure
461 class DummyCmp:
462 def __cmp__(self, dst):
463 raise ExpectedError
464 s = [DummyCmp(), DummyCmp(), None]
465
466 # __cmp__ failure on outer object
467 self.assertRaises(ExpectedError, gulp, s, func=id)
468 # __cmp__ failure on inner object
469 self.assertRaises(ExpectedError, gulp, s)
470
471 # keyfunc failure
472 def keyfunc(obj):
473 if keyfunc.skip > 0:
474 keyfunc.skip -= 1
475 return obj
476 else:
477 raise ExpectedError
478
479 # keyfunc failure on outer object
480 keyfunc.skip = 0
481 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
482 keyfunc.skip = 1
483 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
484
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000485 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000486 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000487 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000488 self.assertEqual(list(ifilter(bool, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000489 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000490 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000491 self.assertRaises(TypeError, ifilter, lambda x:x)
492 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000493 self.assertRaises(TypeError, ifilter, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000494 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000495
496 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000497 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
498 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000499 self.assertEqual(list(ifilterfalse(bool, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000500 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000501 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000502 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
503 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000504 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000505 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000506
507 def test_izip(self):
508 ans = [(x,y) for x, y in izip('abc',count())]
509 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000510 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
511 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
Raymond Hettinger02420702003-06-29 20:36:23 +0000512 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000513 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000514 self.assertEqual(list(izip()), zip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000515 self.assertRaises(TypeError, izip, 3)
516 self.assertRaises(TypeError, izip, range(3), 3)
517 # Check tuple re-use (implementation detail)
518 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
519 zip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000520 self.assertEqual([pair for pair in izip('abc', 'def')],
521 zip('abc', 'def'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000522 ids = map(id, izip('abc', 'def'))
523 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000524 ids = map(id, list(izip('abc', 'def')))
525 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000526
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000527 def test_iziplongest(self):
528 for args in [
529 ['abc', range(6)],
530 [range(6), 'abc'],
531 [range(1000), range(2000,2100), range(3000,3050)],
532 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
533 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
534 ]:
535 target = map(None, *args)
536 self.assertEqual(list(izip_longest(*args)), target)
537 self.assertEqual(list(izip_longest(*args, **{})), target)
538 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
539 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Tim Petersea5962f2007-03-12 18:07:52 +0000540
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000541 self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
542
543 self.assertEqual(list(izip_longest()), zip())
544 self.assertEqual(list(izip_longest([])), zip([]))
545 self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
Tim Petersea5962f2007-03-12 18:07:52 +0000546
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000547 self.assertEqual(list(izip_longest('abc', 'defg', **{})), map(None, 'abc', 'defg')) # empty keyword dict
548 self.assertRaises(TypeError, izip_longest, 3)
549 self.assertRaises(TypeError, izip_longest, range(3), 3)
550
551 for stmt in [
552 "izip_longest('abc', fv=1)",
Tim Petersea5962f2007-03-12 18:07:52 +0000553 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000554 ]:
555 try:
556 eval(stmt, globals(), locals())
557 except TypeError:
558 pass
559 else:
560 self.fail('Did not raise Type in: ' + stmt)
Tim Petersea5962f2007-03-12 18:07:52 +0000561
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000562 # Check tuple re-use (implementation detail)
563 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
564 zip('abc', 'def'))
565 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
566 zip('abc', 'def'))
567 ids = map(id, izip_longest('abc', 'def'))
568 self.assertEqual(min(ids), max(ids))
569 ids = map(id, list(izip_longest('abc', 'def')))
570 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
571
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000572 def test_product(self):
573 for args, result in [
Raymond Hettingerd553d852008-03-04 04:17:08 +0000574 ([], [()]), # zero iterables
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000575 (['ab'], [('a',), ('b',)]), # one iterable
576 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
577 ([range(0), range(2), range(3)], []), # first iterable with zero length
578 ([range(2), range(0), range(3)], []), # middle iterable with zero length
579 ([range(2), range(3), range(0)], []), # last iterable with zero length
580 ]:
581 self.assertEqual(list(product(*args)), result)
Raymond Hettinger08ff6822008-02-29 02:21:48 +0000582 for r in range(4):
583 self.assertEqual(list(product(*(args*r))),
584 list(product(*args, **dict(repeat=r))))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000585 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
586 self.assertRaises(TypeError, product, range(6), None)
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000587
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000588 def product1(*args, **kwds):
589 pools = map(tuple, args) * kwds.get('repeat', 1)
590 n = len(pools)
591 if n == 0:
592 yield ()
593 return
594 if any(len(pool) == 0 for pool in pools):
595 return
596 indices = [0] * n
597 yield tuple(pool[i] for pool, i in zip(pools, indices))
598 while 1:
599 for i in reversed(range(n)): # right to left
600 if indices[i] == len(pools[i]) - 1:
601 continue
602 indices[i] += 1
603 for j in range(i+1, n):
604 indices[j] = 0
605 yield tuple(pool[i] for pool, i in zip(pools, indices))
606 break
607 else:
608 return
609
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000610 def product2(*args, **kwds):
611 'Pure python version used in docs'
612 pools = map(tuple, args) * kwds.get('repeat', 1)
613 result = [[]]
614 for pool in pools:
615 result = [x+[y] for x in result for y in pool]
616 for prod in result:
617 yield tuple(prod)
618
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000619 argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
620 set('abcdefg'), range(11), tuple(range(13))]
621 for i in range(100):
622 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Raymond Hettingerd553d852008-03-04 04:17:08 +0000623 expected_len = prod(map(len, args))
624 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000625 self.assertEqual(list(product(*args)), list(product1(*args)))
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000626 self.assertEqual(list(product(*args)), list(product2(*args)))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000627 args = map(iter, args)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000628 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000629
Raymond Hettinger73d79632008-02-23 02:20:41 +0000630 # Test implementation detail: tuple re-use
631 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
632 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000633
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000634 def test_repeat(self):
Raymond Hettinger182edae2009-02-19 02:38:25 +0000635 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000636 self.assertEqual(zip(xrange(3),repeat('a')),
637 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000638 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000639 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000640 self.assertEqual(list(repeat('a', 0)), [])
641 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000642 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000643 self.assertRaises(TypeError, repeat, None, 3, 4)
644 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000645 r = repeat(1+0j)
646 self.assertEqual(repr(r), 'repeat((1+0j))')
647 r = repeat(1+0j, 5)
648 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
649 list(r)
650 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000651
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000652 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000653 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
654 [0**1, 1**2, 2**3])
655 self.assertEqual(list(imap(None, 'abc', range(5))),
656 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000657 self.assertEqual(list(imap(None, 'abc', count())),
658 [('a',0),('b',1),('c',2)])
659 self.assertEqual(take(2,imap(None, 'abc', count())),
660 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000661 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000662 self.assertRaises(TypeError, imap)
663 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000664 self.assertRaises(TypeError, imap(10, range(5)).next)
665 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
666 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000667
668 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000669 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
670 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000671 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
672 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000673 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger47317092008-01-17 03:02:14 +0000674 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
675 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000676 self.assertRaises(TypeError, starmap)
677 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
678 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
679 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
680 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000681
682 def test_islice(self):
683 for args in [ # islice(args) should agree with range(args)
684 (10, 20, 3),
685 (10, 3, 20),
686 (10, 20),
687 (10, 3),
688 (20,)
689 ]:
690 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
691
692 for args, tgtargs in [ # Stop when seqn is exhausted
693 ((10, 110, 3), ((10, 100, 3))),
694 ((10, 110), ((10, 100))),
695 ((110,), (100,))
696 ]:
697 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
698
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000699 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000700 self.assertEqual(list(islice(xrange(10), None)), range(10))
Raymond Hettingerb2594052004-12-05 09:25:51 +0000701 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
702 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000703 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
704 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
705
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000706 # Test number of items consumed SF #1171417
707 it = iter(range(10))
708 self.assertEqual(list(islice(it, 3)), range(3))
709 self.assertEqual(list(it), range(3, 10))
710
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000711 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000712 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000713 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
714 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
715 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
716 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
717 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000718 self.assertRaises(ValueError, islice, xrange(10), 'a')
719 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
720 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
721 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
722 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +0000723 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000724
725 def test_takewhile(self):
726 data = [1, 3, 5, 20, 2, 4, 6, 8]
727 underten = lambda x: x<10
728 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000729 self.assertEqual(list(takewhile(underten, [])), [])
730 self.assertRaises(TypeError, takewhile)
731 self.assertRaises(TypeError, takewhile, operator.pow)
732 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
733 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
734 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000735 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
736 self.assertEqual(list(t), [1, 1, 1])
737 self.assertRaises(StopIteration, t.next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000738
739 def test_dropwhile(self):
740 data = [1, 3, 5, 20, 2, 4, 6, 8]
741 underten = lambda x: x<10
742 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000743 self.assertEqual(list(dropwhile(underten, [])), [])
744 self.assertRaises(TypeError, dropwhile)
745 self.assertRaises(TypeError, dropwhile, operator.pow)
746 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
747 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
748 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000749
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000750 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000751 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000752 def irange(n):
753 for i in xrange(n):
754 yield i
755
756 a, b = tee([]) # test empty iterator
757 self.assertEqual(list(a), [])
758 self.assertEqual(list(b), [])
759
760 a, b = tee(irange(n)) # test 100% interleaved
761 self.assertEqual(zip(a,b), zip(range(n),range(n)))
762
763 a, b = tee(irange(n)) # test 0% interleaved
764 self.assertEqual(list(a), range(n))
765 self.assertEqual(list(b), range(n))
766
767 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000768 for i in xrange(100):
769 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000770 del a
771 self.assertEqual(list(b), range(n))
772
773 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000774 for i in xrange(100):
775 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000776 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000777 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000778
779 for j in xrange(5): # test randomly interleaved
780 order = [0]*n + [1]*n
781 random.shuffle(order)
782 lists = ([], [])
783 its = tee(irange(n))
784 for i in order:
785 value = its[i].next()
786 lists[i].append(value)
787 self.assertEqual(lists[0], range(n))
788 self.assertEqual(lists[1], range(n))
789
Raymond Hettingerad983e72003-11-12 14:32:26 +0000790 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000791 self.assertRaises(TypeError, tee)
792 self.assertRaises(TypeError, tee, 3)
793 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000794 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000795
Raymond Hettingerad983e72003-11-12 14:32:26 +0000796 # tee object should be instantiable
797 a, b = tee('abc')
798 c = type(a)('def')
799 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000800
Raymond Hettingerad983e72003-11-12 14:32:26 +0000801 # test long-lagged and multi-way split
802 a, b, c = tee(xrange(2000), 3)
803 for i in xrange(100):
804 self.assertEqual(a.next(), i)
805 self.assertEqual(list(b), range(2000))
806 self.assertEqual([c.next(), c.next()], range(2))
807 self.assertEqual(list(a), range(100,2000))
808 self.assertEqual(list(c), range(2,2000))
809
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000810 # test values of n
811 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Neal Norwitz69e88972006-09-02 02:58:13 +0000812 self.assertRaises(ValueError, tee, [], -1)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000813 for n in xrange(5):
814 result = tee('abc', n)
815 self.assertEqual(type(result), tuple)
816 self.assertEqual(len(result), n)
817 self.assertEqual(map(list, result), [list('abc')]*n)
818
Raymond Hettingerad983e72003-11-12 14:32:26 +0000819 # tee pass-through to copyable iterator
820 a, b = tee('abc')
821 c, d = tee(a)
822 self.assert_(a is c)
823
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000824 # test tee_new
825 t1, t2 = tee('abc')
826 tnew = type(t1)
827 self.assertRaises(TypeError, tnew)
828 self.assertRaises(TypeError, tnew, 10)
829 t3 = tnew(t1)
830 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000831
Raymond Hettingera9f60922004-10-17 16:40:14 +0000832 # test that tee objects are weak referencable
833 a, b = tee(xrange(10))
834 p = proxy(a)
835 self.assertEqual(getattr(p, '__class__'), type(b))
836 del a
837 self.assertRaises(ReferenceError, getattr, p, '__class__')
838
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000839 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000840 self.assertRaises(StopIteration, izip().next)
841
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000842 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000843 self.assertRaises(StopIteration, f([]).next)
844 self.assertRaises(StopIteration, f(StopNow()).next)
845
846 self.assertRaises(StopIteration, islice([], None).next)
847 self.assertRaises(StopIteration, islice(StopNow(), None).next)
848
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000849 p, q = tee([])
850 self.assertRaises(StopIteration, p.next)
851 self.assertRaises(StopIteration, q.next)
852 p, q = tee(StopNow())
853 self.assertRaises(StopIteration, p.next)
854 self.assertRaises(StopIteration, q.next)
855
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000856 self.assertRaises(StopIteration, repeat(None, 0).next)
857
858 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
859 self.assertRaises(StopIteration, f(lambda x:x, []).next)
860 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
861
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000862class TestExamples(unittest.TestCase):
863
864 def test_chain(self):
865 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
866
867 def test_chain_from_iterable(self):
868 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
869
870 def test_combinations(self):
871 self.assertEqual(list(combinations('ABCD', 2)),
872 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
873 self.assertEqual(list(combinations(range(4), 3)),
874 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
875
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000876 def test_combinations_with_replacement(self):
877 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
878 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
879
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000880 def test_compress(self):
881 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
882
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000883 def test_count(self):
884 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
885
886 def test_cycle(self):
887 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
888
889 def test_dropwhile(self):
890 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
891
892 def test_groupby(self):
893 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
894 list('ABCDAB'))
895 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
896 [list('AAAA'), list('BBB'), list('CC'), list('D')])
897
898 def test_ifilter(self):
899 self.assertEqual(list(ifilter(lambda x: x%2, range(10))), [1,3,5,7,9])
900
901 def test_ifilterfalse(self):
902 self.assertEqual(list(ifilterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
903
904 def test_imap(self):
905 self.assertEqual(list(imap(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
906
907 def test_islice(self):
908 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
909 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
910 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
911 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
912
913 def test_izip(self):
914 self.assertEqual(list(izip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
915
916 def test_izip_longest(self):
917 self.assertEqual(list(izip_longest('ABCD', 'xy', fillvalue='-')),
918 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
919
920 def test_permutations(self):
921 self.assertEqual(list(permutations('ABCD', 2)),
922 map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split()))
923 self.assertEqual(list(permutations(range(3))),
924 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
925
926 def test_product(self):
927 self.assertEqual(list(product('ABCD', 'xy')),
928 map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split()))
929 self.assertEqual(list(product(range(2), repeat=3)),
930 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
931 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
932
933 def test_repeat(self):
934 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
935
936 def test_stapmap(self):
937 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
938 [32, 9, 1000])
939
940 def test_takewhile(self):
941 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
942
943
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000944class TestGC(unittest.TestCase):
945
946 def makecycle(self, iterator, container):
947 container.append(iterator)
948 iterator.next()
949 del container, iterator
950
951 def test_chain(self):
952 a = []
953 self.makecycle(chain(a), a)
954
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000955 def test_chain_from_iterable(self):
956 a = []
957 self.makecycle(chain.from_iterable([a]), a)
958
959 def test_combinations(self):
960 a = []
961 self.makecycle(combinations([1,2,a,3], 3), a)
962
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000963 def test_combinations_with_replacement(self):
964 a = []
965 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
966
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000967 def test_compress(self):
968 a = []
969 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
970
Raymond Hettingerb21d8102009-02-16 20:39:12 +0000971 def test_count(self):
972 a = []
973 Int = type('Int', (int,), dict(x=a))
974 self.makecycle(count(Int(0), Int(1)), a)
975
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000976 def test_cycle(self):
977 a = []
978 self.makecycle(cycle([a]*2), a)
979
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000980 def test_dropwhile(self):
981 a = []
982 self.makecycle(dropwhile(bool, [0, a, a]), a)
983
984 def test_groupby(self):
985 a = []
986 self.makecycle(groupby([a]*2, lambda x:x), a)
987
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000988 def test_issue2246(self):
989 # Issue 2246 -- the _grouper iterator was not included in GC
990 n = 10
991 keyfunc = lambda x: x
992 for i, j in groupby(xrange(n), key=keyfunc):
993 keyfunc.__dict__.setdefault('x',[]).append(j)
994
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000995 def test_ifilter(self):
996 a = []
997 self.makecycle(ifilter(lambda x:True, [a]*2), a)
998
999 def test_ifilterfalse(self):
1000 a = []
1001 self.makecycle(ifilterfalse(lambda x:False, a), a)
1002
1003 def test_izip(self):
1004 a = []
1005 self.makecycle(izip([a]*2, [a]*3), a)
1006
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001007 def test_izip_longest(self):
1008 a = []
1009 self.makecycle(izip_longest([a]*2, [a]*3), a)
1010 b = [a, None]
1011 self.makecycle(izip_longest([a]*2, [a]*3, fillvalue=b), a)
1012
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001013 def test_imap(self):
1014 a = []
1015 self.makecycle(imap(lambda x:x, [a]*2), a)
1016
1017 def test_islice(self):
1018 a = []
1019 self.makecycle(islice([a]*2, None), a)
1020
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001021 def test_permutations(self):
1022 a = []
1023 self.makecycle(permutations([1,2,a,3], 3), a)
1024
1025 def test_product(self):
1026 a = []
1027 self.makecycle(product([1,2,a,3], repeat=3), a)
1028
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001029 def test_repeat(self):
1030 a = []
1031 self.makecycle(repeat(a), a)
1032
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001033 def test_starmap(self):
1034 a = []
1035 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1036
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001037 def test_takewhile(self):
1038 a = []
1039 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1040
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001041def R(seqn):
1042 'Regular generator'
1043 for i in seqn:
1044 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001045
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001046class G:
1047 'Sequence using __getitem__'
1048 def __init__(self, seqn):
1049 self.seqn = seqn
1050 def __getitem__(self, i):
1051 return self.seqn[i]
1052
1053class I:
1054 'Sequence using iterator protocol'
1055 def __init__(self, seqn):
1056 self.seqn = seqn
1057 self.i = 0
1058 def __iter__(self):
1059 return self
1060 def next(self):
1061 if self.i >= len(self.seqn): raise StopIteration
1062 v = self.seqn[self.i]
1063 self.i += 1
1064 return v
1065
1066class Ig:
1067 'Sequence using iterator protocol defined with a generator'
1068 def __init__(self, seqn):
1069 self.seqn = seqn
1070 self.i = 0
1071 def __iter__(self):
1072 for val in self.seqn:
1073 yield val
1074
1075class X:
1076 'Missing __getitem__ and __iter__'
1077 def __init__(self, seqn):
1078 self.seqn = seqn
1079 self.i = 0
1080 def next(self):
1081 if self.i >= len(self.seqn): raise StopIteration
1082 v = self.seqn[self.i]
1083 self.i += 1
1084 return v
1085
1086class N:
1087 'Iterator missing next()'
1088 def __init__(self, seqn):
1089 self.seqn = seqn
1090 self.i = 0
1091 def __iter__(self):
1092 return self
1093
1094class E:
1095 'Test propagation of exceptions'
1096 def __init__(self, seqn):
1097 self.seqn = seqn
1098 self.i = 0
1099 def __iter__(self):
1100 return self
1101 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001102 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001103
1104class S:
1105 'Test immediate stop'
1106 def __init__(self, seqn):
1107 pass
1108 def __iter__(self):
1109 return self
1110 def next(self):
1111 raise StopIteration
1112
1113def L(seqn):
1114 'Test multiple tiers of iterators'
1115 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
1116
1117
1118class TestVariousIteratorArgs(unittest.TestCase):
1119
1120 def test_chain(self):
1121 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1122 for g in (G, I, Ig, S, L, R):
1123 self.assertEqual(list(chain(g(s))), list(g(s)))
1124 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001125 self.assertRaises(TypeError, list, chain(X(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001126 self.assertRaises(TypeError, list, chain(N(s)))
1127 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1128
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001129 def test_compress(self):
1130 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1131 n = len(s)
1132 for g in (G, I, Ig, S, L, R):
1133 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1134 self.assertRaises(TypeError, compress, X(s), repeat(1))
1135 self.assertRaises(TypeError, list, compress(N(s), repeat(1)))
1136 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1137
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001138 def test_product(self):
1139 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1140 self.assertRaises(TypeError, product, X(s))
1141 self.assertRaises(TypeError, product, N(s))
1142 self.assertRaises(ZeroDivisionError, product, E(s))
1143
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001144 def test_cycle(self):
1145 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1146 for g in (G, I, Ig, S, L, R):
1147 tgtlen = len(s) * 3
1148 expected = list(g(s))*3
1149 actual = list(islice(cycle(g(s)), tgtlen))
1150 self.assertEqual(actual, expected)
1151 self.assertRaises(TypeError, cycle, X(s))
1152 self.assertRaises(TypeError, list, cycle(N(s)))
1153 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1154
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001155 def test_groupby(self):
1156 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1157 for g in (G, I, Ig, S, L, R):
1158 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1159 self.assertRaises(TypeError, groupby, X(s))
1160 self.assertRaises(TypeError, list, groupby(N(s)))
1161 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1162
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001163 def test_ifilter(self):
1164 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1165 for g in (G, I, Ig, S, L, R):
1166 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
1167 self.assertRaises(TypeError, ifilter, isEven, X(s))
1168 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
1169 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
1170
1171 def test_ifilterfalse(self):
1172 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1173 for g in (G, I, Ig, S, L, R):
1174 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
1175 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
1176 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
1177 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
1178
1179 def test_izip(self):
1180 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1181 for g in (G, I, Ig, S, L, R):
1182 self.assertEqual(list(izip(g(s))), zip(g(s)))
1183 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
1184 self.assertRaises(TypeError, izip, X(s))
1185 self.assertRaises(TypeError, list, izip(N(s)))
1186 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
1187
Raymond Hettingerd36862c2007-02-21 05:20:38 +00001188 def test_iziplongest(self):
1189 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1190 for g in (G, I, Ig, S, L, R):
1191 self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
1192 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
1193 self.assertRaises(TypeError, izip_longest, X(s))
1194 self.assertRaises(TypeError, list, izip_longest(N(s)))
1195 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
1196
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001197 def test_imap(self):
1198 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1199 for g in (G, I, Ig, S, L, R):
1200 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
1201 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
1202 self.assertRaises(TypeError, imap, onearg, X(s))
1203 self.assertRaises(TypeError, list, imap(onearg, N(s)))
1204 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
1205
1206 def test_islice(self):
1207 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1208 for g in (G, I, Ig, S, L, R):
1209 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1210 self.assertRaises(TypeError, islice, X(s), 10)
1211 self.assertRaises(TypeError, list, islice(N(s), 10))
1212 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1213
1214 def test_starmap(self):
1215 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1216 for g in (G, I, Ig, S, L, R):
1217 ss = zip(s, s)
1218 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
1219 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
1220 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
1221 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1222
1223 def test_takewhile(self):
1224 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1225 for g in (G, I, Ig, S, L, R):
1226 tgt = []
1227 for elem in g(s):
1228 if not isEven(elem): break
1229 tgt.append(elem)
1230 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1231 self.assertRaises(TypeError, takewhile, isEven, X(s))
1232 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
1233 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1234
1235 def test_dropwhile(self):
1236 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1237 for g in (G, I, Ig, S, L, R):
1238 tgt = []
1239 for elem in g(s):
1240 if not tgt and isOdd(elem): continue
1241 tgt.append(elem)
1242 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1243 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
1244 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
1245 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1246
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001247 def test_tee(self):
1248 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1249 for g in (G, I, Ig, S, L, R):
1250 it1, it2 = tee(g(s))
1251 self.assertEqual(list(it1), list(g(s)))
1252 self.assertEqual(list(it2), list(g(s)))
1253 self.assertRaises(TypeError, tee, X(s))
1254 self.assertRaises(TypeError, list, tee(N(s))[0])
1255 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1256
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001257class LengthTransparency(unittest.TestCase):
1258
1259 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001260 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001261 self.assertEqual(len(repeat(None, 50)), 50)
1262 self.assertRaises(TypeError, len, repeat(None))
1263
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001264class RegressionTests(unittest.TestCase):
1265
1266 def test_sf_793826(self):
1267 # Fix Armin Rigo's successful efforts to wreak havoc
1268
1269 def mutatingtuple(tuple1, f, tuple2):
1270 # this builds a tuple t which is a copy of tuple1,
1271 # then calls f(t), then mutates t to be equal to tuple2
1272 # (needs len(tuple1) == len(tuple2)).
1273 def g(value, first=[1]):
1274 if first:
1275 del first[:]
1276 f(z.next())
1277 return value
1278 items = list(tuple2)
1279 items[1:1] = list(tuple1)
1280 gen = imap(g, items)
1281 z = izip(*[gen]*len(tuple1))
1282 z.next()
1283
1284 def f(t):
1285 global T
1286 T = t
1287 first[:] = list(T)
1288
1289 first = []
1290 mutatingtuple((1,2,3), f, (4,5,6))
1291 second = list(T)
1292 self.assertEqual(first, second)
1293
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001294
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001295 def test_sf_950057(self):
1296 # Make sure that chain() and cycle() catch exceptions immediately
1297 # rather than when shifting between input sources
1298
1299 def gen1():
1300 hist.append(0)
1301 yield 1
1302 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001303 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001304 hist.append(2)
1305
1306 def gen2(x):
1307 hist.append(3)
1308 yield 2
1309 hist.append(4)
1310 if x:
1311 raise StopIteration
1312
1313 hist = []
1314 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1315 self.assertEqual(hist, [0,1])
1316
1317 hist = []
1318 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1319 self.assertEqual(hist, [0,1])
1320
1321 hist = []
1322 self.assertRaises(AssertionError, list, cycle(gen1()))
1323 self.assertEqual(hist, [0,1])
1324
Georg Brandlb84c1372007-01-21 10:28:43 +00001325class SubclassWithKwargsTest(unittest.TestCase):
1326 def test_keywords_in_subclass(self):
1327 # count is not subclassable...
1328 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001329 starmap, islice, takewhile, dropwhile, cycle, compress):
Georg Brandlb84c1372007-01-21 10:28:43 +00001330 class Subclass(cls):
1331 def __init__(self, newarg=None, *args):
1332 cls.__init__(self, *args)
1333 try:
1334 Subclass(newarg=1)
1335 except TypeError, err:
1336 # we expect type errors because of wrong argument count
1337 self.failIf("does not take keyword arguments" in err.args[0])
1338
1339
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001340libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001341
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001342
1343>>> amounts = [120.15, 764.05, 823.14]
1344>>> for checknum, amount in izip(count(1200), amounts):
1345... print 'Check %d is for $%.2f' % (checknum, amount)
1346...
1347Check 1200 is for $120.15
1348Check 1201 is for $764.05
1349Check 1202 is for $823.14
1350
1351>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001352>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1353... print cube
1354...
13551
13568
135727
1358
1359>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001360>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001361... print name.title()
1362...
1363Alex
1364Laura
1365Martin
1366Walter
1367Samuele
1368
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001369>>> from operator import itemgetter
1370>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Armin Rigoa3f09272006-05-28 19:13:17 +00001371>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001372>>> for k, g in groupby(di, itemgetter(1)):
1373... print k, map(itemgetter(0), g)
1374...
13751 ['a', 'c', 'e']
13762 ['b', 'd', 'f']
13773 ['g']
1378
Raymond Hettinger734fb572004-01-20 20:04:40 +00001379# Find runs of consecutive numbers using groupby. The key to the solution
1380# is differencing with a range so that consecutive numbers all appear in
1381# same group.
1382>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
1383>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
1384... print map(operator.itemgetter(1), g)
1385...
1386[1]
1387[4, 5, 6]
1388[10]
1389[15, 16, 17, 18]
1390[22]
1391[25, 26, 27, 28]
1392
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001393>>> def take(n, iterable):
1394... "Return first n items of the iterable as a list"
1395... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001396
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001397>>> def enumerate(iterable, start=0):
1398... return izip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001399
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001400>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001401... "Return function(0), function(1), ..."
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001402... return imap(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001403
Raymond Hettingerf9bce832009-02-19 05:34:35 +00001404>>> def nth(iterable, n, default=None):
1405... "Returns the nth item or a default value"
1406... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001407
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001408>>> def quantify(iterable, pred=bool):
1409... "Count how many times the predicate is true"
1410... return sum(imap(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001411
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001412>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001413... "Returns the sequence elements and then returns None indefinitely"
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001414... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001415
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001416>>> def ncycles(iterable, n):
1417... "Returns the seqeuence elements n times"
1418... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001419
1420>>> def dotproduct(vec1, vec2):
1421... return sum(imap(operator.mul, vec1, vec2))
1422
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001423>>> def flatten(listOfLists):
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001424... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001425
1426>>> def repeatfunc(func, times=None, *args):
1427... "Repeat calls to func with specified arguments."
1428... " Example: repeatfunc(random.random)"
1429... if times is None:
1430... return starmap(func, repeat(args))
1431... else:
1432... return starmap(func, repeat(args, times))
1433
Raymond Hettingerd591f662003-10-26 15:34:50 +00001434>>> def pairwise(iterable):
1435... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1436... a, b = tee(iterable)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001437... for elem in b:
1438... break
Raymond Hettingerad983e72003-11-12 14:32:26 +00001439... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001440
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001441>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001442... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001443... args = [iter(iterable)] * n
Raymond Hettingerf080e6d2008-07-31 01:19:50 +00001444... return izip_longest(fillvalue=fillvalue, *args)
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001445
1446>>> def roundrobin(*iterables):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001447... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001448... # Recipe credited to George Sakkis
1449... pending = len(iterables)
1450... nexts = cycle(iter(it).next for it in iterables)
1451... while pending:
1452... try:
1453... for next in nexts:
1454... yield next()
1455... except StopIteration:
1456... pending -= 1
1457... nexts = cycle(islice(nexts, pending))
1458
1459>>> def powerset(iterable):
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001460... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1461... s = list(iterable)
1462... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001463
Raymond Hettinger44e15812009-01-02 21:26:45 +00001464>>> def unique_everseen(iterable, key=None):
1465... "List unique elements, preserving order. Remember all elements ever seen."
1466... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1467... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1468... seen = set()
1469... seen_add = seen.add
1470... if key is None:
1471... for element in iterable:
1472... if element not in seen:
1473... seen_add(element)
1474... yield element
1475... else:
1476... for element in iterable:
1477... k = key(element)
1478... if k not in seen:
1479... seen_add(k)
1480... yield element
1481
1482>>> def unique_justseen(iterable, key=None):
1483... "List unique elements, preserving order. Remember only the element just seen."
1484... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1485... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1486... return imap(next, imap(itemgetter(1), groupby(iterable, key)))
1487
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001488This is not part of the examples but it tests to make sure the definitions
1489perform as purported.
1490
Raymond Hettingera098b332003-09-08 23:58:40 +00001491>>> take(10, count())
1492[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1493
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001494>>> list(enumerate('abc'))
1495[(0, 'a'), (1, 'b'), (2, 'c')]
1496
1497>>> list(islice(tabulate(lambda x: 2*x), 4))
1498[0, 2, 4, 6]
1499
1500>>> nth('abcde', 3)
Raymond Hettingerd507afd2009-02-04 10:52:32 +00001501'd'
1502
1503>>> nth('abcde', 9) is None
1504True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001505
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001506>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000150750
1508
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001509>>> a = [[1, 2, 3], [4, 5, 6]]
1510>>> flatten(a)
1511[1, 2, 3, 4, 5, 6]
1512
1513>>> list(repeatfunc(pow, 5, 2, 3))
1514[8, 8, 8, 8, 8]
1515
1516>>> import random
1517>>> take(5, imap(int, repeatfunc(random.random)))
1518[0, 0, 0, 0, 0]
1519
Raymond Hettingerd591f662003-10-26 15:34:50 +00001520>>> list(pairwise('abcd'))
1521[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001522
Raymond Hettingerd591f662003-10-26 15:34:50 +00001523>>> list(pairwise([]))
1524[]
1525
1526>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001527[]
1528
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001529>>> list(islice(padnone('abc'), 0, 6))
1530['a', 'b', 'c', None, None, None]
1531
1532>>> list(ncycles('abc', 3))
1533['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1534
1535>>> dotproduct([1,2,3], [4,5,6])
153632
1537
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001538>>> list(grouper(3, 'abcdefg', 'x'))
1539[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1540
1541>>> list(roundrobin('abc', 'd', 'ef'))
1542['a', 'd', 'e', 'b', 'f', 'c']
1543
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001544>>> list(powerset([1,2,3]))
1545[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001546
Raymond Hettinger560f9a82009-01-27 13:26:35 +00001547>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1548True
1549
1550>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1551True
1552
Raymond Hettinger44e15812009-01-02 21:26:45 +00001553>>> list(unique_everseen('AAAABBBCCDAABBB'))
1554['A', 'B', 'C', 'D']
1555
1556>>> list(unique_everseen('ABBCcAD', str.lower))
1557['A', 'B', 'C', 'D']
1558
1559>>> list(unique_justseen('AAAABBBCCDAABBB'))
1560['A', 'B', 'C', 'D', 'A', 'B']
1561
1562>>> list(unique_justseen('ABBCcAD', str.lower))
1563['A', 'B', 'C', 'A', 'D']
1564
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001565"""
1566
1567__test__ = {'libreftest' : libreftest}
1568
1569def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001570 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Georg Brandlb84c1372007-01-21 10:28:43 +00001571 RegressionTests, LengthTransparency,
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001572 SubclassWithKwargsTest, TestExamples)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001573 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001574
1575 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001576 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001577 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001578 counts = [None] * 5
1579 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001580 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001581 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001582 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001583 print counts
1584
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001585 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001586 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001587
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001588if __name__ == "__main__":
1589 test_main(verbose=True)