blob: 2cdb1d7dc0ba3ecc8c28b4220c4450e28b11fa83 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001import unittest
2from test import test_support
3from itertools import *
Raymond Hettingera9f60922004-10-17 16:40:14 +00004from weakref import proxy
Raymond Hettingeraa044612009-02-12 12:04:26 +00005from decimal import Decimal
Raymond Hettingerdbe3bfb2009-02-12 12:43:01 +00006from fractions import Fraction
Raymond Hettinger2012f172003-02-07 05:32:58 +00007import sys
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00008import operator
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00009import random
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +000010maxsize = test_support.MAX_Py_ssize_t
11minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000012
13def onearg(x):
14 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000015 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000016
17def errfunc(*args):
18 'Test function that raises an error'
19 raise ValueError
20
21def gen3():
22 'Non-restartable source sequence'
23 for i in (0, 1, 2):
24 yield i
25
26def isEven(x):
27 'Test predicate'
28 return x%2==0
29
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000030def isOdd(x):
31 'Test predicate'
32 return x%2==1
33
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000034class StopNow:
35 'Class emulating an empty iterable.'
36 def __iter__(self):
37 return self
38 def next(self):
39 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000040
Raymond Hettinger02420702003-06-29 20:36:23 +000041def take(n, seq):
42 'Convenience function for partially consuming a long of infinite iterable'
43 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000044
Raymond Hettingerd553d852008-03-04 04:17:08 +000045def prod(iterable):
46 return reduce(operator.mul, iterable, 1)
47
Raymond Hettinger93e804d2008-02-26 23:40:50 +000048def fact(n):
49 'Factorial'
Raymond Hettingerd553d852008-03-04 04:17:08 +000050 return prod(range(1, n+1))
51
Raymond Hettinger96ef8112003-02-01 00:10:11 +000052class TestBasicOps(unittest.TestCase):
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000053 def test_chain(self):
Raymond Hettingerad47fa12008-03-06 20:52:01 +000054
55 def chain2(*iterables):
56 'Pure python version in the docs'
57 for it in iterables:
58 for element in it:
59 yield element
60
61 for c in (chain, chain2):
62 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
63 self.assertEqual(list(c('abc')), list('abc'))
64 self.assertEqual(list(c('')), [])
65 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
66 self.assertRaises(TypeError, list,c(2, 3))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000067
Raymond Hettingerb4cbc982008-02-28 22:46:41 +000068 def test_chain_from_iterable(self):
69 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
70 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
71 self.assertEqual(list(chain.from_iterable([''])), [])
72 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
73 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
74
Raymond Hettinger93e804d2008-02-26 23:40:50 +000075 def test_combinations(self):
Raymond Hettinger5b913e32009-01-08 06:39:04 +000076 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Raymond Hettinger93e804d2008-02-26 23:40:50 +000077 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Raymond Hettingerd553d852008-03-04 04:17:08 +000078 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Raymond Hettinger93e804d2008-02-26 23:40:50 +000079 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Raymond Hettinger5b913e32009-01-08 06:39:04 +000080 self.assertEqual(list(combinations('abc', 32)), []) # r > n
Raymond Hettinger93e804d2008-02-26 23:40:50 +000081 self.assertEqual(list(combinations(range(4), 3)),
82 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
Raymond Hettingerd553d852008-03-04 04:17:08 +000083
84 def combinations1(iterable, r):
85 'Pure python version shown in the docs'
86 pool = tuple(iterable)
87 n = len(pool)
Raymond Hettinger5b913e32009-01-08 06:39:04 +000088 if r > n:
89 return
Raymond Hettingerd553d852008-03-04 04:17:08 +000090 indices = range(r)
91 yield tuple(pool[i] for i in indices)
92 while 1:
93 for i in reversed(range(r)):
94 if indices[i] != i + n - r:
95 break
96 else:
97 return
98 indices[i] += 1
99 for j in range(i+1, r):
100 indices[j] = indices[j-1] + 1
101 yield tuple(pool[i] for i in indices)
102
103 def combinations2(iterable, r):
104 'Pure python version shown in the docs'
105 pool = tuple(iterable)
106 n = len(pool)
107 for indices in permutations(range(n), r):
108 if sorted(indices) == list(indices):
109 yield tuple(pool[i] for i in indices)
110
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000111 def combinations3(iterable, r):
112 'Pure python version from cwr()'
113 pool = tuple(iterable)
114 n = len(pool)
115 for indices in combinations_with_replacement(range(n), r):
116 if len(set(indices)) == r:
117 yield tuple(pool[i] for i in indices)
118
Raymond Hettingerd553d852008-03-04 04:17:08 +0000119 for n in range(7):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000120 values = [5*x-12 for x in range(n)]
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000121 for r in range(n+2):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000122 result = list(combinations(values, r))
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000123 self.assertEqual(len(result), 0 if r>n else fact(n) / fact(r) / fact(n-r)) # right number of combs
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000124 self.assertEqual(len(result), len(set(result))) # no repeats
125 self.assertEqual(result, sorted(result)) # lexicographic order
126 for c in result:
127 self.assertEqual(len(c), r) # r-length combinations
128 self.assertEqual(len(set(c)), r) # no duplicate elements
129 self.assertEqual(list(c), sorted(c)) # keep original ordering
130 self.assert_(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000131 self.assertEqual(list(c),
132 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Raymond Hettingerd553d852008-03-04 04:17:08 +0000133 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000134 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000135 self.assertEqual(result, list(combinations3(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000136
137 # Test implementation detail: tuple re-use
138 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
139 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
140
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000141 def test_combinations_with_replacement(self):
142 cwr = combinations_with_replacement
143 self.assertRaises(TypeError, cwr, 'abc') # missing r argument
144 self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
145 self.assertRaises(TypeError, cwr, None) # pool is not iterable
146 self.assertRaises(ValueError, cwr, 'abc', -2) # r is negative
147 self.assertEqual(list(cwr('ABC', 2)),
148 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
149
150 def cwr1(iterable, r):
151 'Pure python version shown in the docs'
152 # number items returned: (n+r-1)! / r! / (n-1)! when n>0
153 pool = tuple(iterable)
154 n = len(pool)
155 if not n and r:
156 return
157 indices = [0] * r
158 yield tuple(pool[i] for i in indices)
159 while 1:
160 for i in reversed(range(r)):
161 if indices[i] != n - 1:
162 break
163 else:
164 return
165 indices[i:] = [indices[i] + 1] * (r - i)
166 yield tuple(pool[i] for i in indices)
167
168 def cwr2(iterable, r):
169 'Pure python version shown in the docs'
170 pool = tuple(iterable)
171 n = len(pool)
172 for indices in product(range(n), repeat=r):
173 if sorted(indices) == list(indices):
174 yield tuple(pool[i] for i in indices)
175
176 def numcombs(n, r):
177 if not n:
178 return 0 if r else 1
179 return fact(n+r-1) / fact(r)/ fact(n-1)
180
181 for n in range(7):
182 values = [5*x-12 for x in range(n)]
183 for r in range(n+2):
184 result = list(cwr(values, r))
185
186 self.assertEqual(len(result), numcombs(n, r)) # right number of combs
187 self.assertEqual(len(result), len(set(result))) # no repeats
188 self.assertEqual(result, sorted(result)) # lexicographic order
189
190 regular_combs = list(combinations(values, r)) # compare to combs without replacement
191 if n == 0 or r <= 1:
192 self.assertEquals(result, regular_combs) # cases that should be identical
193 else:
194 self.assert_(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
195
196 for c in result:
197 self.assertEqual(len(c), r) # r-length combinations
198 noruns = [k for k,v in groupby(c)] # combo without consecutive repeats
199 self.assertEqual(len(noruns), len(set(noruns))) # no repeats other than consecutive
200 self.assertEqual(list(c), sorted(c)) # keep original ordering
201 self.assert_(all(e in values for e in c)) # elements taken from input iterable
202 self.assertEqual(noruns,
203 [e for e in values if e in c]) # comb is a subsequence of the input iterable
204 self.assertEqual(result, list(cwr1(values, r))) # matches first pure python version
205 self.assertEqual(result, list(cwr2(values, r))) # matches second pure python version
206
207 # Test implementation detail: tuple re-use
208 self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
209 self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
210
Raymond Hettingerd553d852008-03-04 04:17:08 +0000211 def test_permutations(self):
212 self.assertRaises(TypeError, permutations) # too few arguments
213 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000214 self.assertRaises(TypeError, permutations, None) # pool is not iterable
215 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000216 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000217 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Raymond Hettingerd553d852008-03-04 04:17:08 +0000218 self.assertEqual(list(permutations(range(3), 2)),
219 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
220
221 def permutations1(iterable, r=None):
222 'Pure python version shown in the docs'
223 pool = tuple(iterable)
224 n = len(pool)
225 r = n if r is None else r
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000226 if r > n:
227 return
Raymond Hettingerd553d852008-03-04 04:17:08 +0000228 indices = range(n)
Raymond Hettingere70bb8d2008-03-23 00:55:46 +0000229 cycles = range(n, n-r, -1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000230 yield tuple(pool[i] for i in indices[:r])
231 while n:
232 for i in reversed(range(r)):
233 cycles[i] -= 1
234 if cycles[i] == 0:
235 indices[i:] = indices[i+1:] + indices[i:i+1]
236 cycles[i] = n - i
237 else:
238 j = cycles[i]
239 indices[i], indices[-j] = indices[-j], indices[i]
240 yield tuple(pool[i] for i in indices[:r])
241 break
242 else:
243 return
244
245 def permutations2(iterable, r=None):
246 'Pure python version shown in the docs'
247 pool = tuple(iterable)
248 n = len(pool)
249 r = n if r is None else r
250 for indices in product(range(n), repeat=r):
251 if len(set(indices)) == r:
252 yield tuple(pool[i] for i in indices)
253
254 for n in range(7):
255 values = [5*x-12 for x in range(n)]
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000256 for r in range(n+2):
Raymond Hettingerd553d852008-03-04 04:17:08 +0000257 result = list(permutations(values, r))
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000258 self.assertEqual(len(result), 0 if r>n else fact(n) / fact(n-r)) # right number of perms
Raymond Hettingerd553d852008-03-04 04:17:08 +0000259 self.assertEqual(len(result), len(set(result))) # no repeats
260 self.assertEqual(result, sorted(result)) # lexicographic order
261 for p in result:
262 self.assertEqual(len(p), r) # r-length permutations
263 self.assertEqual(len(set(p)), r) # no duplicate elements
264 self.assert_(all(e in values for e in p)) # elements taken from input iterable
265 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000266 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000267 if r == n:
268 self.assertEqual(result, list(permutations(values, None))) # test r as None
269 self.assertEqual(result, list(permutations(values))) # test default r
270
271 # Test implementation detail: tuple re-use
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000272 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000273 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000274
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000275 def test_combinatorics(self):
276 # Test relationships between product(), permutations(),
277 # combinations() and combinations_with_replacement().
278
Raymond Hettinger2f6c2e02009-01-27 10:36:14 +0000279 for n in range(6):
280 s = 'ABCDEFG'[:n]
281 for r in range(8):
282 prod = list(product(s, repeat=r))
283 cwr = list(combinations_with_replacement(s, r))
284 perm = list(permutations(s, r))
285 comb = list(combinations(s, r))
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000286
Raymond Hettinger2f6c2e02009-01-27 10:36:14 +0000287 # Check size
288 self.assertEquals(len(prod), n**r)
289 self.assertEquals(len(cwr), (fact(n+r-1) / fact(r)/ fact(n-1)) if n else (not r))
290 self.assertEquals(len(perm), 0 if r>n else fact(n) / fact(n-r))
291 self.assertEquals(len(comb), 0 if r>n else fact(n) / fact(r) / fact(n-r))
292
293 # Check lexicographic order without repeated tuples
294 self.assertEquals(prod, sorted(set(prod)))
295 self.assertEquals(cwr, sorted(set(cwr)))
296 self.assertEquals(perm, sorted(set(perm)))
297 self.assertEquals(comb, sorted(set(comb)))
298
299 # Check interrelationships
300 self.assertEquals(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
301 self.assertEquals(perm, [t for t in prod if len(set(t))==r]) # perm: prods with no dups
302 self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
303 self.assertEqual(comb, [t for t in cwr if len(set(t))==r]) # comb: cwrs without dups
304 self.assertEqual(comb, filter(set(cwr).__contains__, perm)) # comb: perm that is a cwr
305 self.assertEqual(comb, filter(set(perm).__contains__, cwr)) # comb: cwr that is a perm
306 self.assertEqual(comb, sorted(set(cwr) & set(perm))) # comb: both a cwr and a perm
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000307
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000308 def test_compress(self):
Raymond Hettinger2e2909f2009-02-19 02:15:14 +0000309 self.assertEqual(list(compress(data='ABCDEF', selectors=[1,0,1,0,1,1])), list('ACEF'))
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000310 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
311 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
312 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
313 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
314 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
315 n = 10000
316 data = chain.from_iterable(repeat(range(6), n))
317 selectors = chain.from_iterable(repeat((0, 1)))
318 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
319 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
320 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
321 self.assertRaises(TypeError, compress, range(6)) # too few args
322 self.assertRaises(TypeError, compress, range(6), None) # too many args
323
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000324 def test_count(self):
325 self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
326 self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000327 self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000328 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
329 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000330 self.assertRaises(TypeError, count, 2, 3, 4)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000331 self.assertRaises(TypeError, count, 'a')
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000332 self.assertEqual(list(islice(count(maxsize-5), 10)), range(maxsize-5, maxsize+5))
333 self.assertEqual(list(islice(count(-maxsize-5), 10)), range(-maxsize-5, -maxsize+5))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000334 c = count(3)
335 self.assertEqual(repr(c), 'count(3)')
336 c.next()
337 self.assertEqual(repr(c), 'count(4)')
Jack Diederich36234e82006-09-21 17:50:26 +0000338 c = count(-9)
339 self.assertEqual(repr(c), 'count(-9)')
340 c.next()
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000341 self.assertEqual(repr(count(10.25)), 'count(10.25)')
Jack Diederich36234e82006-09-21 17:50:26 +0000342 self.assertEqual(c.next(), -8)
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000343 for i in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 10, sys.maxint-5, sys.maxint+5):
Raymond Hettinger3a0de082007-10-12 17:53:11 +0000344 # Test repr (ignoring the L in longs)
345 r1 = repr(count(i)).replace('L', '')
346 r2 = 'count(%r)'.__mod__(i).replace('L', '')
347 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000348
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000349 def test_count_with_stride(self):
350 self.assertEqual(zip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingera4038032009-02-14 00:25:51 +0000351 self.assertEqual(zip('abc',count(start=2,step=3)),
352 [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000353 self.assertEqual(zip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
354 self.assertEqual(zip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
355 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
356 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
357 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettingeraa044612009-02-12 12:04:26 +0000358 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
359 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerdbe3bfb2009-02-12 12:43:01 +0000360 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
361 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000362 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
363 c = count(3, 5)
364 self.assertEqual(repr(c), 'count(3, 5)')
365 c.next()
366 self.assertEqual(repr(c), 'count(8, 5)')
367 c = count(-9, 0)
368 self.assertEqual(repr(c), 'count(-9, 0)')
369 c.next()
370 self.assertEqual(repr(c), 'count(-9, 0)')
371 c = count(-9, -3)
372 self.assertEqual(repr(c), 'count(-9, -3)')
373 c.next()
374 self.assertEqual(repr(c), 'count(-12, -3)')
375 self.assertEqual(repr(c), 'count(-12, -3)')
376 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
377 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
378 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
379 for i in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 10, sys.maxint-5, sys.maxint+5):
380 for j in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 1, 10, sys.maxint-5, sys.maxint+5):
381 # Test repr (ignoring the L in longs)
382 r1 = repr(count(i, j)).replace('L', '')
383 if j == 1:
384 r2 = ('count(%r)' % i).replace('L', '')
385 else:
386 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
387 self.assertEqual(r1, r2)
388
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000389 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000390 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000391 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000392 self.assertRaises(TypeError, cycle)
393 self.assertRaises(TypeError, cycle, 5)
394 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000395
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000396 def test_groupby(self):
397 # Check whether it accepts arguments correctly
398 self.assertEqual([], list(groupby([])))
399 self.assertEqual([], list(groupby([], key=id)))
400 self.assertRaises(TypeError, list, groupby('abc', []))
401 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000402 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000403
404 # Check normal input
405 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
406 (2,15,22), (3,16,23), (3,17,23)]
407 dup = []
408 for k, g in groupby(s, lambda r:r[0]):
409 for elem in g:
410 self.assertEqual(k, elem[0])
411 dup.append(elem)
412 self.assertEqual(s, dup)
413
414 # Check nested case
415 dup = []
416 for k, g in groupby(s, lambda r:r[0]):
417 for ik, ig in groupby(g, lambda r:r[2]):
418 for elem in ig:
419 self.assertEqual(k, elem[0])
420 self.assertEqual(ik, elem[2])
421 dup.append(elem)
422 self.assertEqual(s, dup)
423
424 # Check case where inner iterator is not used
425 keys = [k for k, g in groupby(s, lambda r:r[0])]
426 expectedkeys = set([r[0] for r in s])
427 self.assertEqual(set(keys), expectedkeys)
428 self.assertEqual(len(keys), len(expectedkeys))
429
430 # Exercise pipes and filters style
431 s = 'abracadabra'
432 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000433 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000434 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
435 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000436 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000437 self.assertEqual(r, ['a', 'b', 'r'])
438 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000439 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000440 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
441 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000442 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000443 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
444
445 # iter.next failure
446 class ExpectedError(Exception):
447 pass
448 def delayed_raise(n=0):
449 for i in range(n):
450 yield 'yo'
451 raise ExpectedError
452 def gulp(iterable, keyp=None, func=list):
453 return [func(g) for k, g in groupby(iterable, keyp)]
454
455 # iter.next failure on outer object
456 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
457 # iter.next failure on inner object
458 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
459
460 # __cmp__ failure
461 class DummyCmp:
462 def __cmp__(self, dst):
463 raise ExpectedError
464 s = [DummyCmp(), DummyCmp(), None]
465
466 # __cmp__ failure on outer object
467 self.assertRaises(ExpectedError, gulp, s, func=id)
468 # __cmp__ failure on inner object
469 self.assertRaises(ExpectedError, gulp, s)
470
471 # keyfunc failure
472 def keyfunc(obj):
473 if keyfunc.skip > 0:
474 keyfunc.skip -= 1
475 return obj
476 else:
477 raise ExpectedError
478
479 # keyfunc failure on outer object
480 keyfunc.skip = 0
481 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
482 keyfunc.skip = 1
483 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
484
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000485 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000486 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000487 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000488 self.assertEqual(list(ifilter(bool, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000489 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000490 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000491 self.assertRaises(TypeError, ifilter, lambda x:x)
492 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000493 self.assertRaises(TypeError, ifilter, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000494 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000495
496 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000497 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
498 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000499 self.assertEqual(list(ifilterfalse(bool, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000500 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000501 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000502 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
503 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000504 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000505 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000506
507 def test_izip(self):
508 ans = [(x,y) for x, y in izip('abc',count())]
509 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000510 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
511 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
Raymond Hettinger02420702003-06-29 20:36:23 +0000512 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000513 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000514 self.assertEqual(list(izip()), zip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000515 self.assertRaises(TypeError, izip, 3)
516 self.assertRaises(TypeError, izip, range(3), 3)
517 # Check tuple re-use (implementation detail)
518 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
519 zip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000520 self.assertEqual([pair for pair in izip('abc', 'def')],
521 zip('abc', 'def'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000522 ids = map(id, izip('abc', 'def'))
523 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000524 ids = map(id, list(izip('abc', 'def')))
525 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000526
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000527 def test_iziplongest(self):
528 for args in [
529 ['abc', range(6)],
530 [range(6), 'abc'],
531 [range(1000), range(2000,2100), range(3000,3050)],
532 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
533 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
534 ]:
535 target = map(None, *args)
536 self.assertEqual(list(izip_longest(*args)), target)
537 self.assertEqual(list(izip_longest(*args, **{})), target)
538 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
539 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Tim Petersea5962f2007-03-12 18:07:52 +0000540
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000541 self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
542
543 self.assertEqual(list(izip_longest()), zip())
544 self.assertEqual(list(izip_longest([])), zip([]))
545 self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
Tim Petersea5962f2007-03-12 18:07:52 +0000546
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000547 self.assertEqual(list(izip_longest('abc', 'defg', **{})), map(None, 'abc', 'defg')) # empty keyword dict
548 self.assertRaises(TypeError, izip_longest, 3)
549 self.assertRaises(TypeError, izip_longest, range(3), 3)
550
551 for stmt in [
552 "izip_longest('abc', fv=1)",
Tim Petersea5962f2007-03-12 18:07:52 +0000553 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000554 ]:
555 try:
556 eval(stmt, globals(), locals())
557 except TypeError:
558 pass
559 else:
560 self.fail('Did not raise Type in: ' + stmt)
Tim Petersea5962f2007-03-12 18:07:52 +0000561
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000562 # Check tuple re-use (implementation detail)
563 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
564 zip('abc', 'def'))
565 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
566 zip('abc', 'def'))
567 ids = map(id, izip_longest('abc', 'def'))
568 self.assertEqual(min(ids), max(ids))
569 ids = map(id, list(izip_longest('abc', 'def')))
570 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
571
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000572 def test_product(self):
573 for args, result in [
Raymond Hettingerd553d852008-03-04 04:17:08 +0000574 ([], [()]), # zero iterables
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000575 (['ab'], [('a',), ('b',)]), # one iterable
576 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
577 ([range(0), range(2), range(3)], []), # first iterable with zero length
578 ([range(2), range(0), range(3)], []), # middle iterable with zero length
579 ([range(2), range(3), range(0)], []), # last iterable with zero length
580 ]:
581 self.assertEqual(list(product(*args)), result)
Raymond Hettinger08ff6822008-02-29 02:21:48 +0000582 for r in range(4):
583 self.assertEqual(list(product(*(args*r))),
584 list(product(*args, **dict(repeat=r))))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000585 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
586 self.assertRaises(TypeError, product, range(6), None)
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000587
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000588 def product1(*args, **kwds):
589 pools = map(tuple, args) * kwds.get('repeat', 1)
590 n = len(pools)
591 if n == 0:
592 yield ()
593 return
594 if any(len(pool) == 0 for pool in pools):
595 return
596 indices = [0] * n
597 yield tuple(pool[i] for pool, i in zip(pools, indices))
598 while 1:
599 for i in reversed(range(n)): # right to left
600 if indices[i] == len(pools[i]) - 1:
601 continue
602 indices[i] += 1
603 for j in range(i+1, n):
604 indices[j] = 0
605 yield tuple(pool[i] for pool, i in zip(pools, indices))
606 break
607 else:
608 return
609
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000610 def product2(*args, **kwds):
611 'Pure python version used in docs'
612 pools = map(tuple, args) * kwds.get('repeat', 1)
613 result = [[]]
614 for pool in pools:
615 result = [x+[y] for x in result for y in pool]
616 for prod in result:
617 yield tuple(prod)
618
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000619 argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
620 set('abcdefg'), range(11), tuple(range(13))]
621 for i in range(100):
622 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Raymond Hettingerd553d852008-03-04 04:17:08 +0000623 expected_len = prod(map(len, args))
624 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000625 self.assertEqual(list(product(*args)), list(product1(*args)))
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000626 self.assertEqual(list(product(*args)), list(product2(*args)))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000627 args = map(iter, args)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000628 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000629
Raymond Hettinger73d79632008-02-23 02:20:41 +0000630 # Test implementation detail: tuple re-use
631 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
632 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000633
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000634 def test_repeat(self):
635 self.assertEqual(zip(xrange(3),repeat('a')),
636 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000637 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000638 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000639 self.assertEqual(list(repeat('a', 0)), [])
640 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000641 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000642 self.assertRaises(TypeError, repeat, None, 3, 4)
643 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000644 r = repeat(1+0j)
645 self.assertEqual(repr(r), 'repeat((1+0j))')
646 r = repeat(1+0j, 5)
647 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
648 list(r)
649 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000650
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000651 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000652 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
653 [0**1, 1**2, 2**3])
654 self.assertEqual(list(imap(None, 'abc', range(5))),
655 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000656 self.assertEqual(list(imap(None, 'abc', count())),
657 [('a',0),('b',1),('c',2)])
658 self.assertEqual(take(2,imap(None, 'abc', count())),
659 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000660 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000661 self.assertRaises(TypeError, imap)
662 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000663 self.assertRaises(TypeError, imap(10, range(5)).next)
664 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
665 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000666
667 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000668 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
669 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000670 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
671 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000672 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger47317092008-01-17 03:02:14 +0000673 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
674 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000675 self.assertRaises(TypeError, starmap)
676 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
677 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
678 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
679 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000680
681 def test_islice(self):
682 for args in [ # islice(args) should agree with range(args)
683 (10, 20, 3),
684 (10, 3, 20),
685 (10, 20),
686 (10, 3),
687 (20,)
688 ]:
689 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
690
691 for args, tgtargs in [ # Stop when seqn is exhausted
692 ((10, 110, 3), ((10, 100, 3))),
693 ((10, 110), ((10, 100))),
694 ((110,), (100,))
695 ]:
696 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
697
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000698 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000699 self.assertEqual(list(islice(xrange(10), None)), range(10))
Raymond Hettingerb2594052004-12-05 09:25:51 +0000700 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
701 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000702 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
703 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
704
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000705 # Test number of items consumed SF #1171417
706 it = iter(range(10))
707 self.assertEqual(list(islice(it, 3)), range(3))
708 self.assertEqual(list(it), range(3, 10))
709
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000710 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000711 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000712 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
713 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
714 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
715 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
716 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000717 self.assertRaises(ValueError, islice, xrange(10), 'a')
718 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
719 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
720 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
721 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +0000722 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000723
724 def test_takewhile(self):
725 data = [1, 3, 5, 20, 2, 4, 6, 8]
726 underten = lambda x: x<10
727 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000728 self.assertEqual(list(takewhile(underten, [])), [])
729 self.assertRaises(TypeError, takewhile)
730 self.assertRaises(TypeError, takewhile, operator.pow)
731 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
732 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
733 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000734 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
735 self.assertEqual(list(t), [1, 1, 1])
736 self.assertRaises(StopIteration, t.next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000737
738 def test_dropwhile(self):
739 data = [1, 3, 5, 20, 2, 4, 6, 8]
740 underten = lambda x: x<10
741 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000742 self.assertEqual(list(dropwhile(underten, [])), [])
743 self.assertRaises(TypeError, dropwhile)
744 self.assertRaises(TypeError, dropwhile, operator.pow)
745 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
746 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
747 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000748
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000749 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000750 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000751 def irange(n):
752 for i in xrange(n):
753 yield i
754
755 a, b = tee([]) # test empty iterator
756 self.assertEqual(list(a), [])
757 self.assertEqual(list(b), [])
758
759 a, b = tee(irange(n)) # test 100% interleaved
760 self.assertEqual(zip(a,b), zip(range(n),range(n)))
761
762 a, b = tee(irange(n)) # test 0% interleaved
763 self.assertEqual(list(a), range(n))
764 self.assertEqual(list(b), range(n))
765
766 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000767 for i in xrange(100):
768 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000769 del a
770 self.assertEqual(list(b), range(n))
771
772 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000773 for i in xrange(100):
774 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000775 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000776 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000777
778 for j in xrange(5): # test randomly interleaved
779 order = [0]*n + [1]*n
780 random.shuffle(order)
781 lists = ([], [])
782 its = tee(irange(n))
783 for i in order:
784 value = its[i].next()
785 lists[i].append(value)
786 self.assertEqual(lists[0], range(n))
787 self.assertEqual(lists[1], range(n))
788
Raymond Hettingerad983e72003-11-12 14:32:26 +0000789 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000790 self.assertRaises(TypeError, tee)
791 self.assertRaises(TypeError, tee, 3)
792 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000793 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000794
Raymond Hettingerad983e72003-11-12 14:32:26 +0000795 # tee object should be instantiable
796 a, b = tee('abc')
797 c = type(a)('def')
798 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000799
Raymond Hettingerad983e72003-11-12 14:32:26 +0000800 # test long-lagged and multi-way split
801 a, b, c = tee(xrange(2000), 3)
802 for i in xrange(100):
803 self.assertEqual(a.next(), i)
804 self.assertEqual(list(b), range(2000))
805 self.assertEqual([c.next(), c.next()], range(2))
806 self.assertEqual(list(a), range(100,2000))
807 self.assertEqual(list(c), range(2,2000))
808
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000809 # test values of n
810 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Neal Norwitz69e88972006-09-02 02:58:13 +0000811 self.assertRaises(ValueError, tee, [], -1)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000812 for n in xrange(5):
813 result = tee('abc', n)
814 self.assertEqual(type(result), tuple)
815 self.assertEqual(len(result), n)
816 self.assertEqual(map(list, result), [list('abc')]*n)
817
Raymond Hettingerad983e72003-11-12 14:32:26 +0000818 # tee pass-through to copyable iterator
819 a, b = tee('abc')
820 c, d = tee(a)
821 self.assert_(a is c)
822
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000823 # test tee_new
824 t1, t2 = tee('abc')
825 tnew = type(t1)
826 self.assertRaises(TypeError, tnew)
827 self.assertRaises(TypeError, tnew, 10)
828 t3 = tnew(t1)
829 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000830
Raymond Hettingera9f60922004-10-17 16:40:14 +0000831 # test that tee objects are weak referencable
832 a, b = tee(xrange(10))
833 p = proxy(a)
834 self.assertEqual(getattr(p, '__class__'), type(b))
835 del a
836 self.assertRaises(ReferenceError, getattr, p, '__class__')
837
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000838 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000839 self.assertRaises(StopIteration, izip().next)
840
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000841 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000842 self.assertRaises(StopIteration, f([]).next)
843 self.assertRaises(StopIteration, f(StopNow()).next)
844
845 self.assertRaises(StopIteration, islice([], None).next)
846 self.assertRaises(StopIteration, islice(StopNow(), None).next)
847
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000848 p, q = tee([])
849 self.assertRaises(StopIteration, p.next)
850 self.assertRaises(StopIteration, q.next)
851 p, q = tee(StopNow())
852 self.assertRaises(StopIteration, p.next)
853 self.assertRaises(StopIteration, q.next)
854
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000855 self.assertRaises(StopIteration, repeat(None, 0).next)
856
857 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
858 self.assertRaises(StopIteration, f(lambda x:x, []).next)
859 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
860
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000861class TestExamples(unittest.TestCase):
862
863 def test_chain(self):
864 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
865
866 def test_chain_from_iterable(self):
867 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
868
869 def test_combinations(self):
870 self.assertEqual(list(combinations('ABCD', 2)),
871 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
872 self.assertEqual(list(combinations(range(4), 3)),
873 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
874
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000875 def test_combinations_with_replacement(self):
876 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
877 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
878
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000879 def test_compress(self):
880 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
881
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000882 def test_count(self):
883 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
884
885 def test_cycle(self):
886 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
887
888 def test_dropwhile(self):
889 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
890
891 def test_groupby(self):
892 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
893 list('ABCDAB'))
894 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
895 [list('AAAA'), list('BBB'), list('CC'), list('D')])
896
897 def test_ifilter(self):
898 self.assertEqual(list(ifilter(lambda x: x%2, range(10))), [1,3,5,7,9])
899
900 def test_ifilterfalse(self):
901 self.assertEqual(list(ifilterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
902
903 def test_imap(self):
904 self.assertEqual(list(imap(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
905
906 def test_islice(self):
907 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
908 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
909 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
910 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
911
912 def test_izip(self):
913 self.assertEqual(list(izip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
914
915 def test_izip_longest(self):
916 self.assertEqual(list(izip_longest('ABCD', 'xy', fillvalue='-')),
917 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
918
919 def test_permutations(self):
920 self.assertEqual(list(permutations('ABCD', 2)),
921 map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split()))
922 self.assertEqual(list(permutations(range(3))),
923 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
924
925 def test_product(self):
926 self.assertEqual(list(product('ABCD', 'xy')),
927 map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split()))
928 self.assertEqual(list(product(range(2), repeat=3)),
929 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
930 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
931
932 def test_repeat(self):
933 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
934
935 def test_stapmap(self):
936 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
937 [32, 9, 1000])
938
939 def test_takewhile(self):
940 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
941
942
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000943class TestGC(unittest.TestCase):
944
945 def makecycle(self, iterator, container):
946 container.append(iterator)
947 iterator.next()
948 del container, iterator
949
950 def test_chain(self):
951 a = []
952 self.makecycle(chain(a), a)
953
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000954 def test_chain_from_iterable(self):
955 a = []
956 self.makecycle(chain.from_iterable([a]), a)
957
958 def test_combinations(self):
959 a = []
960 self.makecycle(combinations([1,2,a,3], 3), a)
961
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000962 def test_combinations_with_replacement(self):
963 a = []
964 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
965
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000966 def test_compress(self):
967 a = []
968 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
969
Raymond Hettingerb21d8102009-02-16 20:39:12 +0000970 def test_count(self):
971 a = []
972 Int = type('Int', (int,), dict(x=a))
973 self.makecycle(count(Int(0), Int(1)), a)
974
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000975 def test_cycle(self):
976 a = []
977 self.makecycle(cycle([a]*2), a)
978
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000979 def test_dropwhile(self):
980 a = []
981 self.makecycle(dropwhile(bool, [0, a, a]), a)
982
983 def test_groupby(self):
984 a = []
985 self.makecycle(groupby([a]*2, lambda x:x), a)
986
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000987 def test_issue2246(self):
988 # Issue 2246 -- the _grouper iterator was not included in GC
989 n = 10
990 keyfunc = lambda x: x
991 for i, j in groupby(xrange(n), key=keyfunc):
992 keyfunc.__dict__.setdefault('x',[]).append(j)
993
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000994 def test_ifilter(self):
995 a = []
996 self.makecycle(ifilter(lambda x:True, [a]*2), a)
997
998 def test_ifilterfalse(self):
999 a = []
1000 self.makecycle(ifilterfalse(lambda x:False, a), a)
1001
1002 def test_izip(self):
1003 a = []
1004 self.makecycle(izip([a]*2, [a]*3), a)
1005
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001006 def test_izip_longest(self):
1007 a = []
1008 self.makecycle(izip_longest([a]*2, [a]*3), a)
1009 b = [a, None]
1010 self.makecycle(izip_longest([a]*2, [a]*3, fillvalue=b), a)
1011
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001012 def test_imap(self):
1013 a = []
1014 self.makecycle(imap(lambda x:x, [a]*2), a)
1015
1016 def test_islice(self):
1017 a = []
1018 self.makecycle(islice([a]*2, None), a)
1019
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001020 def test_permutations(self):
1021 a = []
1022 self.makecycle(permutations([1,2,a,3], 3), a)
1023
1024 def test_product(self):
1025 a = []
1026 self.makecycle(product([1,2,a,3], repeat=3), a)
1027
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001028 def test_repeat(self):
1029 a = []
1030 self.makecycle(repeat(a), a)
1031
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001032 def test_starmap(self):
1033 a = []
1034 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1035
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001036 def test_takewhile(self):
1037 a = []
1038 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1039
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001040def R(seqn):
1041 'Regular generator'
1042 for i in seqn:
1043 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001044
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001045class G:
1046 'Sequence using __getitem__'
1047 def __init__(self, seqn):
1048 self.seqn = seqn
1049 def __getitem__(self, i):
1050 return self.seqn[i]
1051
1052class I:
1053 'Sequence using iterator protocol'
1054 def __init__(self, seqn):
1055 self.seqn = seqn
1056 self.i = 0
1057 def __iter__(self):
1058 return self
1059 def next(self):
1060 if self.i >= len(self.seqn): raise StopIteration
1061 v = self.seqn[self.i]
1062 self.i += 1
1063 return v
1064
1065class Ig:
1066 'Sequence using iterator protocol defined with a generator'
1067 def __init__(self, seqn):
1068 self.seqn = seqn
1069 self.i = 0
1070 def __iter__(self):
1071 for val in self.seqn:
1072 yield val
1073
1074class X:
1075 'Missing __getitem__ and __iter__'
1076 def __init__(self, seqn):
1077 self.seqn = seqn
1078 self.i = 0
1079 def next(self):
1080 if self.i >= len(self.seqn): raise StopIteration
1081 v = self.seqn[self.i]
1082 self.i += 1
1083 return v
1084
1085class N:
1086 'Iterator missing next()'
1087 def __init__(self, seqn):
1088 self.seqn = seqn
1089 self.i = 0
1090 def __iter__(self):
1091 return self
1092
1093class E:
1094 'Test propagation of exceptions'
1095 def __init__(self, seqn):
1096 self.seqn = seqn
1097 self.i = 0
1098 def __iter__(self):
1099 return self
1100 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001101 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001102
1103class S:
1104 'Test immediate stop'
1105 def __init__(self, seqn):
1106 pass
1107 def __iter__(self):
1108 return self
1109 def next(self):
1110 raise StopIteration
1111
1112def L(seqn):
1113 'Test multiple tiers of iterators'
1114 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
1115
1116
1117class TestVariousIteratorArgs(unittest.TestCase):
1118
1119 def test_chain(self):
1120 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1121 for g in (G, I, Ig, S, L, R):
1122 self.assertEqual(list(chain(g(s))), list(g(s)))
1123 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001124 self.assertRaises(TypeError, list, chain(X(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001125 self.assertRaises(TypeError, list, chain(N(s)))
1126 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1127
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001128 def test_compress(self):
1129 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1130 n = len(s)
1131 for g in (G, I, Ig, S, L, R):
1132 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1133 self.assertRaises(TypeError, compress, X(s), repeat(1))
1134 self.assertRaises(TypeError, list, compress(N(s), repeat(1)))
1135 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1136
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001137 def test_product(self):
1138 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1139 self.assertRaises(TypeError, product, X(s))
1140 self.assertRaises(TypeError, product, N(s))
1141 self.assertRaises(ZeroDivisionError, product, E(s))
1142
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001143 def test_cycle(self):
1144 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1145 for g in (G, I, Ig, S, L, R):
1146 tgtlen = len(s) * 3
1147 expected = list(g(s))*3
1148 actual = list(islice(cycle(g(s)), tgtlen))
1149 self.assertEqual(actual, expected)
1150 self.assertRaises(TypeError, cycle, X(s))
1151 self.assertRaises(TypeError, list, cycle(N(s)))
1152 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1153
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001154 def test_groupby(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([k for k, sb in groupby(g(s))], list(g(s)))
1158 self.assertRaises(TypeError, groupby, X(s))
1159 self.assertRaises(TypeError, list, groupby(N(s)))
1160 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1161
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001162 def test_ifilter(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(ifilter(isEven, g(s))), filter(isEven, g(s)))
1166 self.assertRaises(TypeError, ifilter, isEven, X(s))
1167 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
1168 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
1169
1170 def test_ifilterfalse(self):
1171 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1172 for g in (G, I, Ig, S, L, R):
1173 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
1174 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
1175 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
1176 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
1177
1178 def test_izip(self):
1179 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1180 for g in (G, I, Ig, S, L, R):
1181 self.assertEqual(list(izip(g(s))), zip(g(s)))
1182 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
1183 self.assertRaises(TypeError, izip, X(s))
1184 self.assertRaises(TypeError, list, izip(N(s)))
1185 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
1186
Raymond Hettingerd36862c2007-02-21 05:20:38 +00001187 def test_iziplongest(self):
1188 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1189 for g in (G, I, Ig, S, L, R):
1190 self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
1191 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
1192 self.assertRaises(TypeError, izip_longest, X(s))
1193 self.assertRaises(TypeError, list, izip_longest(N(s)))
1194 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
1195
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001196 def test_imap(self):
1197 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1198 for g in (G, I, Ig, S, L, R):
1199 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
1200 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
1201 self.assertRaises(TypeError, imap, onearg, X(s))
1202 self.assertRaises(TypeError, list, imap(onearg, N(s)))
1203 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
1204
1205 def test_islice(self):
1206 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1207 for g in (G, I, Ig, S, L, R):
1208 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1209 self.assertRaises(TypeError, islice, X(s), 10)
1210 self.assertRaises(TypeError, list, islice(N(s), 10))
1211 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1212
1213 def test_starmap(self):
1214 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1215 for g in (G, I, Ig, S, L, R):
1216 ss = zip(s, s)
1217 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
1218 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
1219 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
1220 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1221
1222 def test_takewhile(self):
1223 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1224 for g in (G, I, Ig, S, L, R):
1225 tgt = []
1226 for elem in g(s):
1227 if not isEven(elem): break
1228 tgt.append(elem)
1229 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1230 self.assertRaises(TypeError, takewhile, isEven, X(s))
1231 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
1232 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1233
1234 def test_dropwhile(self):
1235 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1236 for g in (G, I, Ig, S, L, R):
1237 tgt = []
1238 for elem in g(s):
1239 if not tgt and isOdd(elem): continue
1240 tgt.append(elem)
1241 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1242 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
1243 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
1244 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1245
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001246 def test_tee(self):
1247 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1248 for g in (G, I, Ig, S, L, R):
1249 it1, it2 = tee(g(s))
1250 self.assertEqual(list(it1), list(g(s)))
1251 self.assertEqual(list(it2), list(g(s)))
1252 self.assertRaises(TypeError, tee, X(s))
1253 self.assertRaises(TypeError, list, tee(N(s))[0])
1254 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1255
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001256class LengthTransparency(unittest.TestCase):
1257
1258 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001259 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001260 self.assertEqual(len(repeat(None, 50)), 50)
1261 self.assertRaises(TypeError, len, repeat(None))
1262
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001263class RegressionTests(unittest.TestCase):
1264
1265 def test_sf_793826(self):
1266 # Fix Armin Rigo's successful efforts to wreak havoc
1267
1268 def mutatingtuple(tuple1, f, tuple2):
1269 # this builds a tuple t which is a copy of tuple1,
1270 # then calls f(t), then mutates t to be equal to tuple2
1271 # (needs len(tuple1) == len(tuple2)).
1272 def g(value, first=[1]):
1273 if first:
1274 del first[:]
1275 f(z.next())
1276 return value
1277 items = list(tuple2)
1278 items[1:1] = list(tuple1)
1279 gen = imap(g, items)
1280 z = izip(*[gen]*len(tuple1))
1281 z.next()
1282
1283 def f(t):
1284 global T
1285 T = t
1286 first[:] = list(T)
1287
1288 first = []
1289 mutatingtuple((1,2,3), f, (4,5,6))
1290 second = list(T)
1291 self.assertEqual(first, second)
1292
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001293
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001294 def test_sf_950057(self):
1295 # Make sure that chain() and cycle() catch exceptions immediately
1296 # rather than when shifting between input sources
1297
1298 def gen1():
1299 hist.append(0)
1300 yield 1
1301 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001302 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001303 hist.append(2)
1304
1305 def gen2(x):
1306 hist.append(3)
1307 yield 2
1308 hist.append(4)
1309 if x:
1310 raise StopIteration
1311
1312 hist = []
1313 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1314 self.assertEqual(hist, [0,1])
1315
1316 hist = []
1317 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1318 self.assertEqual(hist, [0,1])
1319
1320 hist = []
1321 self.assertRaises(AssertionError, list, cycle(gen1()))
1322 self.assertEqual(hist, [0,1])
1323
Georg Brandlb84c1372007-01-21 10:28:43 +00001324class SubclassWithKwargsTest(unittest.TestCase):
1325 def test_keywords_in_subclass(self):
1326 # count is not subclassable...
1327 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001328 starmap, islice, takewhile, dropwhile, cycle, compress):
Georg Brandlb84c1372007-01-21 10:28:43 +00001329 class Subclass(cls):
1330 def __init__(self, newarg=None, *args):
1331 cls.__init__(self, *args)
1332 try:
1333 Subclass(newarg=1)
1334 except TypeError, err:
1335 # we expect type errors because of wrong argument count
1336 self.failIf("does not take keyword arguments" in err.args[0])
1337
1338
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001339libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001340
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001341
1342>>> amounts = [120.15, 764.05, 823.14]
1343>>> for checknum, amount in izip(count(1200), amounts):
1344... print 'Check %d is for $%.2f' % (checknum, amount)
1345...
1346Check 1200 is for $120.15
1347Check 1201 is for $764.05
1348Check 1202 is for $823.14
1349
1350>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001351>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1352... print cube
1353...
13541
13558
135627
1357
1358>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001359>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001360... print name.title()
1361...
1362Alex
1363Laura
1364Martin
1365Walter
1366Samuele
1367
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001368>>> from operator import itemgetter
1369>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Armin Rigoa3f09272006-05-28 19:13:17 +00001370>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001371>>> for k, g in groupby(di, itemgetter(1)):
1372... print k, map(itemgetter(0), g)
1373...
13741 ['a', 'c', 'e']
13752 ['b', 'd', 'f']
13763 ['g']
1377
Raymond Hettinger734fb572004-01-20 20:04:40 +00001378# Find runs of consecutive numbers using groupby. The key to the solution
1379# is differencing with a range so that consecutive numbers all appear in
1380# same group.
1381>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
1382>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
1383... print map(operator.itemgetter(1), g)
1384...
1385[1]
1386[4, 5, 6]
1387[10]
1388[15, 16, 17, 18]
1389[22]
1390[25, 26, 27, 28]
1391
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001392>>> def take(n, iterable):
1393... "Return first n items of the iterable as a list"
1394... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001395
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001396>>> def enumerate(iterable, start=0):
1397... return izip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001398
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001399>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001400... "Return function(0), function(1), ..."
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001401... return imap(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001402
1403>>> def nth(iterable, n):
Raymond Hettingerd507afd2009-02-04 10:52:32 +00001404... "Returns the nth item or None"
1405... return next(islice(iterable, n, None), None)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001406
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001407>>> def quantify(iterable, pred=bool):
1408... "Count how many times the predicate is true"
1409... return sum(imap(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001410
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001411>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001412... "Returns the sequence elements and then returns None indefinitely"
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001413... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001414
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001415>>> def ncycles(iterable, n):
1416... "Returns the seqeuence elements n times"
1417... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001418
1419>>> def dotproduct(vec1, vec2):
1420... return sum(imap(operator.mul, vec1, vec2))
1421
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001422>>> def flatten(listOfLists):
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001423... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001424
1425>>> def repeatfunc(func, times=None, *args):
1426... "Repeat calls to func with specified arguments."
1427... " Example: repeatfunc(random.random)"
1428... if times is None:
1429... return starmap(func, repeat(args))
1430... else:
1431... return starmap(func, repeat(args, times))
1432
Raymond Hettingerd591f662003-10-26 15:34:50 +00001433>>> def pairwise(iterable):
1434... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1435... a, b = tee(iterable)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001436... for elem in b:
1437... break
Raymond Hettingerad983e72003-11-12 14:32:26 +00001438... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001439
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001440>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001441... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001442... args = [iter(iterable)] * n
Raymond Hettingerf080e6d2008-07-31 01:19:50 +00001443... return izip_longest(fillvalue=fillvalue, *args)
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001444
1445>>> def roundrobin(*iterables):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001446... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001447... # Recipe credited to George Sakkis
1448... pending = len(iterables)
1449... nexts = cycle(iter(it).next for it in iterables)
1450... while pending:
1451... try:
1452... for next in nexts:
1453... yield next()
1454... except StopIteration:
1455... pending -= 1
1456... nexts = cycle(islice(nexts, pending))
1457
1458>>> def powerset(iterable):
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001459... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1460... s = list(iterable)
1461... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001462
Raymond Hettinger44e15812009-01-02 21:26:45 +00001463>>> def unique_everseen(iterable, key=None):
1464... "List unique elements, preserving order. Remember all elements ever seen."
1465... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1466... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1467... seen = set()
1468... seen_add = seen.add
1469... if key is None:
1470... for element in iterable:
1471... if element not in seen:
1472... seen_add(element)
1473... yield element
1474... else:
1475... for element in iterable:
1476... k = key(element)
1477... if k not in seen:
1478... seen_add(k)
1479... yield element
1480
1481>>> def unique_justseen(iterable, key=None):
1482... "List unique elements, preserving order. Remember only the element just seen."
1483... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1484... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1485... return imap(next, imap(itemgetter(1), groupby(iterable, key)))
1486
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001487This is not part of the examples but it tests to make sure the definitions
1488perform as purported.
1489
Raymond Hettingera098b332003-09-08 23:58:40 +00001490>>> take(10, count())
1491[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1492
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001493>>> list(enumerate('abc'))
1494[(0, 'a'), (1, 'b'), (2, 'c')]
1495
1496>>> list(islice(tabulate(lambda x: 2*x), 4))
1497[0, 2, 4, 6]
1498
1499>>> nth('abcde', 3)
Raymond Hettingerd507afd2009-02-04 10:52:32 +00001500'd'
1501
1502>>> nth('abcde', 9) is None
1503True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001504
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001505>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000150650
1507
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001508>>> a = [[1, 2, 3], [4, 5, 6]]
1509>>> flatten(a)
1510[1, 2, 3, 4, 5, 6]
1511
1512>>> list(repeatfunc(pow, 5, 2, 3))
1513[8, 8, 8, 8, 8]
1514
1515>>> import random
1516>>> take(5, imap(int, repeatfunc(random.random)))
1517[0, 0, 0, 0, 0]
1518
Raymond Hettingerd591f662003-10-26 15:34:50 +00001519>>> list(pairwise('abcd'))
1520[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001521
Raymond Hettingerd591f662003-10-26 15:34:50 +00001522>>> list(pairwise([]))
1523[]
1524
1525>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001526[]
1527
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001528>>> list(islice(padnone('abc'), 0, 6))
1529['a', 'b', 'c', None, None, None]
1530
1531>>> list(ncycles('abc', 3))
1532['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1533
1534>>> dotproduct([1,2,3], [4,5,6])
153532
1536
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001537>>> list(grouper(3, 'abcdefg', 'x'))
1538[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1539
1540>>> list(roundrobin('abc', 'd', 'ef'))
1541['a', 'd', 'e', 'b', 'f', 'c']
1542
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001543>>> list(powerset([1,2,3]))
1544[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001545
Raymond Hettinger560f9a82009-01-27 13:26:35 +00001546>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1547True
1548
1549>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1550True
1551
Raymond Hettinger44e15812009-01-02 21:26:45 +00001552>>> list(unique_everseen('AAAABBBCCDAABBB'))
1553['A', 'B', 'C', 'D']
1554
1555>>> list(unique_everseen('ABBCcAD', str.lower))
1556['A', 'B', 'C', 'D']
1557
1558>>> list(unique_justseen('AAAABBBCCDAABBB'))
1559['A', 'B', 'C', 'D', 'A', 'B']
1560
1561>>> list(unique_justseen('ABBCcAD', str.lower))
1562['A', 'B', 'C', 'A', 'D']
1563
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001564"""
1565
1566__test__ = {'libreftest' : libreftest}
1567
1568def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001569 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Georg Brandlb84c1372007-01-21 10:28:43 +00001570 RegressionTests, LengthTransparency,
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001571 SubclassWithKwargsTest, TestExamples)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001572 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001573
1574 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001575 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001576 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001577 counts = [None] * 5
1578 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001579 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001580 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001581 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001582 print counts
1583
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001584 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001585 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001586
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001587if __name__ == "__main__":
1588 test_main(verbose=True)