blob: ce5c898788777e185000ffea5ca6ed0029e28172 [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 Hettinger6a5b0272003-10-24 08:45:23 +0000969 def test_cycle(self):
970 a = []
971 self.makecycle(cycle([a]*2), a)
972
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000973 def test_dropwhile(self):
974 a = []
975 self.makecycle(dropwhile(bool, [0, a, a]), a)
976
977 def test_groupby(self):
978 a = []
979 self.makecycle(groupby([a]*2, lambda x:x), a)
980
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000981 def test_issue2246(self):
982 # Issue 2246 -- the _grouper iterator was not included in GC
983 n = 10
984 keyfunc = lambda x: x
985 for i, j in groupby(xrange(n), key=keyfunc):
986 keyfunc.__dict__.setdefault('x',[]).append(j)
987
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000988 def test_ifilter(self):
989 a = []
990 self.makecycle(ifilter(lambda x:True, [a]*2), a)
991
992 def test_ifilterfalse(self):
993 a = []
994 self.makecycle(ifilterfalse(lambda x:False, a), a)
995
996 def test_izip(self):
997 a = []
998 self.makecycle(izip([a]*2, [a]*3), a)
999
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001000 def test_izip_longest(self):
1001 a = []
1002 self.makecycle(izip_longest([a]*2, [a]*3), a)
1003 b = [a, None]
1004 self.makecycle(izip_longest([a]*2, [a]*3, fillvalue=b), a)
1005
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001006 def test_imap(self):
1007 a = []
1008 self.makecycle(imap(lambda x:x, [a]*2), a)
1009
1010 def test_islice(self):
1011 a = []
1012 self.makecycle(islice([a]*2, None), a)
1013
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001014 def test_permutations(self):
1015 a = []
1016 self.makecycle(permutations([1,2,a,3], 3), a)
1017
1018 def test_product(self):
1019 a = []
1020 self.makecycle(product([1,2,a,3], repeat=3), a)
1021
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001022 def test_repeat(self):
1023 a = []
1024 self.makecycle(repeat(a), a)
1025
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001026 def test_starmap(self):
1027 a = []
1028 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1029
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001030 def test_takewhile(self):
1031 a = []
1032 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1033
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001034def R(seqn):
1035 'Regular generator'
1036 for i in seqn:
1037 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001038
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001039class G:
1040 'Sequence using __getitem__'
1041 def __init__(self, seqn):
1042 self.seqn = seqn
1043 def __getitem__(self, i):
1044 return self.seqn[i]
1045
1046class I:
1047 'Sequence using iterator protocol'
1048 def __init__(self, seqn):
1049 self.seqn = seqn
1050 self.i = 0
1051 def __iter__(self):
1052 return self
1053 def next(self):
1054 if self.i >= len(self.seqn): raise StopIteration
1055 v = self.seqn[self.i]
1056 self.i += 1
1057 return v
1058
1059class Ig:
1060 'Sequence using iterator protocol defined with a generator'
1061 def __init__(self, seqn):
1062 self.seqn = seqn
1063 self.i = 0
1064 def __iter__(self):
1065 for val in self.seqn:
1066 yield val
1067
1068class X:
1069 'Missing __getitem__ and __iter__'
1070 def __init__(self, seqn):
1071 self.seqn = seqn
1072 self.i = 0
1073 def next(self):
1074 if self.i >= len(self.seqn): raise StopIteration
1075 v = self.seqn[self.i]
1076 self.i += 1
1077 return v
1078
1079class N:
1080 'Iterator missing next()'
1081 def __init__(self, seqn):
1082 self.seqn = seqn
1083 self.i = 0
1084 def __iter__(self):
1085 return self
1086
1087class E:
1088 'Test propagation of exceptions'
1089 def __init__(self, seqn):
1090 self.seqn = seqn
1091 self.i = 0
1092 def __iter__(self):
1093 return self
1094 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001095 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001096
1097class S:
1098 'Test immediate stop'
1099 def __init__(self, seqn):
1100 pass
1101 def __iter__(self):
1102 return self
1103 def next(self):
1104 raise StopIteration
1105
1106def L(seqn):
1107 'Test multiple tiers of iterators'
1108 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
1109
1110
1111class TestVariousIteratorArgs(unittest.TestCase):
1112
1113 def test_chain(self):
1114 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1115 for g in (G, I, Ig, S, L, R):
1116 self.assertEqual(list(chain(g(s))), list(g(s)))
1117 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001118 self.assertRaises(TypeError, list, chain(X(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001119 self.assertRaises(TypeError, list, chain(N(s)))
1120 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1121
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001122 def test_compress(self):
1123 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1124 n = len(s)
1125 for g in (G, I, Ig, S, L, R):
1126 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1127 self.assertRaises(TypeError, compress, X(s), repeat(1))
1128 self.assertRaises(TypeError, list, compress(N(s), repeat(1)))
1129 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1130
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001131 def test_product(self):
1132 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1133 self.assertRaises(TypeError, product, X(s))
1134 self.assertRaises(TypeError, product, N(s))
1135 self.assertRaises(ZeroDivisionError, product, E(s))
1136
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001137 def test_cycle(self):
1138 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1139 for g in (G, I, Ig, S, L, R):
1140 tgtlen = len(s) * 3
1141 expected = list(g(s))*3
1142 actual = list(islice(cycle(g(s)), tgtlen))
1143 self.assertEqual(actual, expected)
1144 self.assertRaises(TypeError, cycle, X(s))
1145 self.assertRaises(TypeError, list, cycle(N(s)))
1146 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1147
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001148 def test_groupby(self):
1149 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1150 for g in (G, I, Ig, S, L, R):
1151 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1152 self.assertRaises(TypeError, groupby, X(s))
1153 self.assertRaises(TypeError, list, groupby(N(s)))
1154 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1155
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001156 def test_ifilter(self):
1157 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1158 for g in (G, I, Ig, S, L, R):
1159 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
1160 self.assertRaises(TypeError, ifilter, isEven, X(s))
1161 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
1162 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
1163
1164 def test_ifilterfalse(self):
1165 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1166 for g in (G, I, Ig, S, L, R):
1167 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
1168 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
1169 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
1170 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
1171
1172 def test_izip(self):
1173 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1174 for g in (G, I, Ig, S, L, R):
1175 self.assertEqual(list(izip(g(s))), zip(g(s)))
1176 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
1177 self.assertRaises(TypeError, izip, X(s))
1178 self.assertRaises(TypeError, list, izip(N(s)))
1179 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
1180
Raymond Hettingerd36862c2007-02-21 05:20:38 +00001181 def test_iziplongest(self):
1182 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1183 for g in (G, I, Ig, S, L, R):
1184 self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
1185 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
1186 self.assertRaises(TypeError, izip_longest, X(s))
1187 self.assertRaises(TypeError, list, izip_longest(N(s)))
1188 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
1189
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001190 def test_imap(self):
1191 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1192 for g in (G, I, Ig, S, L, R):
1193 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
1194 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
1195 self.assertRaises(TypeError, imap, onearg, X(s))
1196 self.assertRaises(TypeError, list, imap(onearg, N(s)))
1197 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
1198
1199 def test_islice(self):
1200 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1201 for g in (G, I, Ig, S, L, R):
1202 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1203 self.assertRaises(TypeError, islice, X(s), 10)
1204 self.assertRaises(TypeError, list, islice(N(s), 10))
1205 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1206
1207 def test_starmap(self):
1208 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1209 for g in (G, I, Ig, S, L, R):
1210 ss = zip(s, s)
1211 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
1212 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
1213 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
1214 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1215
1216 def test_takewhile(self):
1217 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1218 for g in (G, I, Ig, S, L, R):
1219 tgt = []
1220 for elem in g(s):
1221 if not isEven(elem): break
1222 tgt.append(elem)
1223 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1224 self.assertRaises(TypeError, takewhile, isEven, X(s))
1225 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
1226 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1227
1228 def test_dropwhile(self):
1229 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1230 for g in (G, I, Ig, S, L, R):
1231 tgt = []
1232 for elem in g(s):
1233 if not tgt and isOdd(elem): continue
1234 tgt.append(elem)
1235 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1236 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
1237 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
1238 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1239
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001240 def test_tee(self):
1241 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1242 for g in (G, I, Ig, S, L, R):
1243 it1, it2 = tee(g(s))
1244 self.assertEqual(list(it1), list(g(s)))
1245 self.assertEqual(list(it2), list(g(s)))
1246 self.assertRaises(TypeError, tee, X(s))
1247 self.assertRaises(TypeError, list, tee(N(s))[0])
1248 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1249
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001250class LengthTransparency(unittest.TestCase):
1251
1252 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001253 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001254 self.assertEqual(len(repeat(None, 50)), 50)
1255 self.assertRaises(TypeError, len, repeat(None))
1256
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001257class RegressionTests(unittest.TestCase):
1258
1259 def test_sf_793826(self):
1260 # Fix Armin Rigo's successful efforts to wreak havoc
1261
1262 def mutatingtuple(tuple1, f, tuple2):
1263 # this builds a tuple t which is a copy of tuple1,
1264 # then calls f(t), then mutates t to be equal to tuple2
1265 # (needs len(tuple1) == len(tuple2)).
1266 def g(value, first=[1]):
1267 if first:
1268 del first[:]
1269 f(z.next())
1270 return value
1271 items = list(tuple2)
1272 items[1:1] = list(tuple1)
1273 gen = imap(g, items)
1274 z = izip(*[gen]*len(tuple1))
1275 z.next()
1276
1277 def f(t):
1278 global T
1279 T = t
1280 first[:] = list(T)
1281
1282 first = []
1283 mutatingtuple((1,2,3), f, (4,5,6))
1284 second = list(T)
1285 self.assertEqual(first, second)
1286
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001287
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001288 def test_sf_950057(self):
1289 # Make sure that chain() and cycle() catch exceptions immediately
1290 # rather than when shifting between input sources
1291
1292 def gen1():
1293 hist.append(0)
1294 yield 1
1295 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001296 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001297 hist.append(2)
1298
1299 def gen2(x):
1300 hist.append(3)
1301 yield 2
1302 hist.append(4)
1303 if x:
1304 raise StopIteration
1305
1306 hist = []
1307 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1308 self.assertEqual(hist, [0,1])
1309
1310 hist = []
1311 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1312 self.assertEqual(hist, [0,1])
1313
1314 hist = []
1315 self.assertRaises(AssertionError, list, cycle(gen1()))
1316 self.assertEqual(hist, [0,1])
1317
Georg Brandlb84c1372007-01-21 10:28:43 +00001318class SubclassWithKwargsTest(unittest.TestCase):
1319 def test_keywords_in_subclass(self):
1320 # count is not subclassable...
1321 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001322 starmap, islice, takewhile, dropwhile, cycle, compress):
Georg Brandlb84c1372007-01-21 10:28:43 +00001323 class Subclass(cls):
1324 def __init__(self, newarg=None, *args):
1325 cls.__init__(self, *args)
1326 try:
1327 Subclass(newarg=1)
1328 except TypeError, err:
1329 # we expect type errors because of wrong argument count
1330 self.failIf("does not take keyword arguments" in err.args[0])
1331
1332
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001333libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001334
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001335
1336>>> amounts = [120.15, 764.05, 823.14]
1337>>> for checknum, amount in izip(count(1200), amounts):
1338... print 'Check %d is for $%.2f' % (checknum, amount)
1339...
1340Check 1200 is for $120.15
1341Check 1201 is for $764.05
1342Check 1202 is for $823.14
1343
1344>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001345>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1346... print cube
1347...
13481
13498
135027
1351
1352>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001353>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001354... print name.title()
1355...
1356Alex
1357Laura
1358Martin
1359Walter
1360Samuele
1361
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001362>>> from operator import itemgetter
1363>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Armin Rigoa3f09272006-05-28 19:13:17 +00001364>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001365>>> for k, g in groupby(di, itemgetter(1)):
1366... print k, map(itemgetter(0), g)
1367...
13681 ['a', 'c', 'e']
13692 ['b', 'd', 'f']
13703 ['g']
1371
Raymond Hettinger734fb572004-01-20 20:04:40 +00001372# Find runs of consecutive numbers using groupby. The key to the solution
1373# is differencing with a range so that consecutive numbers all appear in
1374# same group.
1375>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
1376>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
1377... print map(operator.itemgetter(1), g)
1378...
1379[1]
1380[4, 5, 6]
1381[10]
1382[15, 16, 17, 18]
1383[22]
1384[25, 26, 27, 28]
1385
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001386>>> def take(n, iterable):
1387... "Return first n items of the iterable as a list"
1388... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001389
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001390>>> def enumerate(iterable, start=0):
1391... return izip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001392
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001393>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001394... "Return function(0), function(1), ..."
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001395... return imap(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001396
1397>>> def nth(iterable, n):
Raymond Hettingerd507afd2009-02-04 10:52:32 +00001398... "Returns the nth item or None"
1399... return next(islice(iterable, n, None), None)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001400
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001401>>> def quantify(iterable, pred=bool):
1402... "Count how many times the predicate is true"
1403... return sum(imap(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001404
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001405>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001406... "Returns the sequence elements and then returns None indefinitely"
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001407... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001408
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001409>>> def ncycles(iterable, n):
1410... "Returns the seqeuence elements n times"
1411... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001412
1413>>> def dotproduct(vec1, vec2):
1414... return sum(imap(operator.mul, vec1, vec2))
1415
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001416>>> def flatten(listOfLists):
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001417... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001418
1419>>> def repeatfunc(func, times=None, *args):
1420... "Repeat calls to func with specified arguments."
1421... " Example: repeatfunc(random.random)"
1422... if times is None:
1423... return starmap(func, repeat(args))
1424... else:
1425... return starmap(func, repeat(args, times))
1426
Raymond Hettingerd591f662003-10-26 15:34:50 +00001427>>> def pairwise(iterable):
1428... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1429... a, b = tee(iterable)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001430... for elem in b:
1431... break
Raymond Hettingerad983e72003-11-12 14:32:26 +00001432... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001433
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001434>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001435... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001436... args = [iter(iterable)] * n
Raymond Hettingerf080e6d2008-07-31 01:19:50 +00001437... return izip_longest(fillvalue=fillvalue, *args)
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001438
1439>>> def roundrobin(*iterables):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001440... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001441... # Recipe credited to George Sakkis
1442... pending = len(iterables)
1443... nexts = cycle(iter(it).next for it in iterables)
1444... while pending:
1445... try:
1446... for next in nexts:
1447... yield next()
1448... except StopIteration:
1449... pending -= 1
1450... nexts = cycle(islice(nexts, pending))
1451
1452>>> def powerset(iterable):
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001453... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1454... s = list(iterable)
1455... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001456
Raymond Hettinger44e15812009-01-02 21:26:45 +00001457>>> def unique_everseen(iterable, key=None):
1458... "List unique elements, preserving order. Remember all elements ever seen."
1459... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1460... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1461... seen = set()
1462... seen_add = seen.add
1463... if key is None:
1464... for element in iterable:
1465... if element not in seen:
1466... seen_add(element)
1467... yield element
1468... else:
1469... for element in iterable:
1470... k = key(element)
1471... if k not in seen:
1472... seen_add(k)
1473... yield element
1474
1475>>> def unique_justseen(iterable, key=None):
1476... "List unique elements, preserving order. Remember only the element just seen."
1477... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1478... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1479... return imap(next, imap(itemgetter(1), groupby(iterable, key)))
1480
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001481This is not part of the examples but it tests to make sure the definitions
1482perform as purported.
1483
Raymond Hettingera098b332003-09-08 23:58:40 +00001484>>> take(10, count())
1485[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1486
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001487>>> list(enumerate('abc'))
1488[(0, 'a'), (1, 'b'), (2, 'c')]
1489
1490>>> list(islice(tabulate(lambda x: 2*x), 4))
1491[0, 2, 4, 6]
1492
1493>>> nth('abcde', 3)
Raymond Hettingerd507afd2009-02-04 10:52:32 +00001494'd'
1495
1496>>> nth('abcde', 9) is None
1497True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001498
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001499>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000150050
1501
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001502>>> a = [[1, 2, 3], [4, 5, 6]]
1503>>> flatten(a)
1504[1, 2, 3, 4, 5, 6]
1505
1506>>> list(repeatfunc(pow, 5, 2, 3))
1507[8, 8, 8, 8, 8]
1508
1509>>> import random
1510>>> take(5, imap(int, repeatfunc(random.random)))
1511[0, 0, 0, 0, 0]
1512
Raymond Hettingerd591f662003-10-26 15:34:50 +00001513>>> list(pairwise('abcd'))
1514[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001515
Raymond Hettingerd591f662003-10-26 15:34:50 +00001516>>> list(pairwise([]))
1517[]
1518
1519>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001520[]
1521
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001522>>> list(islice(padnone('abc'), 0, 6))
1523['a', 'b', 'c', None, None, None]
1524
1525>>> list(ncycles('abc', 3))
1526['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1527
1528>>> dotproduct([1,2,3], [4,5,6])
152932
1530
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001531>>> list(grouper(3, 'abcdefg', 'x'))
1532[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1533
1534>>> list(roundrobin('abc', 'd', 'ef'))
1535['a', 'd', 'e', 'b', 'f', 'c']
1536
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001537>>> list(powerset([1,2,3]))
1538[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001539
Raymond Hettinger560f9a82009-01-27 13:26:35 +00001540>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1541True
1542
1543>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1544True
1545
Raymond Hettinger44e15812009-01-02 21:26:45 +00001546>>> list(unique_everseen('AAAABBBCCDAABBB'))
1547['A', 'B', 'C', 'D']
1548
1549>>> list(unique_everseen('ABBCcAD', str.lower))
1550['A', 'B', 'C', 'D']
1551
1552>>> list(unique_justseen('AAAABBBCCDAABBB'))
1553['A', 'B', 'C', 'D', 'A', 'B']
1554
1555>>> list(unique_justseen('ABBCcAD', str.lower))
1556['A', 'B', 'C', 'A', 'D']
1557
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001558"""
1559
1560__test__ = {'libreftest' : libreftest}
1561
1562def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001563 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Georg Brandlb84c1372007-01-21 10:28:43 +00001564 RegressionTests, LengthTransparency,
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001565 SubclassWithKwargsTest, TestExamples)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001566 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001567
1568 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001569 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001570 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001571 counts = [None] * 5
1572 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001573 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001574 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001575 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001576 print counts
1577
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001578 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001579 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001580
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001581if __name__ == "__main__":
1582 test_main(verbose=True)