blob: f8bdb8bdec9fe713d0c28c5ac9a2aaca479c186b [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
Benjamin Peterson5c8da862009-06-30 22:57:08 +0000130 self.assertTrue(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:
Benjamin Peterson5c8da862009-06-30 22:57:08 +0000194 self.assertTrue(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000195
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
Benjamin Peterson5c8da862009-06-30 22:57:08 +0000201 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000202 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
Benjamin Peterson5c8da862009-06-30 22:57:08 +0000264 self.assertTrue(all(e in values for e in p)) # elements taken from input iterable
Raymond Hettingerd553d852008-03-04 04:17:08 +0000265 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 Hettingeraa681c72009-02-21 07:17:22 +0000353 self.assertEqual(zip('abc',count(step=-1)),
354 [('a', 0), ('b', -1), ('c', -2)])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000355 self.assertEqual(zip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
356 self.assertEqual(zip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
357 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
358 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
359 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettingeraa044612009-02-12 12:04:26 +0000360 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
361 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerdbe3bfb2009-02-12 12:43:01 +0000362 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
363 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000364 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
365 c = count(3, 5)
366 self.assertEqual(repr(c), 'count(3, 5)')
367 c.next()
368 self.assertEqual(repr(c), 'count(8, 5)')
369 c = count(-9, 0)
370 self.assertEqual(repr(c), 'count(-9, 0)')
371 c.next()
372 self.assertEqual(repr(c), 'count(-9, 0)')
373 c = count(-9, -3)
374 self.assertEqual(repr(c), 'count(-9, -3)')
375 c.next()
376 self.assertEqual(repr(c), 'count(-12, -3)')
377 self.assertEqual(repr(c), 'count(-12, -3)')
378 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
379 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
380 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
381 for i in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 10, sys.maxint-5, sys.maxint+5):
382 for j in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 1, 10, sys.maxint-5, sys.maxint+5):
383 # Test repr (ignoring the L in longs)
384 r1 = repr(count(i, j)).replace('L', '')
385 if j == 1:
386 r2 = ('count(%r)' % i).replace('L', '')
387 else:
388 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
389 self.assertEqual(r1, r2)
390
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000391 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000392 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000393 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000394 self.assertRaises(TypeError, cycle)
395 self.assertRaises(TypeError, cycle, 5)
396 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000397
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000398 def test_groupby(self):
399 # Check whether it accepts arguments correctly
400 self.assertEqual([], list(groupby([])))
401 self.assertEqual([], list(groupby([], key=id)))
402 self.assertRaises(TypeError, list, groupby('abc', []))
403 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000404 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000405
406 # Check normal input
407 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
408 (2,15,22), (3,16,23), (3,17,23)]
409 dup = []
410 for k, g in groupby(s, lambda r:r[0]):
411 for elem in g:
412 self.assertEqual(k, elem[0])
413 dup.append(elem)
414 self.assertEqual(s, dup)
415
416 # Check nested case
417 dup = []
418 for k, g in groupby(s, lambda r:r[0]):
419 for ik, ig in groupby(g, lambda r:r[2]):
420 for elem in ig:
421 self.assertEqual(k, elem[0])
422 self.assertEqual(ik, elem[2])
423 dup.append(elem)
424 self.assertEqual(s, dup)
425
426 # Check case where inner iterator is not used
427 keys = [k for k, g in groupby(s, lambda r:r[0])]
428 expectedkeys = set([r[0] for r in s])
429 self.assertEqual(set(keys), expectedkeys)
430 self.assertEqual(len(keys), len(expectedkeys))
431
432 # Exercise pipes and filters style
433 s = 'abracadabra'
434 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000435 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000436 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
437 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000438 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000439 self.assertEqual(r, ['a', 'b', 'r'])
440 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000441 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000442 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
443 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000444 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000445 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
446
447 # iter.next failure
448 class ExpectedError(Exception):
449 pass
450 def delayed_raise(n=0):
451 for i in range(n):
452 yield 'yo'
453 raise ExpectedError
454 def gulp(iterable, keyp=None, func=list):
455 return [func(g) for k, g in groupby(iterable, keyp)]
456
457 # iter.next failure on outer object
458 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
459 # iter.next failure on inner object
460 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
461
462 # __cmp__ failure
463 class DummyCmp:
464 def __cmp__(self, dst):
465 raise ExpectedError
466 s = [DummyCmp(), DummyCmp(), None]
467
468 # __cmp__ failure on outer object
469 self.assertRaises(ExpectedError, gulp, s, func=id)
470 # __cmp__ failure on inner object
471 self.assertRaises(ExpectedError, gulp, s)
472
473 # keyfunc failure
474 def keyfunc(obj):
475 if keyfunc.skip > 0:
476 keyfunc.skip -= 1
477 return obj
478 else:
479 raise ExpectedError
480
481 # keyfunc failure on outer object
482 keyfunc.skip = 0
483 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
484 keyfunc.skip = 1
485 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
486
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000487 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000488 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000489 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000490 self.assertEqual(list(ifilter(bool, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000491 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000492 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000493 self.assertRaises(TypeError, ifilter, lambda x:x)
494 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000495 self.assertRaises(TypeError, ifilter, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000496 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000497
498 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000499 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
500 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000501 self.assertEqual(list(ifilterfalse(bool, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000502 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000503 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000504 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
505 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000506 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000507 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000508
509 def test_izip(self):
510 ans = [(x,y) for x, y in izip('abc',count())]
511 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000512 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
513 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
Raymond Hettinger02420702003-06-29 20:36:23 +0000514 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000515 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000516 self.assertEqual(list(izip()), zip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000517 self.assertRaises(TypeError, izip, 3)
518 self.assertRaises(TypeError, izip, range(3), 3)
519 # Check tuple re-use (implementation detail)
520 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
521 zip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000522 self.assertEqual([pair for pair in izip('abc', 'def')],
523 zip('abc', 'def'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000524 ids = map(id, izip('abc', 'def'))
525 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000526 ids = map(id, list(izip('abc', 'def')))
527 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000528
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000529 def test_iziplongest(self):
530 for args in [
531 ['abc', range(6)],
532 [range(6), 'abc'],
533 [range(1000), range(2000,2100), range(3000,3050)],
534 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
535 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
536 ]:
537 target = map(None, *args)
538 self.assertEqual(list(izip_longest(*args)), target)
539 self.assertEqual(list(izip_longest(*args, **{})), target)
540 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
541 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Tim Petersea5962f2007-03-12 18:07:52 +0000542
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000543 self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
544
545 self.assertEqual(list(izip_longest()), zip())
546 self.assertEqual(list(izip_longest([])), zip([]))
547 self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
Tim Petersea5962f2007-03-12 18:07:52 +0000548
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000549 self.assertEqual(list(izip_longest('abc', 'defg', **{})), map(None, 'abc', 'defg')) # empty keyword dict
550 self.assertRaises(TypeError, izip_longest, 3)
551 self.assertRaises(TypeError, izip_longest, range(3), 3)
552
553 for stmt in [
554 "izip_longest('abc', fv=1)",
Tim Petersea5962f2007-03-12 18:07:52 +0000555 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000556 ]:
557 try:
558 eval(stmt, globals(), locals())
559 except TypeError:
560 pass
561 else:
562 self.fail('Did not raise Type in: ' + stmt)
Tim Petersea5962f2007-03-12 18:07:52 +0000563
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000564 # Check tuple re-use (implementation detail)
565 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
566 zip('abc', 'def'))
567 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
568 zip('abc', 'def'))
569 ids = map(id, izip_longest('abc', 'def'))
570 self.assertEqual(min(ids), max(ids))
571 ids = map(id, list(izip_longest('abc', 'def')))
572 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
573
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +0000574 def test_bug_7244(self):
575
576 class Repeater(object):
577 # this class is similar to itertools.repeat
578 def __init__(self, o, t, e):
579 self.o = o
580 self.t = int(t)
581 self.e = e
582 def __iter__(self): # its iterator is itself
583 return self
584 def next(self):
585 if self.t > 0:
586 self.t -= 1
587 return self.o
588 else:
589 raise self.e
590
591 # Formerly this code in would fail in debug mode
592 # with Undetected Error and Stop Iteration
593 r1 = Repeater(1, 3, StopIteration)
594 r2 = Repeater(2, 4, StopIteration)
595 def run(r1, r2):
596 result = []
597 for i, j in izip_longest(r1, r2, fillvalue=0):
598 with test_support.captured_output('stdout'):
599 print (i, j)
600 result.append((i, j))
601 return result
602 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
603
604 # Formerly, the RuntimeError would be lost
605 # and StopIteration would stop as expected
606 r1 = Repeater(1, 3, RuntimeError)
607 r2 = Repeater(2, 4, StopIteration)
608 it = izip_longest(r1, r2, fillvalue=0)
609 self.assertEqual(next(it), (1, 2))
610 self.assertEqual(next(it), (1, 2))
611 self.assertEqual(next(it), (1, 2))
612 self.assertRaises(RuntimeError, next, it)
613
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000614 def test_product(self):
615 for args, result in [
Raymond Hettingerd553d852008-03-04 04:17:08 +0000616 ([], [()]), # zero iterables
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000617 (['ab'], [('a',), ('b',)]), # one iterable
618 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
619 ([range(0), range(2), range(3)], []), # first iterable with zero length
620 ([range(2), range(0), range(3)], []), # middle iterable with zero length
621 ([range(2), range(3), range(0)], []), # last iterable with zero length
622 ]:
623 self.assertEqual(list(product(*args)), result)
Raymond Hettinger08ff6822008-02-29 02:21:48 +0000624 for r in range(4):
625 self.assertEqual(list(product(*(args*r))),
626 list(product(*args, **dict(repeat=r))))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000627 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
628 self.assertRaises(TypeError, product, range(6), None)
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000629
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000630 def product1(*args, **kwds):
631 pools = map(tuple, args) * kwds.get('repeat', 1)
632 n = len(pools)
633 if n == 0:
634 yield ()
635 return
636 if any(len(pool) == 0 for pool in pools):
637 return
638 indices = [0] * n
639 yield tuple(pool[i] for pool, i in zip(pools, indices))
640 while 1:
641 for i in reversed(range(n)): # right to left
642 if indices[i] == len(pools[i]) - 1:
643 continue
644 indices[i] += 1
645 for j in range(i+1, n):
646 indices[j] = 0
647 yield tuple(pool[i] for pool, i in zip(pools, indices))
648 break
649 else:
650 return
651
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000652 def product2(*args, **kwds):
653 'Pure python version used in docs'
654 pools = map(tuple, args) * kwds.get('repeat', 1)
655 result = [[]]
656 for pool in pools:
657 result = [x+[y] for x in result for y in pool]
658 for prod in result:
659 yield tuple(prod)
660
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000661 argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
662 set('abcdefg'), range(11), tuple(range(13))]
663 for i in range(100):
664 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Raymond Hettingerd553d852008-03-04 04:17:08 +0000665 expected_len = prod(map(len, args))
666 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000667 self.assertEqual(list(product(*args)), list(product1(*args)))
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000668 self.assertEqual(list(product(*args)), list(product2(*args)))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000669 args = map(iter, args)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000670 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000671
Raymond Hettinger73d79632008-02-23 02:20:41 +0000672 # Test implementation detail: tuple re-use
673 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
674 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000675
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000676 def test_repeat(self):
Raymond Hettinger182edae2009-02-19 02:38:25 +0000677 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000678 self.assertEqual(zip(xrange(3),repeat('a')),
679 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000680 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000681 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000682 self.assertEqual(list(repeat('a', 0)), [])
683 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000684 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000685 self.assertRaises(TypeError, repeat, None, 3, 4)
686 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000687 r = repeat(1+0j)
688 self.assertEqual(repr(r), 'repeat((1+0j))')
689 r = repeat(1+0j, 5)
690 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
691 list(r)
692 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000693
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000694 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000695 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
696 [0**1, 1**2, 2**3])
697 self.assertEqual(list(imap(None, 'abc', range(5))),
698 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000699 self.assertEqual(list(imap(None, 'abc', count())),
700 [('a',0),('b',1),('c',2)])
701 self.assertEqual(take(2,imap(None, 'abc', count())),
702 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000703 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000704 self.assertRaises(TypeError, imap)
705 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000706 self.assertRaises(TypeError, imap(10, range(5)).next)
707 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
708 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000709
710 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000711 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
712 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000713 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
714 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000715 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger47317092008-01-17 03:02:14 +0000716 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
717 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000718 self.assertRaises(TypeError, starmap)
719 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
720 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
721 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
722 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000723
724 def test_islice(self):
725 for args in [ # islice(args) should agree with range(args)
726 (10, 20, 3),
727 (10, 3, 20),
728 (10, 20),
729 (10, 3),
730 (20,)
731 ]:
732 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
733
734 for args, tgtargs in [ # Stop when seqn is exhausted
735 ((10, 110, 3), ((10, 100, 3))),
736 ((10, 110), ((10, 100))),
737 ((110,), (100,))
738 ]:
739 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
740
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000741 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000742 self.assertEqual(list(islice(xrange(10), None)), range(10))
Raymond Hettingerb2594052004-12-05 09:25:51 +0000743 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
744 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000745 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
746 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
747
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000748 # Test number of items consumed SF #1171417
749 it = iter(range(10))
750 self.assertEqual(list(islice(it, 3)), range(3))
751 self.assertEqual(list(it), range(3, 10))
752
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000753 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000754 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000755 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
756 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
757 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
758 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
759 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000760 self.assertRaises(ValueError, islice, xrange(10), 'a')
761 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
762 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
763 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
764 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +0000765 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000766
767 def test_takewhile(self):
768 data = [1, 3, 5, 20, 2, 4, 6, 8]
769 underten = lambda x: x<10
770 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000771 self.assertEqual(list(takewhile(underten, [])), [])
772 self.assertRaises(TypeError, takewhile)
773 self.assertRaises(TypeError, takewhile, operator.pow)
774 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
775 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
776 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000777 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
778 self.assertEqual(list(t), [1, 1, 1])
779 self.assertRaises(StopIteration, t.next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000780
781 def test_dropwhile(self):
782 data = [1, 3, 5, 20, 2, 4, 6, 8]
783 underten = lambda x: x<10
784 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000785 self.assertEqual(list(dropwhile(underten, [])), [])
786 self.assertRaises(TypeError, dropwhile)
787 self.assertRaises(TypeError, dropwhile, operator.pow)
788 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
789 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
790 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000791
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000792 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000793 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000794 def irange(n):
795 for i in xrange(n):
796 yield i
797
798 a, b = tee([]) # test empty iterator
799 self.assertEqual(list(a), [])
800 self.assertEqual(list(b), [])
801
802 a, b = tee(irange(n)) # test 100% interleaved
803 self.assertEqual(zip(a,b), zip(range(n),range(n)))
804
805 a, b = tee(irange(n)) # test 0% interleaved
806 self.assertEqual(list(a), range(n))
807 self.assertEqual(list(b), range(n))
808
809 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000810 for i in xrange(100):
811 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000812 del a
813 self.assertEqual(list(b), range(n))
814
815 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000816 for i in xrange(100):
817 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000818 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000819 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000820
821 for j in xrange(5): # test randomly interleaved
822 order = [0]*n + [1]*n
823 random.shuffle(order)
824 lists = ([], [])
825 its = tee(irange(n))
826 for i in order:
827 value = its[i].next()
828 lists[i].append(value)
829 self.assertEqual(lists[0], range(n))
830 self.assertEqual(lists[1], range(n))
831
Raymond Hettingerad983e72003-11-12 14:32:26 +0000832 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000833 self.assertRaises(TypeError, tee)
834 self.assertRaises(TypeError, tee, 3)
835 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000836 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000837
Raymond Hettingerad983e72003-11-12 14:32:26 +0000838 # tee object should be instantiable
839 a, b = tee('abc')
840 c = type(a)('def')
841 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000842
Raymond Hettingerad983e72003-11-12 14:32:26 +0000843 # test long-lagged and multi-way split
844 a, b, c = tee(xrange(2000), 3)
845 for i in xrange(100):
846 self.assertEqual(a.next(), i)
847 self.assertEqual(list(b), range(2000))
848 self.assertEqual([c.next(), c.next()], range(2))
849 self.assertEqual(list(a), range(100,2000))
850 self.assertEqual(list(c), range(2,2000))
851
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000852 # test values of n
853 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Neal Norwitz69e88972006-09-02 02:58:13 +0000854 self.assertRaises(ValueError, tee, [], -1)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000855 for n in xrange(5):
856 result = tee('abc', n)
857 self.assertEqual(type(result), tuple)
858 self.assertEqual(len(result), n)
859 self.assertEqual(map(list, result), [list('abc')]*n)
860
Raymond Hettingerad983e72003-11-12 14:32:26 +0000861 # tee pass-through to copyable iterator
862 a, b = tee('abc')
863 c, d = tee(a)
Benjamin Peterson5c8da862009-06-30 22:57:08 +0000864 self.assertTrue(a is c)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000865
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000866 # test tee_new
867 t1, t2 = tee('abc')
868 tnew = type(t1)
869 self.assertRaises(TypeError, tnew)
870 self.assertRaises(TypeError, tnew, 10)
871 t3 = tnew(t1)
Benjamin Peterson5c8da862009-06-30 22:57:08 +0000872 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000873
Raymond Hettingera9f60922004-10-17 16:40:14 +0000874 # test that tee objects are weak referencable
875 a, b = tee(xrange(10))
876 p = proxy(a)
877 self.assertEqual(getattr(p, '__class__'), type(b))
878 del a
879 self.assertRaises(ReferenceError, getattr, p, '__class__')
880
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000881 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000882 self.assertRaises(StopIteration, izip().next)
883
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000884 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000885 self.assertRaises(StopIteration, f([]).next)
886 self.assertRaises(StopIteration, f(StopNow()).next)
887
888 self.assertRaises(StopIteration, islice([], None).next)
889 self.assertRaises(StopIteration, islice(StopNow(), None).next)
890
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000891 p, q = tee([])
892 self.assertRaises(StopIteration, p.next)
893 self.assertRaises(StopIteration, q.next)
894 p, q = tee(StopNow())
895 self.assertRaises(StopIteration, p.next)
896 self.assertRaises(StopIteration, q.next)
897
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000898 self.assertRaises(StopIteration, repeat(None, 0).next)
899
900 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
901 self.assertRaises(StopIteration, f(lambda x:x, []).next)
902 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
903
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000904class TestExamples(unittest.TestCase):
905
906 def test_chain(self):
907 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
908
909 def test_chain_from_iterable(self):
910 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
911
912 def test_combinations(self):
913 self.assertEqual(list(combinations('ABCD', 2)),
914 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
915 self.assertEqual(list(combinations(range(4), 3)),
916 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
917
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000918 def test_combinations_with_replacement(self):
919 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
920 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
921
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000922 def test_compress(self):
923 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
924
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000925 def test_count(self):
926 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
927
928 def test_cycle(self):
929 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
930
931 def test_dropwhile(self):
932 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
933
934 def test_groupby(self):
935 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
936 list('ABCDAB'))
937 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
938 [list('AAAA'), list('BBB'), list('CC'), list('D')])
939
940 def test_ifilter(self):
941 self.assertEqual(list(ifilter(lambda x: x%2, range(10))), [1,3,5,7,9])
942
943 def test_ifilterfalse(self):
944 self.assertEqual(list(ifilterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
945
946 def test_imap(self):
947 self.assertEqual(list(imap(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
948
949 def test_islice(self):
950 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
951 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
952 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
953 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
954
955 def test_izip(self):
956 self.assertEqual(list(izip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
957
958 def test_izip_longest(self):
959 self.assertEqual(list(izip_longest('ABCD', 'xy', fillvalue='-')),
960 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
961
962 def test_permutations(self):
963 self.assertEqual(list(permutations('ABCD', 2)),
964 map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split()))
965 self.assertEqual(list(permutations(range(3))),
966 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
967
968 def test_product(self):
969 self.assertEqual(list(product('ABCD', 'xy')),
970 map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split()))
971 self.assertEqual(list(product(range(2), repeat=3)),
972 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
973 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
974
975 def test_repeat(self):
976 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
977
978 def test_stapmap(self):
979 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
980 [32, 9, 1000])
981
982 def test_takewhile(self):
983 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
984
985
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000986class TestGC(unittest.TestCase):
987
988 def makecycle(self, iterator, container):
989 container.append(iterator)
990 iterator.next()
991 del container, iterator
992
993 def test_chain(self):
994 a = []
995 self.makecycle(chain(a), a)
996
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000997 def test_chain_from_iterable(self):
998 a = []
999 self.makecycle(chain.from_iterable([a]), a)
1000
1001 def test_combinations(self):
1002 a = []
1003 self.makecycle(combinations([1,2,a,3], 3), a)
1004
Raymond Hettingerd081abc2009-01-27 02:58:49 +00001005 def test_combinations_with_replacement(self):
1006 a = []
1007 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1008
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001009 def test_compress(self):
1010 a = []
1011 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1012
Raymond Hettingerb21d8102009-02-16 20:39:12 +00001013 def test_count(self):
1014 a = []
1015 Int = type('Int', (int,), dict(x=a))
1016 self.makecycle(count(Int(0), Int(1)), a)
1017
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001018 def test_cycle(self):
1019 a = []
1020 self.makecycle(cycle([a]*2), a)
1021
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001022 def test_dropwhile(self):
1023 a = []
1024 self.makecycle(dropwhile(bool, [0, a, a]), a)
1025
1026 def test_groupby(self):
1027 a = []
1028 self.makecycle(groupby([a]*2, lambda x:x), a)
1029
Raymond Hettingera1ca94a2008-03-06 22:51:36 +00001030 def test_issue2246(self):
1031 # Issue 2246 -- the _grouper iterator was not included in GC
1032 n = 10
1033 keyfunc = lambda x: x
1034 for i, j in groupby(xrange(n), key=keyfunc):
1035 keyfunc.__dict__.setdefault('x',[]).append(j)
1036
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001037 def test_ifilter(self):
1038 a = []
1039 self.makecycle(ifilter(lambda x:True, [a]*2), a)
1040
1041 def test_ifilterfalse(self):
1042 a = []
1043 self.makecycle(ifilterfalse(lambda x:False, a), a)
1044
1045 def test_izip(self):
1046 a = []
1047 self.makecycle(izip([a]*2, [a]*3), a)
1048
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001049 def test_izip_longest(self):
1050 a = []
1051 self.makecycle(izip_longest([a]*2, [a]*3), a)
1052 b = [a, None]
1053 self.makecycle(izip_longest([a]*2, [a]*3, fillvalue=b), a)
1054
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001055 def test_imap(self):
1056 a = []
1057 self.makecycle(imap(lambda x:x, [a]*2), a)
1058
1059 def test_islice(self):
1060 a = []
1061 self.makecycle(islice([a]*2, None), a)
1062
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001063 def test_permutations(self):
1064 a = []
1065 self.makecycle(permutations([1,2,a,3], 3), a)
1066
1067 def test_product(self):
1068 a = []
1069 self.makecycle(product([1,2,a,3], repeat=3), a)
1070
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001071 def test_repeat(self):
1072 a = []
1073 self.makecycle(repeat(a), a)
1074
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001075 def test_starmap(self):
1076 a = []
1077 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1078
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001079 def test_takewhile(self):
1080 a = []
1081 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1082
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001083def R(seqn):
1084 'Regular generator'
1085 for i in seqn:
1086 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001087
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001088class G:
1089 'Sequence using __getitem__'
1090 def __init__(self, seqn):
1091 self.seqn = seqn
1092 def __getitem__(self, i):
1093 return self.seqn[i]
1094
1095class I:
1096 'Sequence using iterator protocol'
1097 def __init__(self, seqn):
1098 self.seqn = seqn
1099 self.i = 0
1100 def __iter__(self):
1101 return self
1102 def next(self):
1103 if self.i >= len(self.seqn): raise StopIteration
1104 v = self.seqn[self.i]
1105 self.i += 1
1106 return v
1107
1108class Ig:
1109 'Sequence using iterator protocol defined with a generator'
1110 def __init__(self, seqn):
1111 self.seqn = seqn
1112 self.i = 0
1113 def __iter__(self):
1114 for val in self.seqn:
1115 yield val
1116
1117class X:
1118 'Missing __getitem__ and __iter__'
1119 def __init__(self, seqn):
1120 self.seqn = seqn
1121 self.i = 0
1122 def next(self):
1123 if self.i >= len(self.seqn): raise StopIteration
1124 v = self.seqn[self.i]
1125 self.i += 1
1126 return v
1127
1128class N:
1129 'Iterator missing next()'
1130 def __init__(self, seqn):
1131 self.seqn = seqn
1132 self.i = 0
1133 def __iter__(self):
1134 return self
1135
1136class E:
1137 'Test propagation of exceptions'
1138 def __init__(self, seqn):
1139 self.seqn = seqn
1140 self.i = 0
1141 def __iter__(self):
1142 return self
1143 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001144 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001145
1146class S:
1147 'Test immediate stop'
1148 def __init__(self, seqn):
1149 pass
1150 def __iter__(self):
1151 return self
1152 def next(self):
1153 raise StopIteration
1154
1155def L(seqn):
1156 'Test multiple tiers of iterators'
1157 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
1158
1159
1160class TestVariousIteratorArgs(unittest.TestCase):
1161
1162 def test_chain(self):
1163 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1164 for g in (G, I, Ig, S, L, R):
1165 self.assertEqual(list(chain(g(s))), list(g(s)))
1166 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001167 self.assertRaises(TypeError, list, chain(X(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001168 self.assertRaises(TypeError, list, chain(N(s)))
1169 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1170
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001171 def test_compress(self):
1172 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1173 n = len(s)
1174 for g in (G, I, Ig, S, L, R):
1175 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1176 self.assertRaises(TypeError, compress, X(s), repeat(1))
1177 self.assertRaises(TypeError, list, compress(N(s), repeat(1)))
1178 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1179
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001180 def test_product(self):
1181 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1182 self.assertRaises(TypeError, product, X(s))
1183 self.assertRaises(TypeError, product, N(s))
1184 self.assertRaises(ZeroDivisionError, product, E(s))
1185
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001186 def test_cycle(self):
1187 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1188 for g in (G, I, Ig, S, L, R):
1189 tgtlen = len(s) * 3
1190 expected = list(g(s))*3
1191 actual = list(islice(cycle(g(s)), tgtlen))
1192 self.assertEqual(actual, expected)
1193 self.assertRaises(TypeError, cycle, X(s))
1194 self.assertRaises(TypeError, list, cycle(N(s)))
1195 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1196
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001197 def test_groupby(self):
1198 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1199 for g in (G, I, Ig, S, L, R):
1200 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1201 self.assertRaises(TypeError, groupby, X(s))
1202 self.assertRaises(TypeError, list, groupby(N(s)))
1203 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1204
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001205 def test_ifilter(self):
1206 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1207 for g in (G, I, Ig, S, L, R):
1208 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
1209 self.assertRaises(TypeError, ifilter, isEven, X(s))
1210 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
1211 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
1212
1213 def test_ifilterfalse(self):
1214 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1215 for g in (G, I, Ig, S, L, R):
1216 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
1217 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
1218 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
1219 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
1220
1221 def test_izip(self):
1222 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1223 for g in (G, I, Ig, S, L, R):
1224 self.assertEqual(list(izip(g(s))), zip(g(s)))
1225 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
1226 self.assertRaises(TypeError, izip, X(s))
1227 self.assertRaises(TypeError, list, izip(N(s)))
1228 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
1229
Raymond Hettingerd36862c2007-02-21 05:20:38 +00001230 def test_iziplongest(self):
1231 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1232 for g in (G, I, Ig, S, L, R):
1233 self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
1234 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
1235 self.assertRaises(TypeError, izip_longest, X(s))
1236 self.assertRaises(TypeError, list, izip_longest(N(s)))
1237 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
1238
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001239 def test_imap(self):
1240 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1241 for g in (G, I, Ig, S, L, R):
1242 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
1243 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
1244 self.assertRaises(TypeError, imap, onearg, X(s))
1245 self.assertRaises(TypeError, list, imap(onearg, N(s)))
1246 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
1247
1248 def test_islice(self):
1249 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1250 for g in (G, I, Ig, S, L, R):
1251 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1252 self.assertRaises(TypeError, islice, X(s), 10)
1253 self.assertRaises(TypeError, list, islice(N(s), 10))
1254 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1255
1256 def test_starmap(self):
1257 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1258 for g in (G, I, Ig, S, L, R):
1259 ss = zip(s, s)
1260 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
1261 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
1262 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
1263 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1264
1265 def test_takewhile(self):
1266 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1267 for g in (G, I, Ig, S, L, R):
1268 tgt = []
1269 for elem in g(s):
1270 if not isEven(elem): break
1271 tgt.append(elem)
1272 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1273 self.assertRaises(TypeError, takewhile, isEven, X(s))
1274 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
1275 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1276
1277 def test_dropwhile(self):
1278 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1279 for g in (G, I, Ig, S, L, R):
1280 tgt = []
1281 for elem in g(s):
1282 if not tgt and isOdd(elem): continue
1283 tgt.append(elem)
1284 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1285 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
1286 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
1287 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1288
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001289 def test_tee(self):
1290 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1291 for g in (G, I, Ig, S, L, R):
1292 it1, it2 = tee(g(s))
1293 self.assertEqual(list(it1), list(g(s)))
1294 self.assertEqual(list(it2), list(g(s)))
1295 self.assertRaises(TypeError, tee, X(s))
1296 self.assertRaises(TypeError, list, tee(N(s))[0])
1297 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1298
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001299class LengthTransparency(unittest.TestCase):
1300
1301 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001302 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001303 self.assertEqual(len(repeat(None, 50)), 50)
1304 self.assertRaises(TypeError, len, repeat(None))
1305
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001306class RegressionTests(unittest.TestCase):
1307
1308 def test_sf_793826(self):
1309 # Fix Armin Rigo's successful efforts to wreak havoc
1310
1311 def mutatingtuple(tuple1, f, tuple2):
1312 # this builds a tuple t which is a copy of tuple1,
1313 # then calls f(t), then mutates t to be equal to tuple2
1314 # (needs len(tuple1) == len(tuple2)).
1315 def g(value, first=[1]):
1316 if first:
1317 del first[:]
1318 f(z.next())
1319 return value
1320 items = list(tuple2)
1321 items[1:1] = list(tuple1)
1322 gen = imap(g, items)
1323 z = izip(*[gen]*len(tuple1))
1324 z.next()
1325
1326 def f(t):
1327 global T
1328 T = t
1329 first[:] = list(T)
1330
1331 first = []
1332 mutatingtuple((1,2,3), f, (4,5,6))
1333 second = list(T)
1334 self.assertEqual(first, second)
1335
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001336
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001337 def test_sf_950057(self):
1338 # Make sure that chain() and cycle() catch exceptions immediately
1339 # rather than when shifting between input sources
1340
1341 def gen1():
1342 hist.append(0)
1343 yield 1
1344 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001345 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001346 hist.append(2)
1347
1348 def gen2(x):
1349 hist.append(3)
1350 yield 2
1351 hist.append(4)
1352 if x:
1353 raise StopIteration
1354
1355 hist = []
1356 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1357 self.assertEqual(hist, [0,1])
1358
1359 hist = []
1360 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1361 self.assertEqual(hist, [0,1])
1362
1363 hist = []
1364 self.assertRaises(AssertionError, list, cycle(gen1()))
1365 self.assertEqual(hist, [0,1])
1366
Georg Brandlb84c1372007-01-21 10:28:43 +00001367class SubclassWithKwargsTest(unittest.TestCase):
1368 def test_keywords_in_subclass(self):
1369 # count is not subclassable...
1370 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001371 starmap, islice, takewhile, dropwhile, cycle, compress):
Georg Brandlb84c1372007-01-21 10:28:43 +00001372 class Subclass(cls):
1373 def __init__(self, newarg=None, *args):
1374 cls.__init__(self, *args)
1375 try:
1376 Subclass(newarg=1)
1377 except TypeError, err:
1378 # we expect type errors because of wrong argument count
Benjamin Peterson5c8da862009-06-30 22:57:08 +00001379 self.assertFalse("does not take keyword arguments" in err.args[0])
Georg Brandlb84c1372007-01-21 10:28:43 +00001380
1381
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001382libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001383
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001384
1385>>> amounts = [120.15, 764.05, 823.14]
1386>>> for checknum, amount in izip(count(1200), amounts):
1387... print 'Check %d is for $%.2f' % (checknum, amount)
1388...
1389Check 1200 is for $120.15
1390Check 1201 is for $764.05
1391Check 1202 is for $823.14
1392
1393>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001394>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1395... print cube
1396...
13971
13988
139927
1400
1401>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001402>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001403... print name.title()
1404...
1405Alex
1406Laura
1407Martin
1408Walter
1409Samuele
1410
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001411>>> from operator import itemgetter
1412>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Armin Rigoa3f09272006-05-28 19:13:17 +00001413>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001414>>> for k, g in groupby(di, itemgetter(1)):
1415... print k, map(itemgetter(0), g)
1416...
14171 ['a', 'c', 'e']
14182 ['b', 'd', 'f']
14193 ['g']
1420
Raymond Hettinger734fb572004-01-20 20:04:40 +00001421# Find runs of consecutive numbers using groupby. The key to the solution
1422# is differencing with a range so that consecutive numbers all appear in
1423# same group.
1424>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
1425>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
1426... print map(operator.itemgetter(1), g)
1427...
1428[1]
1429[4, 5, 6]
1430[10]
1431[15, 16, 17, 18]
1432[22]
1433[25, 26, 27, 28]
1434
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001435>>> def take(n, iterable):
1436... "Return first n items of the iterable as a list"
1437... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001438
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001439>>> def enumerate(iterable, start=0):
1440... return izip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001441
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001442>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001443... "Return function(0), function(1), ..."
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001444... return imap(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001445
Raymond Hettingerf9bce832009-02-19 05:34:35 +00001446>>> def nth(iterable, n, default=None):
1447... "Returns the nth item or a default value"
1448... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001449
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001450>>> def quantify(iterable, pred=bool):
1451... "Count how many times the predicate is true"
1452... return sum(imap(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001453
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001454>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001455... "Returns the sequence elements and then returns None indefinitely"
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001456... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001457
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001458>>> def ncycles(iterable, n):
1459... "Returns the seqeuence elements n times"
1460... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001461
1462>>> def dotproduct(vec1, vec2):
1463... return sum(imap(operator.mul, vec1, vec2))
1464
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001465>>> def flatten(listOfLists):
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001466... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001467
1468>>> def repeatfunc(func, times=None, *args):
1469... "Repeat calls to func with specified arguments."
1470... " Example: repeatfunc(random.random)"
1471... if times is None:
1472... return starmap(func, repeat(args))
1473... else:
1474... return starmap(func, repeat(args, times))
1475
Raymond Hettingerd591f662003-10-26 15:34:50 +00001476>>> def pairwise(iterable):
1477... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1478... a, b = tee(iterable)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001479... for elem in b:
1480... break
Raymond Hettingerad983e72003-11-12 14:32:26 +00001481... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001482
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001483>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001484... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001485... args = [iter(iterable)] * n
Raymond Hettingerf080e6d2008-07-31 01:19:50 +00001486... return izip_longest(fillvalue=fillvalue, *args)
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001487
1488>>> def roundrobin(*iterables):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001489... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001490... # Recipe credited to George Sakkis
1491... pending = len(iterables)
1492... nexts = cycle(iter(it).next for it in iterables)
1493... while pending:
1494... try:
1495... for next in nexts:
1496... yield next()
1497... except StopIteration:
1498... pending -= 1
1499... nexts = cycle(islice(nexts, pending))
1500
1501>>> def powerset(iterable):
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001502... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1503... s = list(iterable)
1504... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001505
Raymond Hettinger44e15812009-01-02 21:26:45 +00001506>>> def unique_everseen(iterable, key=None):
1507... "List unique elements, preserving order. Remember all elements ever seen."
1508... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1509... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1510... seen = set()
1511... seen_add = seen.add
1512... if key is None:
1513... for element in iterable:
1514... if element not in seen:
1515... seen_add(element)
1516... yield element
1517... else:
1518... for element in iterable:
1519... k = key(element)
1520... if k not in seen:
1521... seen_add(k)
1522... yield element
1523
1524>>> def unique_justseen(iterable, key=None):
1525... "List unique elements, preserving order. Remember only the element just seen."
1526... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1527... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1528... return imap(next, imap(itemgetter(1), groupby(iterable, key)))
1529
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001530This is not part of the examples but it tests to make sure the definitions
1531perform as purported.
1532
Raymond Hettingera098b332003-09-08 23:58:40 +00001533>>> take(10, count())
1534[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1535
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001536>>> list(enumerate('abc'))
1537[(0, 'a'), (1, 'b'), (2, 'c')]
1538
1539>>> list(islice(tabulate(lambda x: 2*x), 4))
1540[0, 2, 4, 6]
1541
1542>>> nth('abcde', 3)
Raymond Hettingerd507afd2009-02-04 10:52:32 +00001543'd'
1544
1545>>> nth('abcde', 9) is None
1546True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001547
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001548>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000154950
1550
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001551>>> a = [[1, 2, 3], [4, 5, 6]]
1552>>> flatten(a)
1553[1, 2, 3, 4, 5, 6]
1554
1555>>> list(repeatfunc(pow, 5, 2, 3))
1556[8, 8, 8, 8, 8]
1557
1558>>> import random
1559>>> take(5, imap(int, repeatfunc(random.random)))
1560[0, 0, 0, 0, 0]
1561
Raymond Hettingerd591f662003-10-26 15:34:50 +00001562>>> list(pairwise('abcd'))
1563[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001564
Raymond Hettingerd591f662003-10-26 15:34:50 +00001565>>> list(pairwise([]))
1566[]
1567
1568>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001569[]
1570
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001571>>> list(islice(padnone('abc'), 0, 6))
1572['a', 'b', 'c', None, None, None]
1573
1574>>> list(ncycles('abc', 3))
1575['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1576
1577>>> dotproduct([1,2,3], [4,5,6])
157832
1579
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001580>>> list(grouper(3, 'abcdefg', 'x'))
1581[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1582
1583>>> list(roundrobin('abc', 'd', 'ef'))
1584['a', 'd', 'e', 'b', 'f', 'c']
1585
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001586>>> list(powerset([1,2,3]))
1587[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001588
Raymond Hettinger560f9a82009-01-27 13:26:35 +00001589>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1590True
1591
1592>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1593True
1594
Raymond Hettinger44e15812009-01-02 21:26:45 +00001595>>> list(unique_everseen('AAAABBBCCDAABBB'))
1596['A', 'B', 'C', 'D']
1597
1598>>> list(unique_everseen('ABBCcAD', str.lower))
1599['A', 'B', 'C', 'D']
1600
1601>>> list(unique_justseen('AAAABBBCCDAABBB'))
1602['A', 'B', 'C', 'D', 'A', 'B']
1603
1604>>> list(unique_justseen('ABBCcAD', str.lower))
1605['A', 'B', 'C', 'A', 'D']
1606
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001607"""
1608
1609__test__ = {'libreftest' : libreftest}
1610
1611def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001612 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Georg Brandlb84c1372007-01-21 10:28:43 +00001613 RegressionTests, LengthTransparency,
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001614 SubclassWithKwargsTest, TestExamples)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001615 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001616
1617 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001618 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001619 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001620 counts = [None] * 5
1621 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001622 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001623 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001624 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001625 print counts
1626
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001627 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001628 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001629
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001630if __name__ == "__main__":
1631 test_main(verbose=True)