blob: 043eaeee792d7ac62511dcfe427400aa398adb64 [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)])
350 self.assertEqual(zip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
351 self.assertEqual(zip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
352 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
353 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
354 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettingeraa044612009-02-12 12:04:26 +0000355 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
356 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerdbe3bfb2009-02-12 12:43:01 +0000357 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
358 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000359 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
360 c = count(3, 5)
361 self.assertEqual(repr(c), 'count(3, 5)')
362 c.next()
363 self.assertEqual(repr(c), 'count(8, 5)')
364 c = count(-9, 0)
365 self.assertEqual(repr(c), 'count(-9, 0)')
366 c.next()
367 self.assertEqual(repr(c), 'count(-9, 0)')
368 c = count(-9, -3)
369 self.assertEqual(repr(c), 'count(-9, -3)')
370 c.next()
371 self.assertEqual(repr(c), 'count(-12, -3)')
372 self.assertEqual(repr(c), 'count(-12, -3)')
373 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
374 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
375 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
376 for i in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 10, sys.maxint-5, sys.maxint+5):
377 for j in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 1, 10, sys.maxint-5, sys.maxint+5):
378 # Test repr (ignoring the L in longs)
379 r1 = repr(count(i, j)).replace('L', '')
380 if j == 1:
381 r2 = ('count(%r)' % i).replace('L', '')
382 else:
383 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
384 self.assertEqual(r1, r2)
385
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000386 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000387 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000388 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000389 self.assertRaises(TypeError, cycle)
390 self.assertRaises(TypeError, cycle, 5)
391 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000392
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000393 def test_groupby(self):
394 # Check whether it accepts arguments correctly
395 self.assertEqual([], list(groupby([])))
396 self.assertEqual([], list(groupby([], key=id)))
397 self.assertRaises(TypeError, list, groupby('abc', []))
398 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000399 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000400
401 # Check normal input
402 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
403 (2,15,22), (3,16,23), (3,17,23)]
404 dup = []
405 for k, g in groupby(s, lambda r:r[0]):
406 for elem in g:
407 self.assertEqual(k, elem[0])
408 dup.append(elem)
409 self.assertEqual(s, dup)
410
411 # Check nested case
412 dup = []
413 for k, g in groupby(s, lambda r:r[0]):
414 for ik, ig in groupby(g, lambda r:r[2]):
415 for elem in ig:
416 self.assertEqual(k, elem[0])
417 self.assertEqual(ik, elem[2])
418 dup.append(elem)
419 self.assertEqual(s, dup)
420
421 # Check case where inner iterator is not used
422 keys = [k for k, g in groupby(s, lambda r:r[0])]
423 expectedkeys = set([r[0] for r in s])
424 self.assertEqual(set(keys), expectedkeys)
425 self.assertEqual(len(keys), len(expectedkeys))
426
427 # Exercise pipes and filters style
428 s = 'abracadabra'
429 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000430 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000431 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
432 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000433 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000434 self.assertEqual(r, ['a', 'b', 'r'])
435 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000436 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000437 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
438 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000439 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000440 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
441
442 # iter.next failure
443 class ExpectedError(Exception):
444 pass
445 def delayed_raise(n=0):
446 for i in range(n):
447 yield 'yo'
448 raise ExpectedError
449 def gulp(iterable, keyp=None, func=list):
450 return [func(g) for k, g in groupby(iterable, keyp)]
451
452 # iter.next failure on outer object
453 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
454 # iter.next failure on inner object
455 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
456
457 # __cmp__ failure
458 class DummyCmp:
459 def __cmp__(self, dst):
460 raise ExpectedError
461 s = [DummyCmp(), DummyCmp(), None]
462
463 # __cmp__ failure on outer object
464 self.assertRaises(ExpectedError, gulp, s, func=id)
465 # __cmp__ failure on inner object
466 self.assertRaises(ExpectedError, gulp, s)
467
468 # keyfunc failure
469 def keyfunc(obj):
470 if keyfunc.skip > 0:
471 keyfunc.skip -= 1
472 return obj
473 else:
474 raise ExpectedError
475
476 # keyfunc failure on outer object
477 keyfunc.skip = 0
478 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
479 keyfunc.skip = 1
480 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
481
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000482 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000483 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000484 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000485 self.assertEqual(list(ifilter(bool, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000486 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000487 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000488 self.assertRaises(TypeError, ifilter, lambda x:x)
489 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000490 self.assertRaises(TypeError, ifilter, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000491 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000492
493 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000494 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
495 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000496 self.assertEqual(list(ifilterfalse(bool, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000497 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000498 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000499 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
500 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000501 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000502 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000503
504 def test_izip(self):
505 ans = [(x,y) for x, y in izip('abc',count())]
506 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000507 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
508 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
Raymond Hettinger02420702003-06-29 20:36:23 +0000509 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000510 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000511 self.assertEqual(list(izip()), zip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000512 self.assertRaises(TypeError, izip, 3)
513 self.assertRaises(TypeError, izip, range(3), 3)
514 # Check tuple re-use (implementation detail)
515 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
516 zip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000517 self.assertEqual([pair for pair in izip('abc', 'def')],
518 zip('abc', 'def'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000519 ids = map(id, izip('abc', 'def'))
520 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000521 ids = map(id, list(izip('abc', 'def')))
522 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000523
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000524 def test_iziplongest(self):
525 for args in [
526 ['abc', range(6)],
527 [range(6), 'abc'],
528 [range(1000), range(2000,2100), range(3000,3050)],
529 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
530 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
531 ]:
532 target = map(None, *args)
533 self.assertEqual(list(izip_longest(*args)), target)
534 self.assertEqual(list(izip_longest(*args, **{})), target)
535 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
536 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Tim Petersea5962f2007-03-12 18:07:52 +0000537
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000538 self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
539
540 self.assertEqual(list(izip_longest()), zip())
541 self.assertEqual(list(izip_longest([])), zip([]))
542 self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
Tim Petersea5962f2007-03-12 18:07:52 +0000543
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000544 self.assertEqual(list(izip_longest('abc', 'defg', **{})), map(None, 'abc', 'defg')) # empty keyword dict
545 self.assertRaises(TypeError, izip_longest, 3)
546 self.assertRaises(TypeError, izip_longest, range(3), 3)
547
548 for stmt in [
549 "izip_longest('abc', fv=1)",
Tim Petersea5962f2007-03-12 18:07:52 +0000550 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000551 ]:
552 try:
553 eval(stmt, globals(), locals())
554 except TypeError:
555 pass
556 else:
557 self.fail('Did not raise Type in: ' + stmt)
Tim Petersea5962f2007-03-12 18:07:52 +0000558
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000559 # Check tuple re-use (implementation detail)
560 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
561 zip('abc', 'def'))
562 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
563 zip('abc', 'def'))
564 ids = map(id, izip_longest('abc', 'def'))
565 self.assertEqual(min(ids), max(ids))
566 ids = map(id, list(izip_longest('abc', 'def')))
567 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
568
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000569 def test_product(self):
570 for args, result in [
Raymond Hettingerd553d852008-03-04 04:17:08 +0000571 ([], [()]), # zero iterables
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000572 (['ab'], [('a',), ('b',)]), # one iterable
573 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
574 ([range(0), range(2), range(3)], []), # first iterable with zero length
575 ([range(2), range(0), range(3)], []), # middle iterable with zero length
576 ([range(2), range(3), range(0)], []), # last iterable with zero length
577 ]:
578 self.assertEqual(list(product(*args)), result)
Raymond Hettinger08ff6822008-02-29 02:21:48 +0000579 for r in range(4):
580 self.assertEqual(list(product(*(args*r))),
581 list(product(*args, **dict(repeat=r))))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000582 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
583 self.assertRaises(TypeError, product, range(6), None)
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000584
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000585 def product1(*args, **kwds):
586 pools = map(tuple, args) * kwds.get('repeat', 1)
587 n = len(pools)
588 if n == 0:
589 yield ()
590 return
591 if any(len(pool) == 0 for pool in pools):
592 return
593 indices = [0] * n
594 yield tuple(pool[i] for pool, i in zip(pools, indices))
595 while 1:
596 for i in reversed(range(n)): # right to left
597 if indices[i] == len(pools[i]) - 1:
598 continue
599 indices[i] += 1
600 for j in range(i+1, n):
601 indices[j] = 0
602 yield tuple(pool[i] for pool, i in zip(pools, indices))
603 break
604 else:
605 return
606
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000607 def product2(*args, **kwds):
608 'Pure python version used in docs'
609 pools = map(tuple, args) * kwds.get('repeat', 1)
610 result = [[]]
611 for pool in pools:
612 result = [x+[y] for x in result for y in pool]
613 for prod in result:
614 yield tuple(prod)
615
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000616 argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
617 set('abcdefg'), range(11), tuple(range(13))]
618 for i in range(100):
619 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Raymond Hettingerd553d852008-03-04 04:17:08 +0000620 expected_len = prod(map(len, args))
621 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000622 self.assertEqual(list(product(*args)), list(product1(*args)))
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000623 self.assertEqual(list(product(*args)), list(product2(*args)))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000624 args = map(iter, args)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000625 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000626
Raymond Hettinger73d79632008-02-23 02:20:41 +0000627 # Test implementation detail: tuple re-use
628 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
629 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000630
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000631 def test_repeat(self):
632 self.assertEqual(zip(xrange(3),repeat('a')),
633 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000634 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000635 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000636 self.assertEqual(list(repeat('a', 0)), [])
637 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000638 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000639 self.assertRaises(TypeError, repeat, None, 3, 4)
640 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000641 r = repeat(1+0j)
642 self.assertEqual(repr(r), 'repeat((1+0j))')
643 r = repeat(1+0j, 5)
644 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
645 list(r)
646 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000647
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000648 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000649 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
650 [0**1, 1**2, 2**3])
651 self.assertEqual(list(imap(None, 'abc', range(5))),
652 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000653 self.assertEqual(list(imap(None, 'abc', count())),
654 [('a',0),('b',1),('c',2)])
655 self.assertEqual(take(2,imap(None, 'abc', count())),
656 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000657 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000658 self.assertRaises(TypeError, imap)
659 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000660 self.assertRaises(TypeError, imap(10, range(5)).next)
661 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
662 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000663
664 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000665 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
666 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000667 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
668 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000669 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger47317092008-01-17 03:02:14 +0000670 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
671 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000672 self.assertRaises(TypeError, starmap)
673 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
674 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
675 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
676 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000677
678 def test_islice(self):
679 for args in [ # islice(args) should agree with range(args)
680 (10, 20, 3),
681 (10, 3, 20),
682 (10, 20),
683 (10, 3),
684 (20,)
685 ]:
686 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
687
688 for args, tgtargs in [ # Stop when seqn is exhausted
689 ((10, 110, 3), ((10, 100, 3))),
690 ((10, 110), ((10, 100))),
691 ((110,), (100,))
692 ]:
693 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
694
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000695 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000696 self.assertEqual(list(islice(xrange(10), None)), range(10))
Raymond Hettingerb2594052004-12-05 09:25:51 +0000697 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
698 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000699 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
700 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
701
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000702 # Test number of items consumed SF #1171417
703 it = iter(range(10))
704 self.assertEqual(list(islice(it, 3)), range(3))
705 self.assertEqual(list(it), range(3, 10))
706
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000707 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000708 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000709 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
710 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
711 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
712 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
713 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000714 self.assertRaises(ValueError, islice, xrange(10), 'a')
715 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
716 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
717 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
718 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +0000719 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000720
721 def test_takewhile(self):
722 data = [1, 3, 5, 20, 2, 4, 6, 8]
723 underten = lambda x: x<10
724 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000725 self.assertEqual(list(takewhile(underten, [])), [])
726 self.assertRaises(TypeError, takewhile)
727 self.assertRaises(TypeError, takewhile, operator.pow)
728 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
729 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
730 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000731 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
732 self.assertEqual(list(t), [1, 1, 1])
733 self.assertRaises(StopIteration, t.next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000734
735 def test_dropwhile(self):
736 data = [1, 3, 5, 20, 2, 4, 6, 8]
737 underten = lambda x: x<10
738 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000739 self.assertEqual(list(dropwhile(underten, [])), [])
740 self.assertRaises(TypeError, dropwhile)
741 self.assertRaises(TypeError, dropwhile, operator.pow)
742 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
743 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
744 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000745
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000746 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000747 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000748 def irange(n):
749 for i in xrange(n):
750 yield i
751
752 a, b = tee([]) # test empty iterator
753 self.assertEqual(list(a), [])
754 self.assertEqual(list(b), [])
755
756 a, b = tee(irange(n)) # test 100% interleaved
757 self.assertEqual(zip(a,b), zip(range(n),range(n)))
758
759 a, b = tee(irange(n)) # test 0% interleaved
760 self.assertEqual(list(a), range(n))
761 self.assertEqual(list(b), range(n))
762
763 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000764 for i in xrange(100):
765 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000766 del a
767 self.assertEqual(list(b), range(n))
768
769 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000770 for i in xrange(100):
771 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000772 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000773 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000774
775 for j in xrange(5): # test randomly interleaved
776 order = [0]*n + [1]*n
777 random.shuffle(order)
778 lists = ([], [])
779 its = tee(irange(n))
780 for i in order:
781 value = its[i].next()
782 lists[i].append(value)
783 self.assertEqual(lists[0], range(n))
784 self.assertEqual(lists[1], range(n))
785
Raymond Hettingerad983e72003-11-12 14:32:26 +0000786 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000787 self.assertRaises(TypeError, tee)
788 self.assertRaises(TypeError, tee, 3)
789 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000790 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000791
Raymond Hettingerad983e72003-11-12 14:32:26 +0000792 # tee object should be instantiable
793 a, b = tee('abc')
794 c = type(a)('def')
795 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000796
Raymond Hettingerad983e72003-11-12 14:32:26 +0000797 # test long-lagged and multi-way split
798 a, b, c = tee(xrange(2000), 3)
799 for i in xrange(100):
800 self.assertEqual(a.next(), i)
801 self.assertEqual(list(b), range(2000))
802 self.assertEqual([c.next(), c.next()], range(2))
803 self.assertEqual(list(a), range(100,2000))
804 self.assertEqual(list(c), range(2,2000))
805
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000806 # test values of n
807 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Neal Norwitz69e88972006-09-02 02:58:13 +0000808 self.assertRaises(ValueError, tee, [], -1)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000809 for n in xrange(5):
810 result = tee('abc', n)
811 self.assertEqual(type(result), tuple)
812 self.assertEqual(len(result), n)
813 self.assertEqual(map(list, result), [list('abc')]*n)
814
Raymond Hettingerad983e72003-11-12 14:32:26 +0000815 # tee pass-through to copyable iterator
816 a, b = tee('abc')
817 c, d = tee(a)
818 self.assert_(a is c)
819
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000820 # test tee_new
821 t1, t2 = tee('abc')
822 tnew = type(t1)
823 self.assertRaises(TypeError, tnew)
824 self.assertRaises(TypeError, tnew, 10)
825 t3 = tnew(t1)
826 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000827
Raymond Hettingera9f60922004-10-17 16:40:14 +0000828 # test that tee objects are weak referencable
829 a, b = tee(xrange(10))
830 p = proxy(a)
831 self.assertEqual(getattr(p, '__class__'), type(b))
832 del a
833 self.assertRaises(ReferenceError, getattr, p, '__class__')
834
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000835 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000836 self.assertRaises(StopIteration, izip().next)
837
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000838 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000839 self.assertRaises(StopIteration, f([]).next)
840 self.assertRaises(StopIteration, f(StopNow()).next)
841
842 self.assertRaises(StopIteration, islice([], None).next)
843 self.assertRaises(StopIteration, islice(StopNow(), None).next)
844
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000845 p, q = tee([])
846 self.assertRaises(StopIteration, p.next)
847 self.assertRaises(StopIteration, q.next)
848 p, q = tee(StopNow())
849 self.assertRaises(StopIteration, p.next)
850 self.assertRaises(StopIteration, q.next)
851
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000852 self.assertRaises(StopIteration, repeat(None, 0).next)
853
854 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
855 self.assertRaises(StopIteration, f(lambda x:x, []).next)
856 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
857
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000858class TestExamples(unittest.TestCase):
859
860 def test_chain(self):
861 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
862
863 def test_chain_from_iterable(self):
864 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
865
866 def test_combinations(self):
867 self.assertEqual(list(combinations('ABCD', 2)),
868 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
869 self.assertEqual(list(combinations(range(4), 3)),
870 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
871
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000872 def test_combinations_with_replacement(self):
873 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
874 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
875
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000876 def test_compress(self):
877 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
878
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000879 def test_count(self):
880 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
881
882 def test_cycle(self):
883 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
884
885 def test_dropwhile(self):
886 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
887
888 def test_groupby(self):
889 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
890 list('ABCDAB'))
891 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
892 [list('AAAA'), list('BBB'), list('CC'), list('D')])
893
894 def test_ifilter(self):
895 self.assertEqual(list(ifilter(lambda x: x%2, range(10))), [1,3,5,7,9])
896
897 def test_ifilterfalse(self):
898 self.assertEqual(list(ifilterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
899
900 def test_imap(self):
901 self.assertEqual(list(imap(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
902
903 def test_islice(self):
904 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
905 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
906 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
907 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
908
909 def test_izip(self):
910 self.assertEqual(list(izip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
911
912 def test_izip_longest(self):
913 self.assertEqual(list(izip_longest('ABCD', 'xy', fillvalue='-')),
914 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
915
916 def test_permutations(self):
917 self.assertEqual(list(permutations('ABCD', 2)),
918 map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split()))
919 self.assertEqual(list(permutations(range(3))),
920 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
921
922 def test_product(self):
923 self.assertEqual(list(product('ABCD', 'xy')),
924 map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split()))
925 self.assertEqual(list(product(range(2), repeat=3)),
926 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
927 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
928
929 def test_repeat(self):
930 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
931
932 def test_stapmap(self):
933 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
934 [32, 9, 1000])
935
936 def test_takewhile(self):
937 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
938
939
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000940class TestGC(unittest.TestCase):
941
942 def makecycle(self, iterator, container):
943 container.append(iterator)
944 iterator.next()
945 del container, iterator
946
947 def test_chain(self):
948 a = []
949 self.makecycle(chain(a), a)
950
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000951 def test_chain_from_iterable(self):
952 a = []
953 self.makecycle(chain.from_iterable([a]), a)
954
955 def test_combinations(self):
956 a = []
957 self.makecycle(combinations([1,2,a,3], 3), a)
958
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000959 def test_combinations_with_replacement(self):
960 a = []
961 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
962
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000963 def test_compress(self):
964 a = []
965 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
966
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000967 def test_cycle(self):
968 a = []
969 self.makecycle(cycle([a]*2), a)
970
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000971 def test_dropwhile(self):
972 a = []
973 self.makecycle(dropwhile(bool, [0, a, a]), a)
974
975 def test_groupby(self):
976 a = []
977 self.makecycle(groupby([a]*2, lambda x:x), a)
978
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000979 def test_issue2246(self):
980 # Issue 2246 -- the _grouper iterator was not included in GC
981 n = 10
982 keyfunc = lambda x: x
983 for i, j in groupby(xrange(n), key=keyfunc):
984 keyfunc.__dict__.setdefault('x',[]).append(j)
985
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000986 def test_ifilter(self):
987 a = []
988 self.makecycle(ifilter(lambda x:True, [a]*2), a)
989
990 def test_ifilterfalse(self):
991 a = []
992 self.makecycle(ifilterfalse(lambda x:False, a), a)
993
994 def test_izip(self):
995 a = []
996 self.makecycle(izip([a]*2, [a]*3), a)
997
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000998 def test_izip_longest(self):
999 a = []
1000 self.makecycle(izip_longest([a]*2, [a]*3), a)
1001 b = [a, None]
1002 self.makecycle(izip_longest([a]*2, [a]*3, fillvalue=b), a)
1003
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001004 def test_imap(self):
1005 a = []
1006 self.makecycle(imap(lambda x:x, [a]*2), a)
1007
1008 def test_islice(self):
1009 a = []
1010 self.makecycle(islice([a]*2, None), a)
1011
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001012 def test_permutations(self):
1013 a = []
1014 self.makecycle(permutations([1,2,a,3], 3), a)
1015
1016 def test_product(self):
1017 a = []
1018 self.makecycle(product([1,2,a,3], repeat=3), a)
1019
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001020 def test_repeat(self):
1021 a = []
1022 self.makecycle(repeat(a), a)
1023
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001024 def test_starmap(self):
1025 a = []
1026 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1027
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001028 def test_takewhile(self):
1029 a = []
1030 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1031
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001032def R(seqn):
1033 'Regular generator'
1034 for i in seqn:
1035 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001036
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001037class G:
1038 'Sequence using __getitem__'
1039 def __init__(self, seqn):
1040 self.seqn = seqn
1041 def __getitem__(self, i):
1042 return self.seqn[i]
1043
1044class I:
1045 'Sequence using iterator protocol'
1046 def __init__(self, seqn):
1047 self.seqn = seqn
1048 self.i = 0
1049 def __iter__(self):
1050 return self
1051 def next(self):
1052 if self.i >= len(self.seqn): raise StopIteration
1053 v = self.seqn[self.i]
1054 self.i += 1
1055 return v
1056
1057class Ig:
1058 'Sequence using iterator protocol defined with a generator'
1059 def __init__(self, seqn):
1060 self.seqn = seqn
1061 self.i = 0
1062 def __iter__(self):
1063 for val in self.seqn:
1064 yield val
1065
1066class X:
1067 'Missing __getitem__ and __iter__'
1068 def __init__(self, seqn):
1069 self.seqn = seqn
1070 self.i = 0
1071 def next(self):
1072 if self.i >= len(self.seqn): raise StopIteration
1073 v = self.seqn[self.i]
1074 self.i += 1
1075 return v
1076
1077class N:
1078 'Iterator missing next()'
1079 def __init__(self, seqn):
1080 self.seqn = seqn
1081 self.i = 0
1082 def __iter__(self):
1083 return self
1084
1085class E:
1086 'Test propagation of exceptions'
1087 def __init__(self, seqn):
1088 self.seqn = seqn
1089 self.i = 0
1090 def __iter__(self):
1091 return self
1092 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001093 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001094
1095class S:
1096 'Test immediate stop'
1097 def __init__(self, seqn):
1098 pass
1099 def __iter__(self):
1100 return self
1101 def next(self):
1102 raise StopIteration
1103
1104def L(seqn):
1105 'Test multiple tiers of iterators'
1106 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
1107
1108
1109class TestVariousIteratorArgs(unittest.TestCase):
1110
1111 def test_chain(self):
1112 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1113 for g in (G, I, Ig, S, L, R):
1114 self.assertEqual(list(chain(g(s))), list(g(s)))
1115 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001116 self.assertRaises(TypeError, list, chain(X(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001117 self.assertRaises(TypeError, list, chain(N(s)))
1118 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1119
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001120 def test_compress(self):
1121 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1122 n = len(s)
1123 for g in (G, I, Ig, S, L, R):
1124 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1125 self.assertRaises(TypeError, compress, X(s), repeat(1))
1126 self.assertRaises(TypeError, list, compress(N(s), repeat(1)))
1127 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1128
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001129 def test_product(self):
1130 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1131 self.assertRaises(TypeError, product, X(s))
1132 self.assertRaises(TypeError, product, N(s))
1133 self.assertRaises(ZeroDivisionError, product, E(s))
1134
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001135 def test_cycle(self):
1136 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1137 for g in (G, I, Ig, S, L, R):
1138 tgtlen = len(s) * 3
1139 expected = list(g(s))*3
1140 actual = list(islice(cycle(g(s)), tgtlen))
1141 self.assertEqual(actual, expected)
1142 self.assertRaises(TypeError, cycle, X(s))
1143 self.assertRaises(TypeError, list, cycle(N(s)))
1144 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1145
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001146 def test_groupby(self):
1147 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1148 for g in (G, I, Ig, S, L, R):
1149 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1150 self.assertRaises(TypeError, groupby, X(s))
1151 self.assertRaises(TypeError, list, groupby(N(s)))
1152 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1153
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001154 def test_ifilter(self):
1155 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1156 for g in (G, I, Ig, S, L, R):
1157 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
1158 self.assertRaises(TypeError, ifilter, isEven, X(s))
1159 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
1160 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
1161
1162 def test_ifilterfalse(self):
1163 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1164 for g in (G, I, Ig, S, L, R):
1165 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
1166 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
1167 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
1168 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
1169
1170 def test_izip(self):
1171 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1172 for g in (G, I, Ig, S, L, R):
1173 self.assertEqual(list(izip(g(s))), zip(g(s)))
1174 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
1175 self.assertRaises(TypeError, izip, X(s))
1176 self.assertRaises(TypeError, list, izip(N(s)))
1177 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
1178
Raymond Hettingerd36862c2007-02-21 05:20:38 +00001179 def test_iziplongest(self):
1180 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1181 for g in (G, I, Ig, S, L, R):
1182 self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
1183 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
1184 self.assertRaises(TypeError, izip_longest, X(s))
1185 self.assertRaises(TypeError, list, izip_longest(N(s)))
1186 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
1187
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001188 def test_imap(self):
1189 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1190 for g in (G, I, Ig, S, L, R):
1191 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
1192 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
1193 self.assertRaises(TypeError, imap, onearg, X(s))
1194 self.assertRaises(TypeError, list, imap(onearg, N(s)))
1195 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
1196
1197 def test_islice(self):
1198 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1199 for g in (G, I, Ig, S, L, R):
1200 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1201 self.assertRaises(TypeError, islice, X(s), 10)
1202 self.assertRaises(TypeError, list, islice(N(s), 10))
1203 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1204
1205 def test_starmap(self):
1206 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1207 for g in (G, I, Ig, S, L, R):
1208 ss = zip(s, s)
1209 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
1210 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
1211 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
1212 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1213
1214 def test_takewhile(self):
1215 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1216 for g in (G, I, Ig, S, L, R):
1217 tgt = []
1218 for elem in g(s):
1219 if not isEven(elem): break
1220 tgt.append(elem)
1221 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1222 self.assertRaises(TypeError, takewhile, isEven, X(s))
1223 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
1224 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1225
1226 def test_dropwhile(self):
1227 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1228 for g in (G, I, Ig, S, L, R):
1229 tgt = []
1230 for elem in g(s):
1231 if not tgt and isOdd(elem): continue
1232 tgt.append(elem)
1233 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1234 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
1235 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
1236 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1237
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001238 def test_tee(self):
1239 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1240 for g in (G, I, Ig, S, L, R):
1241 it1, it2 = tee(g(s))
1242 self.assertEqual(list(it1), list(g(s)))
1243 self.assertEqual(list(it2), list(g(s)))
1244 self.assertRaises(TypeError, tee, X(s))
1245 self.assertRaises(TypeError, list, tee(N(s))[0])
1246 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1247
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001248class LengthTransparency(unittest.TestCase):
1249
1250 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001251 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001252 self.assertEqual(len(repeat(None, 50)), 50)
1253 self.assertRaises(TypeError, len, repeat(None))
1254
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001255class RegressionTests(unittest.TestCase):
1256
1257 def test_sf_793826(self):
1258 # Fix Armin Rigo's successful efforts to wreak havoc
1259
1260 def mutatingtuple(tuple1, f, tuple2):
1261 # this builds a tuple t which is a copy of tuple1,
1262 # then calls f(t), then mutates t to be equal to tuple2
1263 # (needs len(tuple1) == len(tuple2)).
1264 def g(value, first=[1]):
1265 if first:
1266 del first[:]
1267 f(z.next())
1268 return value
1269 items = list(tuple2)
1270 items[1:1] = list(tuple1)
1271 gen = imap(g, items)
1272 z = izip(*[gen]*len(tuple1))
1273 z.next()
1274
1275 def f(t):
1276 global T
1277 T = t
1278 first[:] = list(T)
1279
1280 first = []
1281 mutatingtuple((1,2,3), f, (4,5,6))
1282 second = list(T)
1283 self.assertEqual(first, second)
1284
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001285
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001286 def test_sf_950057(self):
1287 # Make sure that chain() and cycle() catch exceptions immediately
1288 # rather than when shifting between input sources
1289
1290 def gen1():
1291 hist.append(0)
1292 yield 1
1293 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001294 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001295 hist.append(2)
1296
1297 def gen2(x):
1298 hist.append(3)
1299 yield 2
1300 hist.append(4)
1301 if x:
1302 raise StopIteration
1303
1304 hist = []
1305 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1306 self.assertEqual(hist, [0,1])
1307
1308 hist = []
1309 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1310 self.assertEqual(hist, [0,1])
1311
1312 hist = []
1313 self.assertRaises(AssertionError, list, cycle(gen1()))
1314 self.assertEqual(hist, [0,1])
1315
Georg Brandlb84c1372007-01-21 10:28:43 +00001316class SubclassWithKwargsTest(unittest.TestCase):
1317 def test_keywords_in_subclass(self):
1318 # count is not subclassable...
1319 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001320 starmap, islice, takewhile, dropwhile, cycle, compress):
Georg Brandlb84c1372007-01-21 10:28:43 +00001321 class Subclass(cls):
1322 def __init__(self, newarg=None, *args):
1323 cls.__init__(self, *args)
1324 try:
1325 Subclass(newarg=1)
1326 except TypeError, err:
1327 # we expect type errors because of wrong argument count
1328 self.failIf("does not take keyword arguments" in err.args[0])
1329
1330
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001331libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001332
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001333
1334>>> amounts = [120.15, 764.05, 823.14]
1335>>> for checknum, amount in izip(count(1200), amounts):
1336... print 'Check %d is for $%.2f' % (checknum, amount)
1337...
1338Check 1200 is for $120.15
1339Check 1201 is for $764.05
1340Check 1202 is for $823.14
1341
1342>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001343>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1344... print cube
1345...
13461
13478
134827
1349
1350>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001351>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001352... print name.title()
1353...
1354Alex
1355Laura
1356Martin
1357Walter
1358Samuele
1359
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001360>>> from operator import itemgetter
1361>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Armin Rigoa3f09272006-05-28 19:13:17 +00001362>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001363>>> for k, g in groupby(di, itemgetter(1)):
1364... print k, map(itemgetter(0), g)
1365...
13661 ['a', 'c', 'e']
13672 ['b', 'd', 'f']
13683 ['g']
1369
Raymond Hettinger734fb572004-01-20 20:04:40 +00001370# Find runs of consecutive numbers using groupby. The key to the solution
1371# is differencing with a range so that consecutive numbers all appear in
1372# same group.
1373>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
1374>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
1375... print map(operator.itemgetter(1), g)
1376...
1377[1]
1378[4, 5, 6]
1379[10]
1380[15, 16, 17, 18]
1381[22]
1382[25, 26, 27, 28]
1383
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001384>>> def take(n, iterable):
1385... "Return first n items of the iterable as a list"
1386... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001387
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001388>>> def enumerate(iterable, start=0):
1389... return izip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001390
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001391>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001392... "Return function(0), function(1), ..."
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001393... return imap(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001394
1395>>> def nth(iterable, n):
Raymond Hettingerd507afd2009-02-04 10:52:32 +00001396... "Returns the nth item or None"
1397... return next(islice(iterable, n, None), None)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001398
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001399>>> def quantify(iterable, pred=bool):
1400... "Count how many times the predicate is true"
1401... return sum(imap(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001402
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001403>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001404... "Returns the sequence elements and then returns None indefinitely"
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001405... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001406
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001407>>> def ncycles(iterable, n):
1408... "Returns the seqeuence elements n times"
1409... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001410
1411>>> def dotproduct(vec1, vec2):
1412... return sum(imap(operator.mul, vec1, vec2))
1413
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001414>>> def flatten(listOfLists):
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001415... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001416
1417>>> def repeatfunc(func, times=None, *args):
1418... "Repeat calls to func with specified arguments."
1419... " Example: repeatfunc(random.random)"
1420... if times is None:
1421... return starmap(func, repeat(args))
1422... else:
1423... return starmap(func, repeat(args, times))
1424
Raymond Hettingerd591f662003-10-26 15:34:50 +00001425>>> def pairwise(iterable):
1426... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1427... a, b = tee(iterable)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001428... for elem in b:
1429... break
Raymond Hettingerad983e72003-11-12 14:32:26 +00001430... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001431
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001432>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001433... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001434... args = [iter(iterable)] * n
Raymond Hettingerf080e6d2008-07-31 01:19:50 +00001435... return izip_longest(fillvalue=fillvalue, *args)
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001436
1437>>> def roundrobin(*iterables):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001438... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001439... # Recipe credited to George Sakkis
1440... pending = len(iterables)
1441... nexts = cycle(iter(it).next for it in iterables)
1442... while pending:
1443... try:
1444... for next in nexts:
1445... yield next()
1446... except StopIteration:
1447... pending -= 1
1448... nexts = cycle(islice(nexts, pending))
1449
1450>>> def powerset(iterable):
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001451... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1452... s = list(iterable)
1453... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001454
Raymond Hettinger44e15812009-01-02 21:26:45 +00001455>>> def unique_everseen(iterable, key=None):
1456... "List unique elements, preserving order. Remember all elements ever seen."
1457... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1458... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1459... seen = set()
1460... seen_add = seen.add
1461... if key is None:
1462... for element in iterable:
1463... if element not in seen:
1464... seen_add(element)
1465... yield element
1466... else:
1467... for element in iterable:
1468... k = key(element)
1469... if k not in seen:
1470... seen_add(k)
1471... yield element
1472
1473>>> def unique_justseen(iterable, key=None):
1474... "List unique elements, preserving order. Remember only the element just seen."
1475... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1476... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1477... return imap(next, imap(itemgetter(1), groupby(iterable, key)))
1478
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001479This is not part of the examples but it tests to make sure the definitions
1480perform as purported.
1481
Raymond Hettingera098b332003-09-08 23:58:40 +00001482>>> take(10, count())
1483[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1484
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001485>>> list(enumerate('abc'))
1486[(0, 'a'), (1, 'b'), (2, 'c')]
1487
1488>>> list(islice(tabulate(lambda x: 2*x), 4))
1489[0, 2, 4, 6]
1490
1491>>> nth('abcde', 3)
Raymond Hettingerd507afd2009-02-04 10:52:32 +00001492'd'
1493
1494>>> nth('abcde', 9) is None
1495True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001496
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001497>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000149850
1499
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001500>>> a = [[1, 2, 3], [4, 5, 6]]
1501>>> flatten(a)
1502[1, 2, 3, 4, 5, 6]
1503
1504>>> list(repeatfunc(pow, 5, 2, 3))
1505[8, 8, 8, 8, 8]
1506
1507>>> import random
1508>>> take(5, imap(int, repeatfunc(random.random)))
1509[0, 0, 0, 0, 0]
1510
Raymond Hettingerd591f662003-10-26 15:34:50 +00001511>>> list(pairwise('abcd'))
1512[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001513
Raymond Hettingerd591f662003-10-26 15:34:50 +00001514>>> list(pairwise([]))
1515[]
1516
1517>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001518[]
1519
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001520>>> list(islice(padnone('abc'), 0, 6))
1521['a', 'b', 'c', None, None, None]
1522
1523>>> list(ncycles('abc', 3))
1524['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1525
1526>>> dotproduct([1,2,3], [4,5,6])
152732
1528
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001529>>> list(grouper(3, 'abcdefg', 'x'))
1530[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1531
1532>>> list(roundrobin('abc', 'd', 'ef'))
1533['a', 'd', 'e', 'b', 'f', 'c']
1534
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001535>>> list(powerset([1,2,3]))
1536[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001537
Raymond Hettinger560f9a82009-01-27 13:26:35 +00001538>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1539True
1540
1541>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1542True
1543
Raymond Hettinger44e15812009-01-02 21:26:45 +00001544>>> list(unique_everseen('AAAABBBCCDAABBB'))
1545['A', 'B', 'C', 'D']
1546
1547>>> list(unique_everseen('ABBCcAD', str.lower))
1548['A', 'B', 'C', 'D']
1549
1550>>> list(unique_justseen('AAAABBBCCDAABBB'))
1551['A', 'B', 'C', 'D', 'A', 'B']
1552
1553>>> list(unique_justseen('ABBCcAD', str.lower))
1554['A', 'B', 'C', 'A', 'D']
1555
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001556"""
1557
1558__test__ = {'libreftest' : libreftest}
1559
1560def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001561 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Georg Brandlb84c1372007-01-21 10:28:43 +00001562 RegressionTests, LengthTransparency,
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001563 SubclassWithKwargsTest, TestExamples)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001564 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001565
1566 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001567 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001568 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001569 counts = [None] * 5
1570 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001571 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001572 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001573 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001574 print counts
1575
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001576 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001577 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001578
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001579if __name__ == "__main__":
1580 test_main(verbose=True)