blob: 9e6b7c82941033425f9fb331bbaaf0dd693646b8 [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 Hettinger2012f172003-02-07 05:32:58 +00005import sys
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00006import operator
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00007import random
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +00008maxsize = test_support.MAX_Py_ssize_t
9minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000010
11def onearg(x):
12 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000013 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000014
15def errfunc(*args):
16 'Test function that raises an error'
17 raise ValueError
18
19def gen3():
20 'Non-restartable source sequence'
21 for i in (0, 1, 2):
22 yield i
23
24def isEven(x):
25 'Test predicate'
26 return x%2==0
27
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000028def isOdd(x):
29 'Test predicate'
30 return x%2==1
31
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000032class StopNow:
33 'Class emulating an empty iterable.'
34 def __iter__(self):
35 return self
36 def next(self):
37 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000038
Raymond Hettinger02420702003-06-29 20:36:23 +000039def take(n, seq):
40 'Convenience function for partially consuming a long of infinite iterable'
41 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000042
Raymond Hettingerd553d852008-03-04 04:17:08 +000043def prod(iterable):
44 return reduce(operator.mul, iterable, 1)
45
Raymond Hettinger93e804d2008-02-26 23:40:50 +000046def fact(n):
47 'Factorial'
Raymond Hettingerd553d852008-03-04 04:17:08 +000048 return prod(range(1, n+1))
49
Raymond Hettinger96ef8112003-02-01 00:10:11 +000050class TestBasicOps(unittest.TestCase):
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000051 def test_chain(self):
Raymond Hettingerad47fa12008-03-06 20:52:01 +000052
53 def chain2(*iterables):
54 'Pure python version in the docs'
55 for it in iterables:
56 for element in it:
57 yield element
58
59 for c in (chain, chain2):
60 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
61 self.assertEqual(list(c('abc')), list('abc'))
62 self.assertEqual(list(c('')), [])
63 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
64 self.assertRaises(TypeError, list,c(2, 3))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000065
Raymond Hettingerb4cbc982008-02-28 22:46:41 +000066 def test_chain_from_iterable(self):
67 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
68 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
69 self.assertEqual(list(chain.from_iterable([''])), [])
70 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
71 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
72
Raymond Hettinger93e804d2008-02-26 23:40:50 +000073 def test_combinations(self):
Raymond Hettinger5b913e32009-01-08 06:39:04 +000074 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Raymond Hettinger93e804d2008-02-26 23:40:50 +000075 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Raymond Hettingerd553d852008-03-04 04:17:08 +000076 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Raymond Hettinger93e804d2008-02-26 23:40:50 +000077 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Raymond Hettinger5b913e32009-01-08 06:39:04 +000078 self.assertEqual(list(combinations('abc', 32)), []) # r > n
Raymond Hettinger93e804d2008-02-26 23:40:50 +000079 self.assertEqual(list(combinations(range(4), 3)),
80 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
Raymond Hettingerd553d852008-03-04 04:17:08 +000081
82 def combinations1(iterable, r):
83 'Pure python version shown in the docs'
84 pool = tuple(iterable)
85 n = len(pool)
Raymond Hettinger5b913e32009-01-08 06:39:04 +000086 if r > n:
87 return
Raymond Hettingerd553d852008-03-04 04:17:08 +000088 indices = range(r)
89 yield tuple(pool[i] for i in indices)
90 while 1:
91 for i in reversed(range(r)):
92 if indices[i] != i + n - r:
93 break
94 else:
95 return
96 indices[i] += 1
97 for j in range(i+1, r):
98 indices[j] = indices[j-1] + 1
99 yield tuple(pool[i] for i in indices)
100
101 def combinations2(iterable, r):
102 'Pure python version shown in the docs'
103 pool = tuple(iterable)
104 n = len(pool)
105 for indices in permutations(range(n), r):
106 if sorted(indices) == list(indices):
107 yield tuple(pool[i] for i in indices)
108
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000109 def combinations3(iterable, r):
110 'Pure python version from cwr()'
111 pool = tuple(iterable)
112 n = len(pool)
113 for indices in combinations_with_replacement(range(n), r):
114 if len(set(indices)) == r:
115 yield tuple(pool[i] for i in indices)
116
Raymond Hettingerd553d852008-03-04 04:17:08 +0000117 for n in range(7):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000118 values = [5*x-12 for x in range(n)]
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000119 for r in range(n+2):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000120 result = list(combinations(values, r))
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000121 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 +0000122 self.assertEqual(len(result), len(set(result))) # no repeats
123 self.assertEqual(result, sorted(result)) # lexicographic order
124 for c in result:
125 self.assertEqual(len(c), r) # r-length combinations
126 self.assertEqual(len(set(c)), r) # no duplicate elements
127 self.assertEqual(list(c), sorted(c)) # keep original ordering
128 self.assert_(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000129 self.assertEqual(list(c),
130 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Raymond Hettingerd553d852008-03-04 04:17:08 +0000131 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000132 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000133 self.assertEqual(result, list(combinations3(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000134
135 # Test implementation detail: tuple re-use
136 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
137 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
138
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000139 def test_combinations_with_replacement(self):
140 cwr = combinations_with_replacement
141 self.assertRaises(TypeError, cwr, 'abc') # missing r argument
142 self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
143 self.assertRaises(TypeError, cwr, None) # pool is not iterable
144 self.assertRaises(ValueError, cwr, 'abc', -2) # r is negative
145 self.assertEqual(list(cwr('ABC', 2)),
146 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
147
148 def cwr1(iterable, r):
149 'Pure python version shown in the docs'
150 # number items returned: (n+r-1)! / r! / (n-1)! when n>0
151 pool = tuple(iterable)
152 n = len(pool)
153 if not n and r:
154 return
155 indices = [0] * r
156 yield tuple(pool[i] for i in indices)
157 while 1:
158 for i in reversed(range(r)):
159 if indices[i] != n - 1:
160 break
161 else:
162 return
163 indices[i:] = [indices[i] + 1] * (r - i)
164 yield tuple(pool[i] for i in indices)
165
166 def cwr2(iterable, r):
167 'Pure python version shown in the docs'
168 pool = tuple(iterable)
169 n = len(pool)
170 for indices in product(range(n), repeat=r):
171 if sorted(indices) == list(indices):
172 yield tuple(pool[i] for i in indices)
173
174 def numcombs(n, r):
175 if not n:
176 return 0 if r else 1
177 return fact(n+r-1) / fact(r)/ fact(n-1)
178
179 for n in range(7):
180 values = [5*x-12 for x in range(n)]
181 for r in range(n+2):
182 result = list(cwr(values, r))
183
184 self.assertEqual(len(result), numcombs(n, r)) # right number of combs
185 self.assertEqual(len(result), len(set(result))) # no repeats
186 self.assertEqual(result, sorted(result)) # lexicographic order
187
188 regular_combs = list(combinations(values, r)) # compare to combs without replacement
189 if n == 0 or r <= 1:
190 self.assertEquals(result, regular_combs) # cases that should be identical
191 else:
192 self.assert_(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
193
194 for c in result:
195 self.assertEqual(len(c), r) # r-length combinations
196 noruns = [k for k,v in groupby(c)] # combo without consecutive repeats
197 self.assertEqual(len(noruns), len(set(noruns))) # no repeats other than consecutive
198 self.assertEqual(list(c), sorted(c)) # keep original ordering
199 self.assert_(all(e in values for e in c)) # elements taken from input iterable
200 self.assertEqual(noruns,
201 [e for e in values if e in c]) # comb is a subsequence of the input iterable
202 self.assertEqual(result, list(cwr1(values, r))) # matches first pure python version
203 self.assertEqual(result, list(cwr2(values, r))) # matches second pure python version
204
205 # Test implementation detail: tuple re-use
206 self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
207 self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
208
Raymond Hettingerd553d852008-03-04 04:17:08 +0000209 def test_permutations(self):
210 self.assertRaises(TypeError, permutations) # too few arguments
211 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000212 self.assertRaises(TypeError, permutations, None) # pool is not iterable
213 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000214 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000215 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Raymond Hettingerd553d852008-03-04 04:17:08 +0000216 self.assertEqual(list(permutations(range(3), 2)),
217 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
218
219 def permutations1(iterable, r=None):
220 'Pure python version shown in the docs'
221 pool = tuple(iterable)
222 n = len(pool)
223 r = n if r is None else r
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000224 if r > n:
225 return
Raymond Hettingerd553d852008-03-04 04:17:08 +0000226 indices = range(n)
Raymond Hettingere70bb8d2008-03-23 00:55:46 +0000227 cycles = range(n, n-r, -1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000228 yield tuple(pool[i] for i in indices[:r])
229 while n:
230 for i in reversed(range(r)):
231 cycles[i] -= 1
232 if cycles[i] == 0:
233 indices[i:] = indices[i+1:] + indices[i:i+1]
234 cycles[i] = n - i
235 else:
236 j = cycles[i]
237 indices[i], indices[-j] = indices[-j], indices[i]
238 yield tuple(pool[i] for i in indices[:r])
239 break
240 else:
241 return
242
243 def permutations2(iterable, r=None):
244 'Pure python version shown in the docs'
245 pool = tuple(iterable)
246 n = len(pool)
247 r = n if r is None else r
248 for indices in product(range(n), repeat=r):
249 if len(set(indices)) == r:
250 yield tuple(pool[i] for i in indices)
251
252 for n in range(7):
253 values = [5*x-12 for x in range(n)]
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000254 for r in range(n+2):
Raymond Hettingerd553d852008-03-04 04:17:08 +0000255 result = list(permutations(values, r))
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000256 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 +0000257 self.assertEqual(len(result), len(set(result))) # no repeats
258 self.assertEqual(result, sorted(result)) # lexicographic order
259 for p in result:
260 self.assertEqual(len(p), r) # r-length permutations
261 self.assertEqual(len(set(p)), r) # no duplicate elements
262 self.assert_(all(e in values for e in p)) # elements taken from input iterable
263 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000264 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000265 if r == n:
266 self.assertEqual(result, list(permutations(values, None))) # test r as None
267 self.assertEqual(result, list(permutations(values))) # test default r
268
269 # Test implementation detail: tuple re-use
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000270 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000271 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000272
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000273 def test_combinatorics(self):
274 # Test relationships between product(), permutations(),
275 # combinations() and combinations_with_replacement().
276
Raymond Hettinger2f6c2e02009-01-27 10:36:14 +0000277 for n in range(6):
278 s = 'ABCDEFG'[:n]
279 for r in range(8):
280 prod = list(product(s, repeat=r))
281 cwr = list(combinations_with_replacement(s, r))
282 perm = list(permutations(s, r))
283 comb = list(combinations(s, r))
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000284
Raymond Hettinger2f6c2e02009-01-27 10:36:14 +0000285 # Check size
286 self.assertEquals(len(prod), n**r)
287 self.assertEquals(len(cwr), (fact(n+r-1) / fact(r)/ fact(n-1)) if n else (not r))
288 self.assertEquals(len(perm), 0 if r>n else fact(n) / fact(n-r))
289 self.assertEquals(len(comb), 0 if r>n else fact(n) / fact(r) / fact(n-r))
290
291 # Check lexicographic order without repeated tuples
292 self.assertEquals(prod, sorted(set(prod)))
293 self.assertEquals(cwr, sorted(set(cwr)))
294 self.assertEquals(perm, sorted(set(perm)))
295 self.assertEquals(comb, sorted(set(comb)))
296
297 # Check interrelationships
298 self.assertEquals(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
299 self.assertEquals(perm, [t for t in prod if len(set(t))==r]) # perm: prods with no dups
300 self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
301 self.assertEqual(comb, [t for t in cwr if len(set(t))==r]) # comb: cwrs without dups
302 self.assertEqual(comb, filter(set(cwr).__contains__, perm)) # comb: perm that is a cwr
303 self.assertEqual(comb, filter(set(perm).__contains__, cwr)) # comb: cwr that is a perm
304 self.assertEqual(comb, sorted(set(cwr) & set(perm))) # comb: both a cwr and a perm
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000305
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000306 def test_compress(self):
307 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
308 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
309 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
310 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
311 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
312 n = 10000
313 data = chain.from_iterable(repeat(range(6), n))
314 selectors = chain.from_iterable(repeat((0, 1)))
315 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
316 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
317 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
318 self.assertRaises(TypeError, compress, range(6)) # too few args
319 self.assertRaises(TypeError, compress, range(6), None) # too many args
320
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000321 def test_count(self):
322 self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
323 self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000324 self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000325 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
326 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000327 self.assertRaises(TypeError, count, 2, 3, 4)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000328 self.assertRaises(TypeError, count, 'a')
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000329 self.assertEqual(list(islice(count(maxsize-5), 10)), range(maxsize-5, maxsize+5))
330 self.assertEqual(list(islice(count(-maxsize-5), 10)), range(-maxsize-5, -maxsize+5))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000331 c = count(3)
332 self.assertEqual(repr(c), 'count(3)')
333 c.next()
334 self.assertEqual(repr(c), 'count(4)')
Jack Diederich36234e82006-09-21 17:50:26 +0000335 c = count(-9)
336 self.assertEqual(repr(c), 'count(-9)')
337 c.next()
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000338 self.assertEqual(repr(count(10.25)), 'count(10.25)')
Jack Diederich36234e82006-09-21 17:50:26 +0000339 self.assertEqual(c.next(), -8)
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000340 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 +0000341 # Test repr (ignoring the L in longs)
342 r1 = repr(count(i)).replace('L', '')
343 r2 = 'count(%r)'.__mod__(i).replace('L', '')
344 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000345
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000346 def test_count_with_stride(self):
347 self.assertEqual(zip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
348 self.assertEqual(zip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
349 self.assertEqual(zip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
350 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
351 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
352 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
353 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
354 c = count(3, 5)
355 self.assertEqual(repr(c), 'count(3, 5)')
356 c.next()
357 self.assertEqual(repr(c), 'count(8, 5)')
358 c = count(-9, 0)
359 self.assertEqual(repr(c), 'count(-9, 0)')
360 c.next()
361 self.assertEqual(repr(c), 'count(-9, 0)')
362 c = count(-9, -3)
363 self.assertEqual(repr(c), 'count(-9, -3)')
364 c.next()
365 self.assertEqual(repr(c), 'count(-12, -3)')
366 self.assertEqual(repr(c), 'count(-12, -3)')
367 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
368 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
369 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
370 for i in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 10, sys.maxint-5, sys.maxint+5):
371 for j in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 1, 10, sys.maxint-5, sys.maxint+5):
372 # Test repr (ignoring the L in longs)
373 r1 = repr(count(i, j)).replace('L', '')
374 if j == 1:
375 r2 = ('count(%r)' % i).replace('L', '')
376 else:
377 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
378 self.assertEqual(r1, r2)
379
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000380 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000381 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000382 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000383 self.assertRaises(TypeError, cycle)
384 self.assertRaises(TypeError, cycle, 5)
385 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000386
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000387 def test_groupby(self):
388 # Check whether it accepts arguments correctly
389 self.assertEqual([], list(groupby([])))
390 self.assertEqual([], list(groupby([], key=id)))
391 self.assertRaises(TypeError, list, groupby('abc', []))
392 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000393 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000394
395 # Check normal input
396 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
397 (2,15,22), (3,16,23), (3,17,23)]
398 dup = []
399 for k, g in groupby(s, lambda r:r[0]):
400 for elem in g:
401 self.assertEqual(k, elem[0])
402 dup.append(elem)
403 self.assertEqual(s, dup)
404
405 # Check nested case
406 dup = []
407 for k, g in groupby(s, lambda r:r[0]):
408 for ik, ig in groupby(g, lambda r:r[2]):
409 for elem in ig:
410 self.assertEqual(k, elem[0])
411 self.assertEqual(ik, elem[2])
412 dup.append(elem)
413 self.assertEqual(s, dup)
414
415 # Check case where inner iterator is not used
416 keys = [k for k, g in groupby(s, lambda r:r[0])]
417 expectedkeys = set([r[0] for r in s])
418 self.assertEqual(set(keys), expectedkeys)
419 self.assertEqual(len(keys), len(expectedkeys))
420
421 # Exercise pipes and filters style
422 s = 'abracadabra'
423 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000424 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000425 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
426 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000427 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000428 self.assertEqual(r, ['a', 'b', 'r'])
429 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000430 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000431 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
432 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000433 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000434 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
435
436 # iter.next failure
437 class ExpectedError(Exception):
438 pass
439 def delayed_raise(n=0):
440 for i in range(n):
441 yield 'yo'
442 raise ExpectedError
443 def gulp(iterable, keyp=None, func=list):
444 return [func(g) for k, g in groupby(iterable, keyp)]
445
446 # iter.next failure on outer object
447 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
448 # iter.next failure on inner object
449 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
450
451 # __cmp__ failure
452 class DummyCmp:
453 def __cmp__(self, dst):
454 raise ExpectedError
455 s = [DummyCmp(), DummyCmp(), None]
456
457 # __cmp__ failure on outer object
458 self.assertRaises(ExpectedError, gulp, s, func=id)
459 # __cmp__ failure on inner object
460 self.assertRaises(ExpectedError, gulp, s)
461
462 # keyfunc failure
463 def keyfunc(obj):
464 if keyfunc.skip > 0:
465 keyfunc.skip -= 1
466 return obj
467 else:
468 raise ExpectedError
469
470 # keyfunc failure on outer object
471 keyfunc.skip = 0
472 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
473 keyfunc.skip = 1
474 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
475
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000476 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000477 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000478 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000479 self.assertEqual(list(ifilter(bool, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000480 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000481 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000482 self.assertRaises(TypeError, ifilter, lambda x:x)
483 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000484 self.assertRaises(TypeError, ifilter, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000485 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000486
487 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000488 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
489 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000490 self.assertEqual(list(ifilterfalse(bool, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000491 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000492 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000493 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
494 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000495 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000496 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000497
498 def test_izip(self):
499 ans = [(x,y) for x, y in izip('abc',count())]
500 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000501 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
502 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
Raymond Hettinger02420702003-06-29 20:36:23 +0000503 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000504 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000505 self.assertEqual(list(izip()), zip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000506 self.assertRaises(TypeError, izip, 3)
507 self.assertRaises(TypeError, izip, range(3), 3)
508 # Check tuple re-use (implementation detail)
509 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
510 zip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000511 self.assertEqual([pair for pair in izip('abc', 'def')],
512 zip('abc', 'def'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000513 ids = map(id, izip('abc', 'def'))
514 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000515 ids = map(id, list(izip('abc', 'def')))
516 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000517
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000518 def test_iziplongest(self):
519 for args in [
520 ['abc', range(6)],
521 [range(6), 'abc'],
522 [range(1000), range(2000,2100), range(3000,3050)],
523 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
524 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
525 ]:
526 target = map(None, *args)
527 self.assertEqual(list(izip_longest(*args)), target)
528 self.assertEqual(list(izip_longest(*args, **{})), target)
529 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
530 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Tim Petersea5962f2007-03-12 18:07:52 +0000531
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000532 self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
533
534 self.assertEqual(list(izip_longest()), zip())
535 self.assertEqual(list(izip_longest([])), zip([]))
536 self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
Tim Petersea5962f2007-03-12 18:07:52 +0000537
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000538 self.assertEqual(list(izip_longest('abc', 'defg', **{})), map(None, 'abc', 'defg')) # empty keyword dict
539 self.assertRaises(TypeError, izip_longest, 3)
540 self.assertRaises(TypeError, izip_longest, range(3), 3)
541
542 for stmt in [
543 "izip_longest('abc', fv=1)",
Tim Petersea5962f2007-03-12 18:07:52 +0000544 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000545 ]:
546 try:
547 eval(stmt, globals(), locals())
548 except TypeError:
549 pass
550 else:
551 self.fail('Did not raise Type in: ' + stmt)
Tim Petersea5962f2007-03-12 18:07:52 +0000552
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000553 # Check tuple re-use (implementation detail)
554 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
555 zip('abc', 'def'))
556 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
557 zip('abc', 'def'))
558 ids = map(id, izip_longest('abc', 'def'))
559 self.assertEqual(min(ids), max(ids))
560 ids = map(id, list(izip_longest('abc', 'def')))
561 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
562
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000563 def test_product(self):
564 for args, result in [
Raymond Hettingerd553d852008-03-04 04:17:08 +0000565 ([], [()]), # zero iterables
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000566 (['ab'], [('a',), ('b',)]), # one iterable
567 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
568 ([range(0), range(2), range(3)], []), # first iterable with zero length
569 ([range(2), range(0), range(3)], []), # middle iterable with zero length
570 ([range(2), range(3), range(0)], []), # last iterable with zero length
571 ]:
572 self.assertEqual(list(product(*args)), result)
Raymond Hettinger08ff6822008-02-29 02:21:48 +0000573 for r in range(4):
574 self.assertEqual(list(product(*(args*r))),
575 list(product(*args, **dict(repeat=r))))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000576 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
577 self.assertRaises(TypeError, product, range(6), None)
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000578
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000579 def product1(*args, **kwds):
580 pools = map(tuple, args) * kwds.get('repeat', 1)
581 n = len(pools)
582 if n == 0:
583 yield ()
584 return
585 if any(len(pool) == 0 for pool in pools):
586 return
587 indices = [0] * n
588 yield tuple(pool[i] for pool, i in zip(pools, indices))
589 while 1:
590 for i in reversed(range(n)): # right to left
591 if indices[i] == len(pools[i]) - 1:
592 continue
593 indices[i] += 1
594 for j in range(i+1, n):
595 indices[j] = 0
596 yield tuple(pool[i] for pool, i in zip(pools, indices))
597 break
598 else:
599 return
600
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000601 def product2(*args, **kwds):
602 'Pure python version used in docs'
603 pools = map(tuple, args) * kwds.get('repeat', 1)
604 result = [[]]
605 for pool in pools:
606 result = [x+[y] for x in result for y in pool]
607 for prod in result:
608 yield tuple(prod)
609
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000610 argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
611 set('abcdefg'), range(11), tuple(range(13))]
612 for i in range(100):
613 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Raymond Hettingerd553d852008-03-04 04:17:08 +0000614 expected_len = prod(map(len, args))
615 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000616 self.assertEqual(list(product(*args)), list(product1(*args)))
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000617 self.assertEqual(list(product(*args)), list(product2(*args)))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000618 args = map(iter, args)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000619 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000620
Raymond Hettinger73d79632008-02-23 02:20:41 +0000621 # Test implementation detail: tuple re-use
622 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
623 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000624
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000625 def test_repeat(self):
626 self.assertEqual(zip(xrange(3),repeat('a')),
627 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000628 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000629 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000630 self.assertEqual(list(repeat('a', 0)), [])
631 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000632 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000633 self.assertRaises(TypeError, repeat, None, 3, 4)
634 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000635 r = repeat(1+0j)
636 self.assertEqual(repr(r), 'repeat((1+0j))')
637 r = repeat(1+0j, 5)
638 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
639 list(r)
640 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000641
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000642 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000643 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
644 [0**1, 1**2, 2**3])
645 self.assertEqual(list(imap(None, 'abc', range(5))),
646 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000647 self.assertEqual(list(imap(None, 'abc', count())),
648 [('a',0),('b',1),('c',2)])
649 self.assertEqual(take(2,imap(None, 'abc', count())),
650 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000651 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000652 self.assertRaises(TypeError, imap)
653 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000654 self.assertRaises(TypeError, imap(10, range(5)).next)
655 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
656 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000657
658 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000659 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
660 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000661 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
662 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000663 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger47317092008-01-17 03:02:14 +0000664 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
665 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000666 self.assertRaises(TypeError, starmap)
667 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
668 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
669 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
670 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000671
672 def test_islice(self):
673 for args in [ # islice(args) should agree with range(args)
674 (10, 20, 3),
675 (10, 3, 20),
676 (10, 20),
677 (10, 3),
678 (20,)
679 ]:
680 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
681
682 for args, tgtargs in [ # Stop when seqn is exhausted
683 ((10, 110, 3), ((10, 100, 3))),
684 ((10, 110), ((10, 100))),
685 ((110,), (100,))
686 ]:
687 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
688
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000689 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000690 self.assertEqual(list(islice(xrange(10), None)), range(10))
Raymond Hettingerb2594052004-12-05 09:25:51 +0000691 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
692 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000693 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
694 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
695
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000696 # Test number of items consumed SF #1171417
697 it = iter(range(10))
698 self.assertEqual(list(islice(it, 3)), range(3))
699 self.assertEqual(list(it), range(3, 10))
700
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000701 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000702 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000703 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
704 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
705 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
706 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
707 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000708 self.assertRaises(ValueError, islice, xrange(10), 'a')
709 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
710 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
711 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
712 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +0000713 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000714
715 def test_takewhile(self):
716 data = [1, 3, 5, 20, 2, 4, 6, 8]
717 underten = lambda x: x<10
718 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000719 self.assertEqual(list(takewhile(underten, [])), [])
720 self.assertRaises(TypeError, takewhile)
721 self.assertRaises(TypeError, takewhile, operator.pow)
722 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
723 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
724 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000725 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
726 self.assertEqual(list(t), [1, 1, 1])
727 self.assertRaises(StopIteration, t.next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000728
729 def test_dropwhile(self):
730 data = [1, 3, 5, 20, 2, 4, 6, 8]
731 underten = lambda x: x<10
732 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000733 self.assertEqual(list(dropwhile(underten, [])), [])
734 self.assertRaises(TypeError, dropwhile)
735 self.assertRaises(TypeError, dropwhile, operator.pow)
736 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
737 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
738 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000739
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000740 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000741 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000742 def irange(n):
743 for i in xrange(n):
744 yield i
745
746 a, b = tee([]) # test empty iterator
747 self.assertEqual(list(a), [])
748 self.assertEqual(list(b), [])
749
750 a, b = tee(irange(n)) # test 100% interleaved
751 self.assertEqual(zip(a,b), zip(range(n),range(n)))
752
753 a, b = tee(irange(n)) # test 0% interleaved
754 self.assertEqual(list(a), range(n))
755 self.assertEqual(list(b), range(n))
756
757 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000758 for i in xrange(100):
759 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000760 del a
761 self.assertEqual(list(b), range(n))
762
763 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000764 for i in xrange(100):
765 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000766 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000767 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000768
769 for j in xrange(5): # test randomly interleaved
770 order = [0]*n + [1]*n
771 random.shuffle(order)
772 lists = ([], [])
773 its = tee(irange(n))
774 for i in order:
775 value = its[i].next()
776 lists[i].append(value)
777 self.assertEqual(lists[0], range(n))
778 self.assertEqual(lists[1], range(n))
779
Raymond Hettingerad983e72003-11-12 14:32:26 +0000780 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000781 self.assertRaises(TypeError, tee)
782 self.assertRaises(TypeError, tee, 3)
783 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000784 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000785
Raymond Hettingerad983e72003-11-12 14:32:26 +0000786 # tee object should be instantiable
787 a, b = tee('abc')
788 c = type(a)('def')
789 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000790
Raymond Hettingerad983e72003-11-12 14:32:26 +0000791 # test long-lagged and multi-way split
792 a, b, c = tee(xrange(2000), 3)
793 for i in xrange(100):
794 self.assertEqual(a.next(), i)
795 self.assertEqual(list(b), range(2000))
796 self.assertEqual([c.next(), c.next()], range(2))
797 self.assertEqual(list(a), range(100,2000))
798 self.assertEqual(list(c), range(2,2000))
799
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000800 # test values of n
801 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Neal Norwitz69e88972006-09-02 02:58:13 +0000802 self.assertRaises(ValueError, tee, [], -1)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000803 for n in xrange(5):
804 result = tee('abc', n)
805 self.assertEqual(type(result), tuple)
806 self.assertEqual(len(result), n)
807 self.assertEqual(map(list, result), [list('abc')]*n)
808
Raymond Hettingerad983e72003-11-12 14:32:26 +0000809 # tee pass-through to copyable iterator
810 a, b = tee('abc')
811 c, d = tee(a)
812 self.assert_(a is c)
813
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000814 # test tee_new
815 t1, t2 = tee('abc')
816 tnew = type(t1)
817 self.assertRaises(TypeError, tnew)
818 self.assertRaises(TypeError, tnew, 10)
819 t3 = tnew(t1)
820 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000821
Raymond Hettingera9f60922004-10-17 16:40:14 +0000822 # test that tee objects are weak referencable
823 a, b = tee(xrange(10))
824 p = proxy(a)
825 self.assertEqual(getattr(p, '__class__'), type(b))
826 del a
827 self.assertRaises(ReferenceError, getattr, p, '__class__')
828
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000829 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000830 self.assertRaises(StopIteration, izip().next)
831
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000832 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000833 self.assertRaises(StopIteration, f([]).next)
834 self.assertRaises(StopIteration, f(StopNow()).next)
835
836 self.assertRaises(StopIteration, islice([], None).next)
837 self.assertRaises(StopIteration, islice(StopNow(), None).next)
838
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000839 p, q = tee([])
840 self.assertRaises(StopIteration, p.next)
841 self.assertRaises(StopIteration, q.next)
842 p, q = tee(StopNow())
843 self.assertRaises(StopIteration, p.next)
844 self.assertRaises(StopIteration, q.next)
845
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000846 self.assertRaises(StopIteration, repeat(None, 0).next)
847
848 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
849 self.assertRaises(StopIteration, f(lambda x:x, []).next)
850 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
851
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000852class TestExamples(unittest.TestCase):
853
854 def test_chain(self):
855 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
856
857 def test_chain_from_iterable(self):
858 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
859
860 def test_combinations(self):
861 self.assertEqual(list(combinations('ABCD', 2)),
862 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
863 self.assertEqual(list(combinations(range(4), 3)),
864 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
865
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000866 def test_combinations_with_replacement(self):
867 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
868 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
869
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000870 def test_compress(self):
871 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
872
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000873 def test_count(self):
874 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
875
876 def test_cycle(self):
877 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
878
879 def test_dropwhile(self):
880 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
881
882 def test_groupby(self):
883 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
884 list('ABCDAB'))
885 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
886 [list('AAAA'), list('BBB'), list('CC'), list('D')])
887
888 def test_ifilter(self):
889 self.assertEqual(list(ifilter(lambda x: x%2, range(10))), [1,3,5,7,9])
890
891 def test_ifilterfalse(self):
892 self.assertEqual(list(ifilterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
893
894 def test_imap(self):
895 self.assertEqual(list(imap(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
896
897 def test_islice(self):
898 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
899 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
900 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
901 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
902
903 def test_izip(self):
904 self.assertEqual(list(izip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
905
906 def test_izip_longest(self):
907 self.assertEqual(list(izip_longest('ABCD', 'xy', fillvalue='-')),
908 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
909
910 def test_permutations(self):
911 self.assertEqual(list(permutations('ABCD', 2)),
912 map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split()))
913 self.assertEqual(list(permutations(range(3))),
914 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
915
916 def test_product(self):
917 self.assertEqual(list(product('ABCD', 'xy')),
918 map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split()))
919 self.assertEqual(list(product(range(2), repeat=3)),
920 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
921 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
922
923 def test_repeat(self):
924 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
925
926 def test_stapmap(self):
927 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
928 [32, 9, 1000])
929
930 def test_takewhile(self):
931 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
932
933
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000934class TestGC(unittest.TestCase):
935
936 def makecycle(self, iterator, container):
937 container.append(iterator)
938 iterator.next()
939 del container, iterator
940
941 def test_chain(self):
942 a = []
943 self.makecycle(chain(a), a)
944
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000945 def test_chain_from_iterable(self):
946 a = []
947 self.makecycle(chain.from_iterable([a]), a)
948
949 def test_combinations(self):
950 a = []
951 self.makecycle(combinations([1,2,a,3], 3), a)
952
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000953 def test_combinations_with_replacement(self):
954 a = []
955 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
956
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000957 def test_compress(self):
958 a = []
959 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
960
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000961 def test_cycle(self):
962 a = []
963 self.makecycle(cycle([a]*2), a)
964
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000965 def test_dropwhile(self):
966 a = []
967 self.makecycle(dropwhile(bool, [0, a, a]), a)
968
969 def test_groupby(self):
970 a = []
971 self.makecycle(groupby([a]*2, lambda x:x), a)
972
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000973 def test_issue2246(self):
974 # Issue 2246 -- the _grouper iterator was not included in GC
975 n = 10
976 keyfunc = lambda x: x
977 for i, j in groupby(xrange(n), key=keyfunc):
978 keyfunc.__dict__.setdefault('x',[]).append(j)
979
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000980 def test_ifilter(self):
981 a = []
982 self.makecycle(ifilter(lambda x:True, [a]*2), a)
983
984 def test_ifilterfalse(self):
985 a = []
986 self.makecycle(ifilterfalse(lambda x:False, a), a)
987
988 def test_izip(self):
989 a = []
990 self.makecycle(izip([a]*2, [a]*3), a)
991
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000992 def test_izip_longest(self):
993 a = []
994 self.makecycle(izip_longest([a]*2, [a]*3), a)
995 b = [a, None]
996 self.makecycle(izip_longest([a]*2, [a]*3, fillvalue=b), a)
997
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000998 def test_imap(self):
999 a = []
1000 self.makecycle(imap(lambda x:x, [a]*2), a)
1001
1002 def test_islice(self):
1003 a = []
1004 self.makecycle(islice([a]*2, None), a)
1005
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001006 def test_permutations(self):
1007 a = []
1008 self.makecycle(permutations([1,2,a,3], 3), a)
1009
1010 def test_product(self):
1011 a = []
1012 self.makecycle(product([1,2,a,3], repeat=3), a)
1013
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001014 def test_repeat(self):
1015 a = []
1016 self.makecycle(repeat(a), a)
1017
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001018 def test_starmap(self):
1019 a = []
1020 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1021
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001022 def test_takewhile(self):
1023 a = []
1024 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1025
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001026def R(seqn):
1027 'Regular generator'
1028 for i in seqn:
1029 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001030
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001031class G:
1032 'Sequence using __getitem__'
1033 def __init__(self, seqn):
1034 self.seqn = seqn
1035 def __getitem__(self, i):
1036 return self.seqn[i]
1037
1038class I:
1039 'Sequence using iterator protocol'
1040 def __init__(self, seqn):
1041 self.seqn = seqn
1042 self.i = 0
1043 def __iter__(self):
1044 return self
1045 def next(self):
1046 if self.i >= len(self.seqn): raise StopIteration
1047 v = self.seqn[self.i]
1048 self.i += 1
1049 return v
1050
1051class Ig:
1052 'Sequence using iterator protocol defined with a generator'
1053 def __init__(self, seqn):
1054 self.seqn = seqn
1055 self.i = 0
1056 def __iter__(self):
1057 for val in self.seqn:
1058 yield val
1059
1060class X:
1061 'Missing __getitem__ and __iter__'
1062 def __init__(self, seqn):
1063 self.seqn = seqn
1064 self.i = 0
1065 def next(self):
1066 if self.i >= len(self.seqn): raise StopIteration
1067 v = self.seqn[self.i]
1068 self.i += 1
1069 return v
1070
1071class N:
1072 'Iterator missing next()'
1073 def __init__(self, seqn):
1074 self.seqn = seqn
1075 self.i = 0
1076 def __iter__(self):
1077 return self
1078
1079class E:
1080 'Test propagation of exceptions'
1081 def __init__(self, seqn):
1082 self.seqn = seqn
1083 self.i = 0
1084 def __iter__(self):
1085 return self
1086 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001087 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001088
1089class S:
1090 'Test immediate stop'
1091 def __init__(self, seqn):
1092 pass
1093 def __iter__(self):
1094 return self
1095 def next(self):
1096 raise StopIteration
1097
1098def L(seqn):
1099 'Test multiple tiers of iterators'
1100 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
1101
1102
1103class TestVariousIteratorArgs(unittest.TestCase):
1104
1105 def test_chain(self):
1106 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1107 for g in (G, I, Ig, S, L, R):
1108 self.assertEqual(list(chain(g(s))), list(g(s)))
1109 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001110 self.assertRaises(TypeError, list, chain(X(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001111 self.assertRaises(TypeError, list, chain(N(s)))
1112 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1113
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001114 def test_compress(self):
1115 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1116 n = len(s)
1117 for g in (G, I, Ig, S, L, R):
1118 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1119 self.assertRaises(TypeError, compress, X(s), repeat(1))
1120 self.assertRaises(TypeError, list, compress(N(s), repeat(1)))
1121 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1122
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001123 def test_product(self):
1124 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1125 self.assertRaises(TypeError, product, X(s))
1126 self.assertRaises(TypeError, product, N(s))
1127 self.assertRaises(ZeroDivisionError, product, E(s))
1128
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001129 def test_cycle(self):
1130 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1131 for g in (G, I, Ig, S, L, R):
1132 tgtlen = len(s) * 3
1133 expected = list(g(s))*3
1134 actual = list(islice(cycle(g(s)), tgtlen))
1135 self.assertEqual(actual, expected)
1136 self.assertRaises(TypeError, cycle, X(s))
1137 self.assertRaises(TypeError, list, cycle(N(s)))
1138 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1139
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001140 def test_groupby(self):
1141 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1142 for g in (G, I, Ig, S, L, R):
1143 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1144 self.assertRaises(TypeError, groupby, X(s))
1145 self.assertRaises(TypeError, list, groupby(N(s)))
1146 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1147
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001148 def test_ifilter(self):
1149 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1150 for g in (G, I, Ig, S, L, R):
1151 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
1152 self.assertRaises(TypeError, ifilter, isEven, X(s))
1153 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
1154 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
1155
1156 def test_ifilterfalse(self):
1157 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1158 for g in (G, I, Ig, S, L, R):
1159 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
1160 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
1161 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
1162 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
1163
1164 def test_izip(self):
1165 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1166 for g in (G, I, Ig, S, L, R):
1167 self.assertEqual(list(izip(g(s))), zip(g(s)))
1168 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
1169 self.assertRaises(TypeError, izip, X(s))
1170 self.assertRaises(TypeError, list, izip(N(s)))
1171 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
1172
Raymond Hettingerd36862c2007-02-21 05:20:38 +00001173 def test_iziplongest(self):
1174 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1175 for g in (G, I, Ig, S, L, R):
1176 self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
1177 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
1178 self.assertRaises(TypeError, izip_longest, X(s))
1179 self.assertRaises(TypeError, list, izip_longest(N(s)))
1180 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
1181
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001182 def test_imap(self):
1183 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1184 for g in (G, I, Ig, S, L, R):
1185 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
1186 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
1187 self.assertRaises(TypeError, imap, onearg, X(s))
1188 self.assertRaises(TypeError, list, imap(onearg, N(s)))
1189 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
1190
1191 def test_islice(self):
1192 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1193 for g in (G, I, Ig, S, L, R):
1194 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1195 self.assertRaises(TypeError, islice, X(s), 10)
1196 self.assertRaises(TypeError, list, islice(N(s), 10))
1197 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1198
1199 def test_starmap(self):
1200 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1201 for g in (G, I, Ig, S, L, R):
1202 ss = zip(s, s)
1203 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
1204 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
1205 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
1206 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1207
1208 def test_takewhile(self):
1209 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1210 for g in (G, I, Ig, S, L, R):
1211 tgt = []
1212 for elem in g(s):
1213 if not isEven(elem): break
1214 tgt.append(elem)
1215 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1216 self.assertRaises(TypeError, takewhile, isEven, X(s))
1217 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
1218 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1219
1220 def test_dropwhile(self):
1221 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1222 for g in (G, I, Ig, S, L, R):
1223 tgt = []
1224 for elem in g(s):
1225 if not tgt and isOdd(elem): continue
1226 tgt.append(elem)
1227 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1228 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
1229 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
1230 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1231
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001232 def test_tee(self):
1233 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1234 for g in (G, I, Ig, S, L, R):
1235 it1, it2 = tee(g(s))
1236 self.assertEqual(list(it1), list(g(s)))
1237 self.assertEqual(list(it2), list(g(s)))
1238 self.assertRaises(TypeError, tee, X(s))
1239 self.assertRaises(TypeError, list, tee(N(s))[0])
1240 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1241
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001242class LengthTransparency(unittest.TestCase):
1243
1244 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001245 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001246 self.assertEqual(len(repeat(None, 50)), 50)
1247 self.assertRaises(TypeError, len, repeat(None))
1248
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001249class RegressionTests(unittest.TestCase):
1250
1251 def test_sf_793826(self):
1252 # Fix Armin Rigo's successful efforts to wreak havoc
1253
1254 def mutatingtuple(tuple1, f, tuple2):
1255 # this builds a tuple t which is a copy of tuple1,
1256 # then calls f(t), then mutates t to be equal to tuple2
1257 # (needs len(tuple1) == len(tuple2)).
1258 def g(value, first=[1]):
1259 if first:
1260 del first[:]
1261 f(z.next())
1262 return value
1263 items = list(tuple2)
1264 items[1:1] = list(tuple1)
1265 gen = imap(g, items)
1266 z = izip(*[gen]*len(tuple1))
1267 z.next()
1268
1269 def f(t):
1270 global T
1271 T = t
1272 first[:] = list(T)
1273
1274 first = []
1275 mutatingtuple((1,2,3), f, (4,5,6))
1276 second = list(T)
1277 self.assertEqual(first, second)
1278
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001279
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001280 def test_sf_950057(self):
1281 # Make sure that chain() and cycle() catch exceptions immediately
1282 # rather than when shifting between input sources
1283
1284 def gen1():
1285 hist.append(0)
1286 yield 1
1287 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001288 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001289 hist.append(2)
1290
1291 def gen2(x):
1292 hist.append(3)
1293 yield 2
1294 hist.append(4)
1295 if x:
1296 raise StopIteration
1297
1298 hist = []
1299 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1300 self.assertEqual(hist, [0,1])
1301
1302 hist = []
1303 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1304 self.assertEqual(hist, [0,1])
1305
1306 hist = []
1307 self.assertRaises(AssertionError, list, cycle(gen1()))
1308 self.assertEqual(hist, [0,1])
1309
Georg Brandlb84c1372007-01-21 10:28:43 +00001310class SubclassWithKwargsTest(unittest.TestCase):
1311 def test_keywords_in_subclass(self):
1312 # count is not subclassable...
1313 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001314 starmap, islice, takewhile, dropwhile, cycle, compress):
Georg Brandlb84c1372007-01-21 10:28:43 +00001315 class Subclass(cls):
1316 def __init__(self, newarg=None, *args):
1317 cls.__init__(self, *args)
1318 try:
1319 Subclass(newarg=1)
1320 except TypeError, err:
1321 # we expect type errors because of wrong argument count
1322 self.failIf("does not take keyword arguments" in err.args[0])
1323
1324
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001325libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001326
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001327
1328>>> amounts = [120.15, 764.05, 823.14]
1329>>> for checknum, amount in izip(count(1200), amounts):
1330... print 'Check %d is for $%.2f' % (checknum, amount)
1331...
1332Check 1200 is for $120.15
1333Check 1201 is for $764.05
1334Check 1202 is for $823.14
1335
1336>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001337>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1338... print cube
1339...
13401
13418
134227
1343
1344>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001345>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001346... print name.title()
1347...
1348Alex
1349Laura
1350Martin
1351Walter
1352Samuele
1353
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001354>>> from operator import itemgetter
1355>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Armin Rigoa3f09272006-05-28 19:13:17 +00001356>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001357>>> for k, g in groupby(di, itemgetter(1)):
1358... print k, map(itemgetter(0), g)
1359...
13601 ['a', 'c', 'e']
13612 ['b', 'd', 'f']
13623 ['g']
1363
Raymond Hettinger734fb572004-01-20 20:04:40 +00001364# Find runs of consecutive numbers using groupby. The key to the solution
1365# is differencing with a range so that consecutive numbers all appear in
1366# same group.
1367>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
1368>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
1369... print map(operator.itemgetter(1), g)
1370...
1371[1]
1372[4, 5, 6]
1373[10]
1374[15, 16, 17, 18]
1375[22]
1376[25, 26, 27, 28]
1377
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001378>>> def take(n, iterable):
1379... "Return first n items of the iterable as a list"
1380... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001381
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001382>>> def enumerate(iterable, start=0):
1383... return izip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001384
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001385>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001386... "Return function(0), function(1), ..."
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001387... return imap(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001388
1389>>> def nth(iterable, n):
Raymond Hettingerd507afd2009-02-04 10:52:32 +00001390... "Returns the nth item or None"
1391... return next(islice(iterable, n, None), None)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001392
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001393>>> def quantify(iterable, pred=bool):
1394... "Count how many times the predicate is true"
1395... return sum(imap(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001396
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001397>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001398... "Returns the sequence elements and then returns None indefinitely"
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001399... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001400
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001401>>> def ncycles(iterable, n):
1402... "Returns the seqeuence elements n times"
1403... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001404
1405>>> def dotproduct(vec1, vec2):
1406... return sum(imap(operator.mul, vec1, vec2))
1407
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001408>>> def flatten(listOfLists):
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001409... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001410
1411>>> def repeatfunc(func, times=None, *args):
1412... "Repeat calls to func with specified arguments."
1413... " Example: repeatfunc(random.random)"
1414... if times is None:
1415... return starmap(func, repeat(args))
1416... else:
1417... return starmap(func, repeat(args, times))
1418
Raymond Hettingerd591f662003-10-26 15:34:50 +00001419>>> def pairwise(iterable):
1420... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1421... a, b = tee(iterable)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001422... for elem in b:
1423... break
Raymond Hettingerad983e72003-11-12 14:32:26 +00001424... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001425
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001426>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001427... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001428... args = [iter(iterable)] * n
Raymond Hettingerf080e6d2008-07-31 01:19:50 +00001429... return izip_longest(fillvalue=fillvalue, *args)
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001430
1431>>> def roundrobin(*iterables):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001432... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001433... # Recipe credited to George Sakkis
1434... pending = len(iterables)
1435... nexts = cycle(iter(it).next for it in iterables)
1436... while pending:
1437... try:
1438... for next in nexts:
1439... yield next()
1440... except StopIteration:
1441... pending -= 1
1442... nexts = cycle(islice(nexts, pending))
1443
1444>>> def powerset(iterable):
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001445... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1446... s = list(iterable)
1447... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001448
Raymond Hettinger44e15812009-01-02 21:26:45 +00001449>>> def unique_everseen(iterable, key=None):
1450... "List unique elements, preserving order. Remember all elements ever seen."
1451... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1452... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1453... seen = set()
1454... seen_add = seen.add
1455... if key is None:
1456... for element in iterable:
1457... if element not in seen:
1458... seen_add(element)
1459... yield element
1460... else:
1461... for element in iterable:
1462... k = key(element)
1463... if k not in seen:
1464... seen_add(k)
1465... yield element
1466
1467>>> def unique_justseen(iterable, key=None):
1468... "List unique elements, preserving order. Remember only the element just seen."
1469... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1470... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1471... return imap(next, imap(itemgetter(1), groupby(iterable, key)))
1472
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001473This is not part of the examples but it tests to make sure the definitions
1474perform as purported.
1475
Raymond Hettingera098b332003-09-08 23:58:40 +00001476>>> take(10, count())
1477[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1478
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001479>>> list(enumerate('abc'))
1480[(0, 'a'), (1, 'b'), (2, 'c')]
1481
1482>>> list(islice(tabulate(lambda x: 2*x), 4))
1483[0, 2, 4, 6]
1484
1485>>> nth('abcde', 3)
Raymond Hettingerd507afd2009-02-04 10:52:32 +00001486'd'
1487
1488>>> nth('abcde', 9) is None
1489True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001490
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001491>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000149250
1493
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001494>>> a = [[1, 2, 3], [4, 5, 6]]
1495>>> flatten(a)
1496[1, 2, 3, 4, 5, 6]
1497
1498>>> list(repeatfunc(pow, 5, 2, 3))
1499[8, 8, 8, 8, 8]
1500
1501>>> import random
1502>>> take(5, imap(int, repeatfunc(random.random)))
1503[0, 0, 0, 0, 0]
1504
Raymond Hettingerd591f662003-10-26 15:34:50 +00001505>>> list(pairwise('abcd'))
1506[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001507
Raymond Hettingerd591f662003-10-26 15:34:50 +00001508>>> list(pairwise([]))
1509[]
1510
1511>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001512[]
1513
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001514>>> list(islice(padnone('abc'), 0, 6))
1515['a', 'b', 'c', None, None, None]
1516
1517>>> list(ncycles('abc', 3))
1518['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1519
1520>>> dotproduct([1,2,3], [4,5,6])
152132
1522
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001523>>> list(grouper(3, 'abcdefg', 'x'))
1524[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1525
1526>>> list(roundrobin('abc', 'd', 'ef'))
1527['a', 'd', 'e', 'b', 'f', 'c']
1528
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001529>>> list(powerset([1,2,3]))
1530[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001531
Raymond Hettinger560f9a82009-01-27 13:26:35 +00001532>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1533True
1534
1535>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1536True
1537
Raymond Hettinger44e15812009-01-02 21:26:45 +00001538>>> list(unique_everseen('AAAABBBCCDAABBB'))
1539['A', 'B', 'C', 'D']
1540
1541>>> list(unique_everseen('ABBCcAD', str.lower))
1542['A', 'B', 'C', 'D']
1543
1544>>> list(unique_justseen('AAAABBBCCDAABBB'))
1545['A', 'B', 'C', 'D', 'A', 'B']
1546
1547>>> list(unique_justseen('ABBCcAD', str.lower))
1548['A', 'B', 'C', 'A', 'D']
1549
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001550"""
1551
1552__test__ = {'libreftest' : libreftest}
1553
1554def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001555 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Georg Brandlb84c1372007-01-21 10:28:43 +00001556 RegressionTests, LengthTransparency,
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001557 SubclassWithKwargsTest, TestExamples)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001558 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001559
1560 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001561 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001562 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001563 counts = [None] * 5
1564 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001565 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001566 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001567 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001568 print counts
1569
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001570 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001571 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001572
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001573if __name__ == "__main__":
1574 test_main(verbose=True)