blob: a33939eb20f655ea99f56a4688bb44ae055deb93 [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):
309 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
310 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
311 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
312 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
313 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
314 n = 10000
315 data = chain.from_iterable(repeat(range(6), n))
316 selectors = chain.from_iterable(repeat((0, 1)))
317 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
318 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
319 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
320 self.assertRaises(TypeError, compress, range(6)) # too few args
321 self.assertRaises(TypeError, compress, range(6), None) # too many args
322
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000323 def test_count(self):
324 self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
325 self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000326 self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000327 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
328 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000329 self.assertRaises(TypeError, count, 2, 3, 4)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000330 self.assertRaises(TypeError, count, 'a')
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000331 self.assertEqual(list(islice(count(maxsize-5), 10)), range(maxsize-5, maxsize+5))
332 self.assertEqual(list(islice(count(-maxsize-5), 10)), range(-maxsize-5, -maxsize+5))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000333 c = count(3)
334 self.assertEqual(repr(c), 'count(3)')
335 c.next()
336 self.assertEqual(repr(c), 'count(4)')
Jack Diederich36234e82006-09-21 17:50:26 +0000337 c = count(-9)
338 self.assertEqual(repr(c), 'count(-9)')
339 c.next()
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000340 self.assertEqual(repr(count(10.25)), 'count(10.25)')
Jack Diederich36234e82006-09-21 17:50:26 +0000341 self.assertEqual(c.next(), -8)
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000342 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 +0000343 # Test repr (ignoring the L in longs)
344 r1 = repr(count(i)).replace('L', '')
345 r2 = 'count(%r)'.__mod__(i).replace('L', '')
346 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000347
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000348 def test_count_with_stride(self):
349 self.assertEqual(zip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingera4038032009-02-14 00:25:51 +0000350 self.assertEqual(zip('abc',count(start=2,step=3)),
351 [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000352 self.assertEqual(zip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
353 self.assertEqual(zip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
354 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
355 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
356 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettingeraa044612009-02-12 12:04:26 +0000357 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
358 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerdbe3bfb2009-02-12 12:43:01 +0000359 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
360 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000361 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
362 c = count(3, 5)
363 self.assertEqual(repr(c), 'count(3, 5)')
364 c.next()
365 self.assertEqual(repr(c), 'count(8, 5)')
366 c = count(-9, 0)
367 self.assertEqual(repr(c), 'count(-9, 0)')
368 c.next()
369 self.assertEqual(repr(c), 'count(-9, 0)')
370 c = count(-9, -3)
371 self.assertEqual(repr(c), 'count(-9, -3)')
372 c.next()
373 self.assertEqual(repr(c), 'count(-12, -3)')
374 self.assertEqual(repr(c), 'count(-12, -3)')
375 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
376 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
377 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
378 for i in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 10, sys.maxint-5, sys.maxint+5):
379 for j in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 1, 10, sys.maxint-5, sys.maxint+5):
380 # Test repr (ignoring the L in longs)
381 r1 = repr(count(i, j)).replace('L', '')
382 if j == 1:
383 r2 = ('count(%r)' % i).replace('L', '')
384 else:
385 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
386 self.assertEqual(r1, r2)
387
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000388 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000389 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000390 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000391 self.assertRaises(TypeError, cycle)
392 self.assertRaises(TypeError, cycle, 5)
393 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000394
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000395 def test_groupby(self):
396 # Check whether it accepts arguments correctly
397 self.assertEqual([], list(groupby([])))
398 self.assertEqual([], list(groupby([], key=id)))
399 self.assertRaises(TypeError, list, groupby('abc', []))
400 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000401 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000402
403 # Check normal input
404 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
405 (2,15,22), (3,16,23), (3,17,23)]
406 dup = []
407 for k, g in groupby(s, lambda r:r[0]):
408 for elem in g:
409 self.assertEqual(k, elem[0])
410 dup.append(elem)
411 self.assertEqual(s, dup)
412
413 # Check nested case
414 dup = []
415 for k, g in groupby(s, lambda r:r[0]):
416 for ik, ig in groupby(g, lambda r:r[2]):
417 for elem in ig:
418 self.assertEqual(k, elem[0])
419 self.assertEqual(ik, elem[2])
420 dup.append(elem)
421 self.assertEqual(s, dup)
422
423 # Check case where inner iterator is not used
424 keys = [k for k, g in groupby(s, lambda r:r[0])]
425 expectedkeys = set([r[0] for r in s])
426 self.assertEqual(set(keys), expectedkeys)
427 self.assertEqual(len(keys), len(expectedkeys))
428
429 # Exercise pipes and filters style
430 s = 'abracadabra'
431 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000432 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000433 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
434 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000435 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000436 self.assertEqual(r, ['a', 'b', 'r'])
437 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000438 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000439 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
440 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000441 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000442 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
443
444 # iter.next failure
445 class ExpectedError(Exception):
446 pass
447 def delayed_raise(n=0):
448 for i in range(n):
449 yield 'yo'
450 raise ExpectedError
451 def gulp(iterable, keyp=None, func=list):
452 return [func(g) for k, g in groupby(iterable, keyp)]
453
454 # iter.next failure on outer object
455 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
456 # iter.next failure on inner object
457 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
458
459 # __cmp__ failure
460 class DummyCmp:
461 def __cmp__(self, dst):
462 raise ExpectedError
463 s = [DummyCmp(), DummyCmp(), None]
464
465 # __cmp__ failure on outer object
466 self.assertRaises(ExpectedError, gulp, s, func=id)
467 # __cmp__ failure on inner object
468 self.assertRaises(ExpectedError, gulp, s)
469
470 # keyfunc failure
471 def keyfunc(obj):
472 if keyfunc.skip > 0:
473 keyfunc.skip -= 1
474 return obj
475 else:
476 raise ExpectedError
477
478 # keyfunc failure on outer object
479 keyfunc.skip = 0
480 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
481 keyfunc.skip = 1
482 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
483
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000484 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000485 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000486 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000487 self.assertEqual(list(ifilter(bool, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000488 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000489 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000490 self.assertRaises(TypeError, ifilter, lambda x:x)
491 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000492 self.assertRaises(TypeError, ifilter, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000493 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000494
495 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000496 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
497 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000498 self.assertEqual(list(ifilterfalse(bool, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000499 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000500 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000501 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
502 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000503 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000504 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000505
506 def test_izip(self):
507 ans = [(x,y) for x, y in izip('abc',count())]
508 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000509 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
510 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
Raymond Hettinger02420702003-06-29 20:36:23 +0000511 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000512 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000513 self.assertEqual(list(izip()), zip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000514 self.assertRaises(TypeError, izip, 3)
515 self.assertRaises(TypeError, izip, range(3), 3)
516 # Check tuple re-use (implementation detail)
517 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
518 zip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000519 self.assertEqual([pair for pair in izip('abc', 'def')],
520 zip('abc', 'def'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000521 ids = map(id, izip('abc', 'def'))
522 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000523 ids = map(id, list(izip('abc', 'def')))
524 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000525
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000526 def test_iziplongest(self):
527 for args in [
528 ['abc', range(6)],
529 [range(6), 'abc'],
530 [range(1000), range(2000,2100), range(3000,3050)],
531 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
532 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
533 ]:
534 target = map(None, *args)
535 self.assertEqual(list(izip_longest(*args)), target)
536 self.assertEqual(list(izip_longest(*args, **{})), target)
537 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
538 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Tim Petersea5962f2007-03-12 18:07:52 +0000539
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000540 self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
541
542 self.assertEqual(list(izip_longest()), zip())
543 self.assertEqual(list(izip_longest([])), zip([]))
544 self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
Tim Petersea5962f2007-03-12 18:07:52 +0000545
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000546 self.assertEqual(list(izip_longest('abc', 'defg', **{})), map(None, 'abc', 'defg')) # empty keyword dict
547 self.assertRaises(TypeError, izip_longest, 3)
548 self.assertRaises(TypeError, izip_longest, range(3), 3)
549
550 for stmt in [
551 "izip_longest('abc', fv=1)",
Tim Petersea5962f2007-03-12 18:07:52 +0000552 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000553 ]:
554 try:
555 eval(stmt, globals(), locals())
556 except TypeError:
557 pass
558 else:
559 self.fail('Did not raise Type in: ' + stmt)
Tim Petersea5962f2007-03-12 18:07:52 +0000560
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000561 # Check tuple re-use (implementation detail)
562 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
563 zip('abc', 'def'))
564 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
565 zip('abc', 'def'))
566 ids = map(id, izip_longest('abc', 'def'))
567 self.assertEqual(min(ids), max(ids))
568 ids = map(id, list(izip_longest('abc', 'def')))
569 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
570
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000571 def test_product(self):
572 for args, result in [
Raymond Hettingerd553d852008-03-04 04:17:08 +0000573 ([], [()]), # zero iterables
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000574 (['ab'], [('a',), ('b',)]), # one iterable
575 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
576 ([range(0), range(2), range(3)], []), # first iterable with zero length
577 ([range(2), range(0), range(3)], []), # middle iterable with zero length
578 ([range(2), range(3), range(0)], []), # last iterable with zero length
579 ]:
580 self.assertEqual(list(product(*args)), result)
Raymond Hettinger08ff6822008-02-29 02:21:48 +0000581 for r in range(4):
582 self.assertEqual(list(product(*(args*r))),
583 list(product(*args, **dict(repeat=r))))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000584 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
585 self.assertRaises(TypeError, product, range(6), None)
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000586
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000587 def product1(*args, **kwds):
588 pools = map(tuple, args) * kwds.get('repeat', 1)
589 n = len(pools)
590 if n == 0:
591 yield ()
592 return
593 if any(len(pool) == 0 for pool in pools):
594 return
595 indices = [0] * n
596 yield tuple(pool[i] for pool, i in zip(pools, indices))
597 while 1:
598 for i in reversed(range(n)): # right to left
599 if indices[i] == len(pools[i]) - 1:
600 continue
601 indices[i] += 1
602 for j in range(i+1, n):
603 indices[j] = 0
604 yield tuple(pool[i] for pool, i in zip(pools, indices))
605 break
606 else:
607 return
608
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000609 def product2(*args, **kwds):
610 'Pure python version used in docs'
611 pools = map(tuple, args) * kwds.get('repeat', 1)
612 result = [[]]
613 for pool in pools:
614 result = [x+[y] for x in result for y in pool]
615 for prod in result:
616 yield tuple(prod)
617
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000618 argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
619 set('abcdefg'), range(11), tuple(range(13))]
620 for i in range(100):
621 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Raymond Hettingerd553d852008-03-04 04:17:08 +0000622 expected_len = prod(map(len, args))
623 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000624 self.assertEqual(list(product(*args)), list(product1(*args)))
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000625 self.assertEqual(list(product(*args)), list(product2(*args)))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000626 args = map(iter, args)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000627 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000628
Raymond Hettinger73d79632008-02-23 02:20:41 +0000629 # Test implementation detail: tuple re-use
630 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
631 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000632
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000633 def test_repeat(self):
634 self.assertEqual(zip(xrange(3),repeat('a')),
635 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000636 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000637 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000638 self.assertEqual(list(repeat('a', 0)), [])
639 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000640 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000641 self.assertRaises(TypeError, repeat, None, 3, 4)
642 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000643 r = repeat(1+0j)
644 self.assertEqual(repr(r), 'repeat((1+0j))')
645 r = repeat(1+0j, 5)
646 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
647 list(r)
648 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000649
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000650 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000651 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
652 [0**1, 1**2, 2**3])
653 self.assertEqual(list(imap(None, 'abc', range(5))),
654 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000655 self.assertEqual(list(imap(None, 'abc', count())),
656 [('a',0),('b',1),('c',2)])
657 self.assertEqual(take(2,imap(None, 'abc', count())),
658 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000659 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000660 self.assertRaises(TypeError, imap)
661 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000662 self.assertRaises(TypeError, imap(10, range(5)).next)
663 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
664 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000665
666 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000667 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
668 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000669 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
670 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000671 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger47317092008-01-17 03:02:14 +0000672 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
673 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000674 self.assertRaises(TypeError, starmap)
675 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
676 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
677 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
678 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000679
680 def test_islice(self):
681 for args in [ # islice(args) should agree with range(args)
682 (10, 20, 3),
683 (10, 3, 20),
684 (10, 20),
685 (10, 3),
686 (20,)
687 ]:
688 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
689
690 for args, tgtargs in [ # Stop when seqn is exhausted
691 ((10, 110, 3), ((10, 100, 3))),
692 ((10, 110), ((10, 100))),
693 ((110,), (100,))
694 ]:
695 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
696
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000697 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000698 self.assertEqual(list(islice(xrange(10), None)), range(10))
Raymond Hettingerb2594052004-12-05 09:25:51 +0000699 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
700 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000701 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
702 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
703
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000704 # Test number of items consumed SF #1171417
705 it = iter(range(10))
706 self.assertEqual(list(islice(it, 3)), range(3))
707 self.assertEqual(list(it), range(3, 10))
708
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000709 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000710 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000711 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
712 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
713 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
714 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
715 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000716 self.assertRaises(ValueError, islice, xrange(10), 'a')
717 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
718 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
719 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
720 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +0000721 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000722
723 def test_takewhile(self):
724 data = [1, 3, 5, 20, 2, 4, 6, 8]
725 underten = lambda x: x<10
726 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000727 self.assertEqual(list(takewhile(underten, [])), [])
728 self.assertRaises(TypeError, takewhile)
729 self.assertRaises(TypeError, takewhile, operator.pow)
730 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
731 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
732 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000733 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
734 self.assertEqual(list(t), [1, 1, 1])
735 self.assertRaises(StopIteration, t.next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000736
737 def test_dropwhile(self):
738 data = [1, 3, 5, 20, 2, 4, 6, 8]
739 underten = lambda x: x<10
740 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000741 self.assertEqual(list(dropwhile(underten, [])), [])
742 self.assertRaises(TypeError, dropwhile)
743 self.assertRaises(TypeError, dropwhile, operator.pow)
744 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
745 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
746 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000747
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000748 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000749 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000750 def irange(n):
751 for i in xrange(n):
752 yield i
753
754 a, b = tee([]) # test empty iterator
755 self.assertEqual(list(a), [])
756 self.assertEqual(list(b), [])
757
758 a, b = tee(irange(n)) # test 100% interleaved
759 self.assertEqual(zip(a,b), zip(range(n),range(n)))
760
761 a, b = tee(irange(n)) # test 0% interleaved
762 self.assertEqual(list(a), range(n))
763 self.assertEqual(list(b), range(n))
764
765 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000766 for i in xrange(100):
767 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000768 del a
769 self.assertEqual(list(b), range(n))
770
771 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000772 for i in xrange(100):
773 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000774 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000775 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000776
777 for j in xrange(5): # test randomly interleaved
778 order = [0]*n + [1]*n
779 random.shuffle(order)
780 lists = ([], [])
781 its = tee(irange(n))
782 for i in order:
783 value = its[i].next()
784 lists[i].append(value)
785 self.assertEqual(lists[0], range(n))
786 self.assertEqual(lists[1], range(n))
787
Raymond Hettingerad983e72003-11-12 14:32:26 +0000788 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000789 self.assertRaises(TypeError, tee)
790 self.assertRaises(TypeError, tee, 3)
791 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000792 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000793
Raymond Hettingerad983e72003-11-12 14:32:26 +0000794 # tee object should be instantiable
795 a, b = tee('abc')
796 c = type(a)('def')
797 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000798
Raymond Hettingerad983e72003-11-12 14:32:26 +0000799 # test long-lagged and multi-way split
800 a, b, c = tee(xrange(2000), 3)
801 for i in xrange(100):
802 self.assertEqual(a.next(), i)
803 self.assertEqual(list(b), range(2000))
804 self.assertEqual([c.next(), c.next()], range(2))
805 self.assertEqual(list(a), range(100,2000))
806 self.assertEqual(list(c), range(2,2000))
807
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000808 # test values of n
809 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Neal Norwitz69e88972006-09-02 02:58:13 +0000810 self.assertRaises(ValueError, tee, [], -1)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000811 for n in xrange(5):
812 result = tee('abc', n)
813 self.assertEqual(type(result), tuple)
814 self.assertEqual(len(result), n)
815 self.assertEqual(map(list, result), [list('abc')]*n)
816
Raymond Hettingerad983e72003-11-12 14:32:26 +0000817 # tee pass-through to copyable iterator
818 a, b = tee('abc')
819 c, d = tee(a)
820 self.assert_(a is c)
821
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000822 # test tee_new
823 t1, t2 = tee('abc')
824 tnew = type(t1)
825 self.assertRaises(TypeError, tnew)
826 self.assertRaises(TypeError, tnew, 10)
827 t3 = tnew(t1)
828 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000829
Raymond Hettingera9f60922004-10-17 16:40:14 +0000830 # test that tee objects are weak referencable
831 a, b = tee(xrange(10))
832 p = proxy(a)
833 self.assertEqual(getattr(p, '__class__'), type(b))
834 del a
835 self.assertRaises(ReferenceError, getattr, p, '__class__')
836
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000837 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000838 self.assertRaises(StopIteration, izip().next)
839
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000840 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000841 self.assertRaises(StopIteration, f([]).next)
842 self.assertRaises(StopIteration, f(StopNow()).next)
843
844 self.assertRaises(StopIteration, islice([], None).next)
845 self.assertRaises(StopIteration, islice(StopNow(), None).next)
846
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000847 p, q = tee([])
848 self.assertRaises(StopIteration, p.next)
849 self.assertRaises(StopIteration, q.next)
850 p, q = tee(StopNow())
851 self.assertRaises(StopIteration, p.next)
852 self.assertRaises(StopIteration, q.next)
853
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000854 self.assertRaises(StopIteration, repeat(None, 0).next)
855
856 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
857 self.assertRaises(StopIteration, f(lambda x:x, []).next)
858 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
859
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000860class TestExamples(unittest.TestCase):
861
862 def test_chain(self):
863 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
864
865 def test_chain_from_iterable(self):
866 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
867
868 def test_combinations(self):
869 self.assertEqual(list(combinations('ABCD', 2)),
870 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
871 self.assertEqual(list(combinations(range(4), 3)),
872 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
873
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000874 def test_combinations_with_replacement(self):
875 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
876 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
877
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000878 def test_compress(self):
879 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
880
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000881 def test_count(self):
882 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
883
884 def test_cycle(self):
885 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
886
887 def test_dropwhile(self):
888 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
889
890 def test_groupby(self):
891 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
892 list('ABCDAB'))
893 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
894 [list('AAAA'), list('BBB'), list('CC'), list('D')])
895
896 def test_ifilter(self):
897 self.assertEqual(list(ifilter(lambda x: x%2, range(10))), [1,3,5,7,9])
898
899 def test_ifilterfalse(self):
900 self.assertEqual(list(ifilterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
901
902 def test_imap(self):
903 self.assertEqual(list(imap(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
904
905 def test_islice(self):
906 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
907 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
908 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
909 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
910
911 def test_izip(self):
912 self.assertEqual(list(izip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
913
914 def test_izip_longest(self):
915 self.assertEqual(list(izip_longest('ABCD', 'xy', fillvalue='-')),
916 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
917
918 def test_permutations(self):
919 self.assertEqual(list(permutations('ABCD', 2)),
920 map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split()))
921 self.assertEqual(list(permutations(range(3))),
922 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
923
924 def test_product(self):
925 self.assertEqual(list(product('ABCD', 'xy')),
926 map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split()))
927 self.assertEqual(list(product(range(2), repeat=3)),
928 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
929 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
930
931 def test_repeat(self):
932 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
933
934 def test_stapmap(self):
935 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
936 [32, 9, 1000])
937
938 def test_takewhile(self):
939 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
940
941
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000942class TestGC(unittest.TestCase):
943
944 def makecycle(self, iterator, container):
945 container.append(iterator)
946 iterator.next()
947 del container, iterator
948
949 def test_chain(self):
950 a = []
951 self.makecycle(chain(a), a)
952
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000953 def test_chain_from_iterable(self):
954 a = []
955 self.makecycle(chain.from_iterable([a]), a)
956
957 def test_combinations(self):
958 a = []
959 self.makecycle(combinations([1,2,a,3], 3), a)
960
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000961 def test_combinations_with_replacement(self):
962 a = []
963 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
964
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000965 def test_compress(self):
966 a = []
967 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
968
Raymond Hettingerb21d8102009-02-16 20:39:12 +0000969 def test_count(self):
970 a = []
971 Int = type('Int', (int,), dict(x=a))
972 self.makecycle(count(Int(0), Int(1)), a)
973
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000974 def test_cycle(self):
975 a = []
976 self.makecycle(cycle([a]*2), a)
977
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000978 def test_dropwhile(self):
979 a = []
980 self.makecycle(dropwhile(bool, [0, a, a]), a)
981
982 def test_groupby(self):
983 a = []
984 self.makecycle(groupby([a]*2, lambda x:x), a)
985
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000986 def test_issue2246(self):
987 # Issue 2246 -- the _grouper iterator was not included in GC
988 n = 10
989 keyfunc = lambda x: x
990 for i, j in groupby(xrange(n), key=keyfunc):
991 keyfunc.__dict__.setdefault('x',[]).append(j)
992
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000993 def test_ifilter(self):
994 a = []
995 self.makecycle(ifilter(lambda x:True, [a]*2), a)
996
997 def test_ifilterfalse(self):
998 a = []
999 self.makecycle(ifilterfalse(lambda x:False, a), a)
1000
1001 def test_izip(self):
1002 a = []
1003 self.makecycle(izip([a]*2, [a]*3), a)
1004
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001005 def test_izip_longest(self):
1006 a = []
1007 self.makecycle(izip_longest([a]*2, [a]*3), a)
1008 b = [a, None]
1009 self.makecycle(izip_longest([a]*2, [a]*3, fillvalue=b), a)
1010
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001011 def test_imap(self):
1012 a = []
1013 self.makecycle(imap(lambda x:x, [a]*2), a)
1014
1015 def test_islice(self):
1016 a = []
1017 self.makecycle(islice([a]*2, None), a)
1018
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001019 def test_permutations(self):
1020 a = []
1021 self.makecycle(permutations([1,2,a,3], 3), a)
1022
1023 def test_product(self):
1024 a = []
1025 self.makecycle(product([1,2,a,3], repeat=3), a)
1026
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001027 def test_repeat(self):
1028 a = []
1029 self.makecycle(repeat(a), a)
1030
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001031 def test_starmap(self):
1032 a = []
1033 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1034
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001035 def test_takewhile(self):
1036 a = []
1037 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1038
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001039def R(seqn):
1040 'Regular generator'
1041 for i in seqn:
1042 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001043
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001044class G:
1045 'Sequence using __getitem__'
1046 def __init__(self, seqn):
1047 self.seqn = seqn
1048 def __getitem__(self, i):
1049 return self.seqn[i]
1050
1051class I:
1052 'Sequence using iterator protocol'
1053 def __init__(self, seqn):
1054 self.seqn = seqn
1055 self.i = 0
1056 def __iter__(self):
1057 return self
1058 def next(self):
1059 if self.i >= len(self.seqn): raise StopIteration
1060 v = self.seqn[self.i]
1061 self.i += 1
1062 return v
1063
1064class Ig:
1065 'Sequence using iterator protocol defined with a generator'
1066 def __init__(self, seqn):
1067 self.seqn = seqn
1068 self.i = 0
1069 def __iter__(self):
1070 for val in self.seqn:
1071 yield val
1072
1073class X:
1074 'Missing __getitem__ and __iter__'
1075 def __init__(self, seqn):
1076 self.seqn = seqn
1077 self.i = 0
1078 def next(self):
1079 if self.i >= len(self.seqn): raise StopIteration
1080 v = self.seqn[self.i]
1081 self.i += 1
1082 return v
1083
1084class N:
1085 'Iterator missing next()'
1086 def __init__(self, seqn):
1087 self.seqn = seqn
1088 self.i = 0
1089 def __iter__(self):
1090 return self
1091
1092class E:
1093 'Test propagation of exceptions'
1094 def __init__(self, seqn):
1095 self.seqn = seqn
1096 self.i = 0
1097 def __iter__(self):
1098 return self
1099 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001100 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001101
1102class S:
1103 'Test immediate stop'
1104 def __init__(self, seqn):
1105 pass
1106 def __iter__(self):
1107 return self
1108 def next(self):
1109 raise StopIteration
1110
1111def L(seqn):
1112 'Test multiple tiers of iterators'
1113 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
1114
1115
1116class TestVariousIteratorArgs(unittest.TestCase):
1117
1118 def test_chain(self):
1119 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1120 for g in (G, I, Ig, S, L, R):
1121 self.assertEqual(list(chain(g(s))), list(g(s)))
1122 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001123 self.assertRaises(TypeError, list, chain(X(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001124 self.assertRaises(TypeError, list, chain(N(s)))
1125 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1126
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001127 def test_compress(self):
1128 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1129 n = len(s)
1130 for g in (G, I, Ig, S, L, R):
1131 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1132 self.assertRaises(TypeError, compress, X(s), repeat(1))
1133 self.assertRaises(TypeError, list, compress(N(s), repeat(1)))
1134 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1135
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001136 def test_product(self):
1137 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1138 self.assertRaises(TypeError, product, X(s))
1139 self.assertRaises(TypeError, product, N(s))
1140 self.assertRaises(ZeroDivisionError, product, E(s))
1141
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001142 def test_cycle(self):
1143 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1144 for g in (G, I, Ig, S, L, R):
1145 tgtlen = len(s) * 3
1146 expected = list(g(s))*3
1147 actual = list(islice(cycle(g(s)), tgtlen))
1148 self.assertEqual(actual, expected)
1149 self.assertRaises(TypeError, cycle, X(s))
1150 self.assertRaises(TypeError, list, cycle(N(s)))
1151 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1152
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001153 def test_groupby(self):
1154 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1155 for g in (G, I, Ig, S, L, R):
1156 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1157 self.assertRaises(TypeError, groupby, X(s))
1158 self.assertRaises(TypeError, list, groupby(N(s)))
1159 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1160
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001161 def test_ifilter(self):
1162 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1163 for g in (G, I, Ig, S, L, R):
1164 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
1165 self.assertRaises(TypeError, ifilter, isEven, X(s))
1166 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
1167 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
1168
1169 def test_ifilterfalse(self):
1170 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1171 for g in (G, I, Ig, S, L, R):
1172 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
1173 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
1174 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
1175 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
1176
1177 def test_izip(self):
1178 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1179 for g in (G, I, Ig, S, L, R):
1180 self.assertEqual(list(izip(g(s))), zip(g(s)))
1181 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
1182 self.assertRaises(TypeError, izip, X(s))
1183 self.assertRaises(TypeError, list, izip(N(s)))
1184 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
1185
Raymond Hettingerd36862c2007-02-21 05:20:38 +00001186 def test_iziplongest(self):
1187 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1188 for g in (G, I, Ig, S, L, R):
1189 self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
1190 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
1191 self.assertRaises(TypeError, izip_longest, X(s))
1192 self.assertRaises(TypeError, list, izip_longest(N(s)))
1193 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
1194
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001195 def test_imap(self):
1196 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1197 for g in (G, I, Ig, S, L, R):
1198 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
1199 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
1200 self.assertRaises(TypeError, imap, onearg, X(s))
1201 self.assertRaises(TypeError, list, imap(onearg, N(s)))
1202 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
1203
1204 def test_islice(self):
1205 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1206 for g in (G, I, Ig, S, L, R):
1207 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1208 self.assertRaises(TypeError, islice, X(s), 10)
1209 self.assertRaises(TypeError, list, islice(N(s), 10))
1210 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1211
1212 def test_starmap(self):
1213 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1214 for g in (G, I, Ig, S, L, R):
1215 ss = zip(s, s)
1216 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
1217 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
1218 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
1219 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1220
1221 def test_takewhile(self):
1222 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1223 for g in (G, I, Ig, S, L, R):
1224 tgt = []
1225 for elem in g(s):
1226 if not isEven(elem): break
1227 tgt.append(elem)
1228 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1229 self.assertRaises(TypeError, takewhile, isEven, X(s))
1230 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
1231 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1232
1233 def test_dropwhile(self):
1234 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1235 for g in (G, I, Ig, S, L, R):
1236 tgt = []
1237 for elem in g(s):
1238 if not tgt and isOdd(elem): continue
1239 tgt.append(elem)
1240 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1241 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
1242 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
1243 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1244
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001245 def test_tee(self):
1246 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1247 for g in (G, I, Ig, S, L, R):
1248 it1, it2 = tee(g(s))
1249 self.assertEqual(list(it1), list(g(s)))
1250 self.assertEqual(list(it2), list(g(s)))
1251 self.assertRaises(TypeError, tee, X(s))
1252 self.assertRaises(TypeError, list, tee(N(s))[0])
1253 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1254
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001255class LengthTransparency(unittest.TestCase):
1256
1257 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001258 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001259 self.assertEqual(len(repeat(None, 50)), 50)
1260 self.assertRaises(TypeError, len, repeat(None))
1261
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001262class RegressionTests(unittest.TestCase):
1263
1264 def test_sf_793826(self):
1265 # Fix Armin Rigo's successful efforts to wreak havoc
1266
1267 def mutatingtuple(tuple1, f, tuple2):
1268 # this builds a tuple t which is a copy of tuple1,
1269 # then calls f(t), then mutates t to be equal to tuple2
1270 # (needs len(tuple1) == len(tuple2)).
1271 def g(value, first=[1]):
1272 if first:
1273 del first[:]
1274 f(z.next())
1275 return value
1276 items = list(tuple2)
1277 items[1:1] = list(tuple1)
1278 gen = imap(g, items)
1279 z = izip(*[gen]*len(tuple1))
1280 z.next()
1281
1282 def f(t):
1283 global T
1284 T = t
1285 first[:] = list(T)
1286
1287 first = []
1288 mutatingtuple((1,2,3), f, (4,5,6))
1289 second = list(T)
1290 self.assertEqual(first, second)
1291
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001292
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001293 def test_sf_950057(self):
1294 # Make sure that chain() and cycle() catch exceptions immediately
1295 # rather than when shifting between input sources
1296
1297 def gen1():
1298 hist.append(0)
1299 yield 1
1300 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001301 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001302 hist.append(2)
1303
1304 def gen2(x):
1305 hist.append(3)
1306 yield 2
1307 hist.append(4)
1308 if x:
1309 raise StopIteration
1310
1311 hist = []
1312 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1313 self.assertEqual(hist, [0,1])
1314
1315 hist = []
1316 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1317 self.assertEqual(hist, [0,1])
1318
1319 hist = []
1320 self.assertRaises(AssertionError, list, cycle(gen1()))
1321 self.assertEqual(hist, [0,1])
1322
Georg Brandlb84c1372007-01-21 10:28:43 +00001323class SubclassWithKwargsTest(unittest.TestCase):
1324 def test_keywords_in_subclass(self):
1325 # count is not subclassable...
1326 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001327 starmap, islice, takewhile, dropwhile, cycle, compress):
Georg Brandlb84c1372007-01-21 10:28:43 +00001328 class Subclass(cls):
1329 def __init__(self, newarg=None, *args):
1330 cls.__init__(self, *args)
1331 try:
1332 Subclass(newarg=1)
1333 except TypeError, err:
1334 # we expect type errors because of wrong argument count
1335 self.failIf("does not take keyword arguments" in err.args[0])
1336
1337
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001338libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001339
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001340
1341>>> amounts = [120.15, 764.05, 823.14]
1342>>> for checknum, amount in izip(count(1200), amounts):
1343... print 'Check %d is for $%.2f' % (checknum, amount)
1344...
1345Check 1200 is for $120.15
1346Check 1201 is for $764.05
1347Check 1202 is for $823.14
1348
1349>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001350>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1351... print cube
1352...
13531
13548
135527
1356
1357>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001358>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001359... print name.title()
1360...
1361Alex
1362Laura
1363Martin
1364Walter
1365Samuele
1366
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001367>>> from operator import itemgetter
1368>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Armin Rigoa3f09272006-05-28 19:13:17 +00001369>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001370>>> for k, g in groupby(di, itemgetter(1)):
1371... print k, map(itemgetter(0), g)
1372...
13731 ['a', 'c', 'e']
13742 ['b', 'd', 'f']
13753 ['g']
1376
Raymond Hettinger734fb572004-01-20 20:04:40 +00001377# Find runs of consecutive numbers using groupby. The key to the solution
1378# is differencing with a range so that consecutive numbers all appear in
1379# same group.
1380>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
1381>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
1382... print map(operator.itemgetter(1), g)
1383...
1384[1]
1385[4, 5, 6]
1386[10]
1387[15, 16, 17, 18]
1388[22]
1389[25, 26, 27, 28]
1390
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001391>>> def take(n, iterable):
1392... "Return first n items of the iterable as a list"
1393... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001394
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001395>>> def enumerate(iterable, start=0):
1396... return izip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001397
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001398>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001399... "Return function(0), function(1), ..."
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001400... return imap(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001401
1402>>> def nth(iterable, n):
Raymond Hettingerd507afd2009-02-04 10:52:32 +00001403... "Returns the nth item or None"
1404... return next(islice(iterable, n, None), None)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001405
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001406>>> def quantify(iterable, pred=bool):
1407... "Count how many times the predicate is true"
1408... return sum(imap(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001409
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001410>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001411... "Returns the sequence elements and then returns None indefinitely"
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001412... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001413
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001414>>> def ncycles(iterable, n):
1415... "Returns the seqeuence elements n times"
1416... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001417
1418>>> def dotproduct(vec1, vec2):
1419... return sum(imap(operator.mul, vec1, vec2))
1420
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001421>>> def flatten(listOfLists):
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001422... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001423
1424>>> def repeatfunc(func, times=None, *args):
1425... "Repeat calls to func with specified arguments."
1426... " Example: repeatfunc(random.random)"
1427... if times is None:
1428... return starmap(func, repeat(args))
1429... else:
1430... return starmap(func, repeat(args, times))
1431
Raymond Hettingerd591f662003-10-26 15:34:50 +00001432>>> def pairwise(iterable):
1433... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1434... a, b = tee(iterable)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001435... for elem in b:
1436... break
Raymond Hettingerad983e72003-11-12 14:32:26 +00001437... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001438
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001439>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001440... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001441... args = [iter(iterable)] * n
Raymond Hettingerf080e6d2008-07-31 01:19:50 +00001442... return izip_longest(fillvalue=fillvalue, *args)
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001443
1444>>> def roundrobin(*iterables):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001445... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001446... # Recipe credited to George Sakkis
1447... pending = len(iterables)
1448... nexts = cycle(iter(it).next for it in iterables)
1449... while pending:
1450... try:
1451... for next in nexts:
1452... yield next()
1453... except StopIteration:
1454... pending -= 1
1455... nexts = cycle(islice(nexts, pending))
1456
1457>>> def powerset(iterable):
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001458... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1459... s = list(iterable)
1460... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001461
Raymond Hettinger44e15812009-01-02 21:26:45 +00001462>>> def unique_everseen(iterable, key=None):
1463... "List unique elements, preserving order. Remember all elements ever seen."
1464... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1465... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1466... seen = set()
1467... seen_add = seen.add
1468... if key is None:
1469... for element in iterable:
1470... if element not in seen:
1471... seen_add(element)
1472... yield element
1473... else:
1474... for element in iterable:
1475... k = key(element)
1476... if k not in seen:
1477... seen_add(k)
1478... yield element
1479
1480>>> def unique_justseen(iterable, key=None):
1481... "List unique elements, preserving order. Remember only the element just seen."
1482... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1483... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1484... return imap(next, imap(itemgetter(1), groupby(iterable, key)))
1485
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001486This is not part of the examples but it tests to make sure the definitions
1487perform as purported.
1488
Raymond Hettingera098b332003-09-08 23:58:40 +00001489>>> take(10, count())
1490[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1491
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001492>>> list(enumerate('abc'))
1493[(0, 'a'), (1, 'b'), (2, 'c')]
1494
1495>>> list(islice(tabulate(lambda x: 2*x), 4))
1496[0, 2, 4, 6]
1497
1498>>> nth('abcde', 3)
Raymond Hettingerd507afd2009-02-04 10:52:32 +00001499'd'
1500
1501>>> nth('abcde', 9) is None
1502True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001503
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001504>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000150550
1506
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001507>>> a = [[1, 2, 3], [4, 5, 6]]
1508>>> flatten(a)
1509[1, 2, 3, 4, 5, 6]
1510
1511>>> list(repeatfunc(pow, 5, 2, 3))
1512[8, 8, 8, 8, 8]
1513
1514>>> import random
1515>>> take(5, imap(int, repeatfunc(random.random)))
1516[0, 0, 0, 0, 0]
1517
Raymond Hettingerd591f662003-10-26 15:34:50 +00001518>>> list(pairwise('abcd'))
1519[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001520
Raymond Hettingerd591f662003-10-26 15:34:50 +00001521>>> list(pairwise([]))
1522[]
1523
1524>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001525[]
1526
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001527>>> list(islice(padnone('abc'), 0, 6))
1528['a', 'b', 'c', None, None, None]
1529
1530>>> list(ncycles('abc', 3))
1531['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1532
1533>>> dotproduct([1,2,3], [4,5,6])
153432
1535
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001536>>> list(grouper(3, 'abcdefg', 'x'))
1537[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1538
1539>>> list(roundrobin('abc', 'd', 'ef'))
1540['a', 'd', 'e', 'b', 'f', 'c']
1541
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001542>>> list(powerset([1,2,3]))
1543[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001544
Raymond Hettinger560f9a82009-01-27 13:26:35 +00001545>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1546True
1547
1548>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1549True
1550
Raymond Hettinger44e15812009-01-02 21:26:45 +00001551>>> list(unique_everseen('AAAABBBCCDAABBB'))
1552['A', 'B', 'C', 'D']
1553
1554>>> list(unique_everseen('ABBCcAD', str.lower))
1555['A', 'B', 'C', 'D']
1556
1557>>> list(unique_justseen('AAAABBBCCDAABBB'))
1558['A', 'B', 'C', 'D', 'A', 'B']
1559
1560>>> list(unique_justseen('ABBCcAD', str.lower))
1561['A', 'B', 'C', 'A', 'D']
1562
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001563"""
1564
1565__test__ = {'libreftest' : libreftest}
1566
1567def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001568 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Georg Brandlb84c1372007-01-21 10:28:43 +00001569 RegressionTests, LengthTransparency,
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001570 SubclassWithKwargsTest, TestExamples)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001571 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001572
1573 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001574 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001575 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001576 counts = [None] * 5
1577 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001578 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001579 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001580 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001581 print counts
1582
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001583 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001584 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001585
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001586if __name__ == "__main__":
1587 test_main(verbose=True)