blob: b79f70f2f4b494fbd663fe538aece58b38baab35 [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 Hettinger96ef8112003-02-01 00:10:11 +0000327 self.assertRaises(TypeError, count, 2, 3)
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()
338 self.assertEqual(c.next(), -8)
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000339 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 +0000340 # Test repr (ignoring the L in longs)
341 r1 = repr(count(i)).replace('L', '')
342 r2 = 'count(%r)'.__mod__(i).replace('L', '')
343 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000344
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000345 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000346 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000347 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000348 self.assertRaises(TypeError, cycle)
349 self.assertRaises(TypeError, cycle, 5)
350 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000351
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000352 def test_groupby(self):
353 # Check whether it accepts arguments correctly
354 self.assertEqual([], list(groupby([])))
355 self.assertEqual([], list(groupby([], key=id)))
356 self.assertRaises(TypeError, list, groupby('abc', []))
357 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000358 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000359
360 # Check normal input
361 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
362 (2,15,22), (3,16,23), (3,17,23)]
363 dup = []
364 for k, g in groupby(s, lambda r:r[0]):
365 for elem in g:
366 self.assertEqual(k, elem[0])
367 dup.append(elem)
368 self.assertEqual(s, dup)
369
370 # Check nested case
371 dup = []
372 for k, g in groupby(s, lambda r:r[0]):
373 for ik, ig in groupby(g, lambda r:r[2]):
374 for elem in ig:
375 self.assertEqual(k, elem[0])
376 self.assertEqual(ik, elem[2])
377 dup.append(elem)
378 self.assertEqual(s, dup)
379
380 # Check case where inner iterator is not used
381 keys = [k for k, g in groupby(s, lambda r:r[0])]
382 expectedkeys = set([r[0] for r in s])
383 self.assertEqual(set(keys), expectedkeys)
384 self.assertEqual(len(keys), len(expectedkeys))
385
386 # Exercise pipes and filters style
387 s = 'abracadabra'
388 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000389 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000390 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
391 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000392 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000393 self.assertEqual(r, ['a', 'b', 'r'])
394 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000395 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000396 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
397 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000398 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000399 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
400
401 # iter.next failure
402 class ExpectedError(Exception):
403 pass
404 def delayed_raise(n=0):
405 for i in range(n):
406 yield 'yo'
407 raise ExpectedError
408 def gulp(iterable, keyp=None, func=list):
409 return [func(g) for k, g in groupby(iterable, keyp)]
410
411 # iter.next failure on outer object
412 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
413 # iter.next failure on inner object
414 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
415
416 # __cmp__ failure
417 class DummyCmp:
418 def __cmp__(self, dst):
419 raise ExpectedError
420 s = [DummyCmp(), DummyCmp(), None]
421
422 # __cmp__ failure on outer object
423 self.assertRaises(ExpectedError, gulp, s, func=id)
424 # __cmp__ failure on inner object
425 self.assertRaises(ExpectedError, gulp, s)
426
427 # keyfunc failure
428 def keyfunc(obj):
429 if keyfunc.skip > 0:
430 keyfunc.skip -= 1
431 return obj
432 else:
433 raise ExpectedError
434
435 # keyfunc failure on outer object
436 keyfunc.skip = 0
437 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
438 keyfunc.skip = 1
439 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
440
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000441 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000442 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000443 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000444 self.assertEqual(list(ifilter(bool, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000445 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000446 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000447 self.assertRaises(TypeError, ifilter, lambda x:x)
448 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000449 self.assertRaises(TypeError, ifilter, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000450 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000451
452 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000453 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
454 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000455 self.assertEqual(list(ifilterfalse(bool, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000456 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000457 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000458 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
459 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000460 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000461 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000462
463 def test_izip(self):
464 ans = [(x,y) for x, y in izip('abc',count())]
465 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000466 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
467 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
Raymond Hettinger02420702003-06-29 20:36:23 +0000468 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000469 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000470 self.assertEqual(list(izip()), zip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000471 self.assertRaises(TypeError, izip, 3)
472 self.assertRaises(TypeError, izip, range(3), 3)
473 # Check tuple re-use (implementation detail)
474 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
475 zip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000476 self.assertEqual([pair for pair in izip('abc', 'def')],
477 zip('abc', 'def'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000478 ids = map(id, izip('abc', 'def'))
479 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000480 ids = map(id, list(izip('abc', 'def')))
481 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000482
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000483 def test_iziplongest(self):
484 for args in [
485 ['abc', range(6)],
486 [range(6), 'abc'],
487 [range(1000), range(2000,2100), range(3000,3050)],
488 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
489 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
490 ]:
491 target = map(None, *args)
492 self.assertEqual(list(izip_longest(*args)), target)
493 self.assertEqual(list(izip_longest(*args, **{})), target)
494 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
495 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Tim Petersea5962f2007-03-12 18:07:52 +0000496
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000497 self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
498
499 self.assertEqual(list(izip_longest()), zip())
500 self.assertEqual(list(izip_longest([])), zip([]))
501 self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
Tim Petersea5962f2007-03-12 18:07:52 +0000502
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000503 self.assertEqual(list(izip_longest('abc', 'defg', **{})), map(None, 'abc', 'defg')) # empty keyword dict
504 self.assertRaises(TypeError, izip_longest, 3)
505 self.assertRaises(TypeError, izip_longest, range(3), 3)
506
507 for stmt in [
508 "izip_longest('abc', fv=1)",
Tim Petersea5962f2007-03-12 18:07:52 +0000509 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000510 ]:
511 try:
512 eval(stmt, globals(), locals())
513 except TypeError:
514 pass
515 else:
516 self.fail('Did not raise Type in: ' + stmt)
Tim Petersea5962f2007-03-12 18:07:52 +0000517
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000518 # Check tuple re-use (implementation detail)
519 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
520 zip('abc', 'def'))
521 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
522 zip('abc', 'def'))
523 ids = map(id, izip_longest('abc', 'def'))
524 self.assertEqual(min(ids), max(ids))
525 ids = map(id, list(izip_longest('abc', 'def')))
526 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
527
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000528 def test_product(self):
529 for args, result in [
Raymond Hettingerd553d852008-03-04 04:17:08 +0000530 ([], [()]), # zero iterables
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000531 (['ab'], [('a',), ('b',)]), # one iterable
532 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
533 ([range(0), range(2), range(3)], []), # first iterable with zero length
534 ([range(2), range(0), range(3)], []), # middle iterable with zero length
535 ([range(2), range(3), range(0)], []), # last iterable with zero length
536 ]:
537 self.assertEqual(list(product(*args)), result)
Raymond Hettinger08ff6822008-02-29 02:21:48 +0000538 for r in range(4):
539 self.assertEqual(list(product(*(args*r))),
540 list(product(*args, **dict(repeat=r))))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000541 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
542 self.assertRaises(TypeError, product, range(6), None)
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000543
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000544 def product1(*args, **kwds):
545 pools = map(tuple, args) * kwds.get('repeat', 1)
546 n = len(pools)
547 if n == 0:
548 yield ()
549 return
550 if any(len(pool) == 0 for pool in pools):
551 return
552 indices = [0] * n
553 yield tuple(pool[i] for pool, i in zip(pools, indices))
554 while 1:
555 for i in reversed(range(n)): # right to left
556 if indices[i] == len(pools[i]) - 1:
557 continue
558 indices[i] += 1
559 for j in range(i+1, n):
560 indices[j] = 0
561 yield tuple(pool[i] for pool, i in zip(pools, indices))
562 break
563 else:
564 return
565
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000566 def product2(*args, **kwds):
567 'Pure python version used in docs'
568 pools = map(tuple, args) * kwds.get('repeat', 1)
569 result = [[]]
570 for pool in pools:
571 result = [x+[y] for x in result for y in pool]
572 for prod in result:
573 yield tuple(prod)
574
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000575 argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
576 set('abcdefg'), range(11), tuple(range(13))]
577 for i in range(100):
578 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Raymond Hettingerd553d852008-03-04 04:17:08 +0000579 expected_len = prod(map(len, args))
580 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000581 self.assertEqual(list(product(*args)), list(product1(*args)))
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000582 self.assertEqual(list(product(*args)), list(product2(*args)))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000583 args = map(iter, args)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000584 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000585
Raymond Hettinger73d79632008-02-23 02:20:41 +0000586 # Test implementation detail: tuple re-use
587 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
588 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000589
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000590 def test_repeat(self):
591 self.assertEqual(zip(xrange(3),repeat('a')),
592 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000593 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000594 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000595 self.assertEqual(list(repeat('a', 0)), [])
596 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000597 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000598 self.assertRaises(TypeError, repeat, None, 3, 4)
599 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000600 r = repeat(1+0j)
601 self.assertEqual(repr(r), 'repeat((1+0j))')
602 r = repeat(1+0j, 5)
603 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
604 list(r)
605 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000606
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000607 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000608 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
609 [0**1, 1**2, 2**3])
610 self.assertEqual(list(imap(None, 'abc', range(5))),
611 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000612 self.assertEqual(list(imap(None, 'abc', count())),
613 [('a',0),('b',1),('c',2)])
614 self.assertEqual(take(2,imap(None, 'abc', count())),
615 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000616 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000617 self.assertRaises(TypeError, imap)
618 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000619 self.assertRaises(TypeError, imap(10, range(5)).next)
620 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
621 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000622
623 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000624 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
625 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000626 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
627 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000628 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger47317092008-01-17 03:02:14 +0000629 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
630 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000631 self.assertRaises(TypeError, starmap)
632 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
633 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
634 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
635 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000636
637 def test_islice(self):
638 for args in [ # islice(args) should agree with range(args)
639 (10, 20, 3),
640 (10, 3, 20),
641 (10, 20),
642 (10, 3),
643 (20,)
644 ]:
645 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
646
647 for args, tgtargs in [ # Stop when seqn is exhausted
648 ((10, 110, 3), ((10, 100, 3))),
649 ((10, 110), ((10, 100))),
650 ((110,), (100,))
651 ]:
652 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
653
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000654 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000655 self.assertEqual(list(islice(xrange(10), None)), range(10))
Raymond Hettingerb2594052004-12-05 09:25:51 +0000656 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
657 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000658 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
659 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
660
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000661 # Test number of items consumed SF #1171417
662 it = iter(range(10))
663 self.assertEqual(list(islice(it, 3)), range(3))
664 self.assertEqual(list(it), range(3, 10))
665
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000666 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000667 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000668 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
669 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
670 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
671 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
672 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000673 self.assertRaises(ValueError, islice, xrange(10), 'a')
674 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
675 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
676 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
677 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +0000678 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000679
680 def test_takewhile(self):
681 data = [1, 3, 5, 20, 2, 4, 6, 8]
682 underten = lambda x: x<10
683 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000684 self.assertEqual(list(takewhile(underten, [])), [])
685 self.assertRaises(TypeError, takewhile)
686 self.assertRaises(TypeError, takewhile, operator.pow)
687 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
688 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
689 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000690 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
691 self.assertEqual(list(t), [1, 1, 1])
692 self.assertRaises(StopIteration, t.next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000693
694 def test_dropwhile(self):
695 data = [1, 3, 5, 20, 2, 4, 6, 8]
696 underten = lambda x: x<10
697 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000698 self.assertEqual(list(dropwhile(underten, [])), [])
699 self.assertRaises(TypeError, dropwhile)
700 self.assertRaises(TypeError, dropwhile, operator.pow)
701 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
702 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
703 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000704
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000705 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000706 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000707 def irange(n):
708 for i in xrange(n):
709 yield i
710
711 a, b = tee([]) # test empty iterator
712 self.assertEqual(list(a), [])
713 self.assertEqual(list(b), [])
714
715 a, b = tee(irange(n)) # test 100% interleaved
716 self.assertEqual(zip(a,b), zip(range(n),range(n)))
717
718 a, b = tee(irange(n)) # test 0% interleaved
719 self.assertEqual(list(a), range(n))
720 self.assertEqual(list(b), range(n))
721
722 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000723 for i in xrange(100):
724 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000725 del a
726 self.assertEqual(list(b), range(n))
727
728 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000729 for i in xrange(100):
730 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000731 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000732 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000733
734 for j in xrange(5): # test randomly interleaved
735 order = [0]*n + [1]*n
736 random.shuffle(order)
737 lists = ([], [])
738 its = tee(irange(n))
739 for i in order:
740 value = its[i].next()
741 lists[i].append(value)
742 self.assertEqual(lists[0], range(n))
743 self.assertEqual(lists[1], range(n))
744
Raymond Hettingerad983e72003-11-12 14:32:26 +0000745 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000746 self.assertRaises(TypeError, tee)
747 self.assertRaises(TypeError, tee, 3)
748 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000749 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000750
Raymond Hettingerad983e72003-11-12 14:32:26 +0000751 # tee object should be instantiable
752 a, b = tee('abc')
753 c = type(a)('def')
754 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000755
Raymond Hettingerad983e72003-11-12 14:32:26 +0000756 # test long-lagged and multi-way split
757 a, b, c = tee(xrange(2000), 3)
758 for i in xrange(100):
759 self.assertEqual(a.next(), i)
760 self.assertEqual(list(b), range(2000))
761 self.assertEqual([c.next(), c.next()], range(2))
762 self.assertEqual(list(a), range(100,2000))
763 self.assertEqual(list(c), range(2,2000))
764
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000765 # test values of n
766 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Neal Norwitz69e88972006-09-02 02:58:13 +0000767 self.assertRaises(ValueError, tee, [], -1)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000768 for n in xrange(5):
769 result = tee('abc', n)
770 self.assertEqual(type(result), tuple)
771 self.assertEqual(len(result), n)
772 self.assertEqual(map(list, result), [list('abc')]*n)
773
Raymond Hettingerad983e72003-11-12 14:32:26 +0000774 # tee pass-through to copyable iterator
775 a, b = tee('abc')
776 c, d = tee(a)
777 self.assert_(a is c)
778
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000779 # test tee_new
780 t1, t2 = tee('abc')
781 tnew = type(t1)
782 self.assertRaises(TypeError, tnew)
783 self.assertRaises(TypeError, tnew, 10)
784 t3 = tnew(t1)
785 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000786
Raymond Hettingera9f60922004-10-17 16:40:14 +0000787 # test that tee objects are weak referencable
788 a, b = tee(xrange(10))
789 p = proxy(a)
790 self.assertEqual(getattr(p, '__class__'), type(b))
791 del a
792 self.assertRaises(ReferenceError, getattr, p, '__class__')
793
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000794 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000795 self.assertRaises(StopIteration, izip().next)
796
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000797 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000798 self.assertRaises(StopIteration, f([]).next)
799 self.assertRaises(StopIteration, f(StopNow()).next)
800
801 self.assertRaises(StopIteration, islice([], None).next)
802 self.assertRaises(StopIteration, islice(StopNow(), None).next)
803
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000804 p, q = tee([])
805 self.assertRaises(StopIteration, p.next)
806 self.assertRaises(StopIteration, q.next)
807 p, q = tee(StopNow())
808 self.assertRaises(StopIteration, p.next)
809 self.assertRaises(StopIteration, q.next)
810
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000811 self.assertRaises(StopIteration, repeat(None, 0).next)
812
813 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
814 self.assertRaises(StopIteration, f(lambda x:x, []).next)
815 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
816
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000817class TestExamples(unittest.TestCase):
818
819 def test_chain(self):
820 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
821
822 def test_chain_from_iterable(self):
823 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
824
825 def test_combinations(self):
826 self.assertEqual(list(combinations('ABCD', 2)),
827 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
828 self.assertEqual(list(combinations(range(4), 3)),
829 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
830
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000831 def test_combinations_with_replacement(self):
832 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
833 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
834
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000835 def test_compress(self):
836 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
837
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000838 def test_count(self):
839 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
840
841 def test_cycle(self):
842 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
843
844 def test_dropwhile(self):
845 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
846
847 def test_groupby(self):
848 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
849 list('ABCDAB'))
850 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
851 [list('AAAA'), list('BBB'), list('CC'), list('D')])
852
853 def test_ifilter(self):
854 self.assertEqual(list(ifilter(lambda x: x%2, range(10))), [1,3,5,7,9])
855
856 def test_ifilterfalse(self):
857 self.assertEqual(list(ifilterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
858
859 def test_imap(self):
860 self.assertEqual(list(imap(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
861
862 def test_islice(self):
863 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
864 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
865 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
866 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
867
868 def test_izip(self):
869 self.assertEqual(list(izip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
870
871 def test_izip_longest(self):
872 self.assertEqual(list(izip_longest('ABCD', 'xy', fillvalue='-')),
873 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
874
875 def test_permutations(self):
876 self.assertEqual(list(permutations('ABCD', 2)),
877 map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split()))
878 self.assertEqual(list(permutations(range(3))),
879 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
880
881 def test_product(self):
882 self.assertEqual(list(product('ABCD', 'xy')),
883 map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split()))
884 self.assertEqual(list(product(range(2), repeat=3)),
885 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
886 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
887
888 def test_repeat(self):
889 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
890
891 def test_stapmap(self):
892 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
893 [32, 9, 1000])
894
895 def test_takewhile(self):
896 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
897
898
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000899class TestGC(unittest.TestCase):
900
901 def makecycle(self, iterator, container):
902 container.append(iterator)
903 iterator.next()
904 del container, iterator
905
906 def test_chain(self):
907 a = []
908 self.makecycle(chain(a), a)
909
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000910 def test_chain_from_iterable(self):
911 a = []
912 self.makecycle(chain.from_iterable([a]), a)
913
914 def test_combinations(self):
915 a = []
916 self.makecycle(combinations([1,2,a,3], 3), a)
917
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000918 def test_combinations_with_replacement(self):
919 a = []
920 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
921
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000922 def test_compress(self):
923 a = []
924 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
925
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000926 def test_cycle(self):
927 a = []
928 self.makecycle(cycle([a]*2), a)
929
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000930 def test_dropwhile(self):
931 a = []
932 self.makecycle(dropwhile(bool, [0, a, a]), a)
933
934 def test_groupby(self):
935 a = []
936 self.makecycle(groupby([a]*2, lambda x:x), a)
937
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000938 def test_issue2246(self):
939 # Issue 2246 -- the _grouper iterator was not included in GC
940 n = 10
941 keyfunc = lambda x: x
942 for i, j in groupby(xrange(n), key=keyfunc):
943 keyfunc.__dict__.setdefault('x',[]).append(j)
944
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000945 def test_ifilter(self):
946 a = []
947 self.makecycle(ifilter(lambda x:True, [a]*2), a)
948
949 def test_ifilterfalse(self):
950 a = []
951 self.makecycle(ifilterfalse(lambda x:False, a), a)
952
953 def test_izip(self):
954 a = []
955 self.makecycle(izip([a]*2, [a]*3), a)
956
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000957 def test_izip_longest(self):
958 a = []
959 self.makecycle(izip_longest([a]*2, [a]*3), a)
960 b = [a, None]
961 self.makecycle(izip_longest([a]*2, [a]*3, fillvalue=b), a)
962
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000963 def test_imap(self):
964 a = []
965 self.makecycle(imap(lambda x:x, [a]*2), a)
966
967 def test_islice(self):
968 a = []
969 self.makecycle(islice([a]*2, None), a)
970
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000971 def test_permutations(self):
972 a = []
973 self.makecycle(permutations([1,2,a,3], 3), a)
974
975 def test_product(self):
976 a = []
977 self.makecycle(product([1,2,a,3], repeat=3), a)
978
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000979 def test_repeat(self):
980 a = []
981 self.makecycle(repeat(a), a)
982
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000983 def test_starmap(self):
984 a = []
985 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
986
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000987 def test_takewhile(self):
988 a = []
989 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
990
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000991def R(seqn):
992 'Regular generator'
993 for i in seqn:
994 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000995
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000996class G:
997 'Sequence using __getitem__'
998 def __init__(self, seqn):
999 self.seqn = seqn
1000 def __getitem__(self, i):
1001 return self.seqn[i]
1002
1003class I:
1004 'Sequence using iterator protocol'
1005 def __init__(self, seqn):
1006 self.seqn = seqn
1007 self.i = 0
1008 def __iter__(self):
1009 return self
1010 def next(self):
1011 if self.i >= len(self.seqn): raise StopIteration
1012 v = self.seqn[self.i]
1013 self.i += 1
1014 return v
1015
1016class Ig:
1017 'Sequence using iterator protocol defined with a generator'
1018 def __init__(self, seqn):
1019 self.seqn = seqn
1020 self.i = 0
1021 def __iter__(self):
1022 for val in self.seqn:
1023 yield val
1024
1025class X:
1026 'Missing __getitem__ and __iter__'
1027 def __init__(self, seqn):
1028 self.seqn = seqn
1029 self.i = 0
1030 def next(self):
1031 if self.i >= len(self.seqn): raise StopIteration
1032 v = self.seqn[self.i]
1033 self.i += 1
1034 return v
1035
1036class N:
1037 'Iterator missing next()'
1038 def __init__(self, seqn):
1039 self.seqn = seqn
1040 self.i = 0
1041 def __iter__(self):
1042 return self
1043
1044class E:
1045 'Test propagation of exceptions'
1046 def __init__(self, seqn):
1047 self.seqn = seqn
1048 self.i = 0
1049 def __iter__(self):
1050 return self
1051 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001052 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001053
1054class S:
1055 'Test immediate stop'
1056 def __init__(self, seqn):
1057 pass
1058 def __iter__(self):
1059 return self
1060 def next(self):
1061 raise StopIteration
1062
1063def L(seqn):
1064 'Test multiple tiers of iterators'
1065 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
1066
1067
1068class TestVariousIteratorArgs(unittest.TestCase):
1069
1070 def test_chain(self):
1071 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1072 for g in (G, I, Ig, S, L, R):
1073 self.assertEqual(list(chain(g(s))), list(g(s)))
1074 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001075 self.assertRaises(TypeError, list, chain(X(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001076 self.assertRaises(TypeError, list, chain(N(s)))
1077 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1078
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001079 def test_compress(self):
1080 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1081 n = len(s)
1082 for g in (G, I, Ig, S, L, R):
1083 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1084 self.assertRaises(TypeError, compress, X(s), repeat(1))
1085 self.assertRaises(TypeError, list, compress(N(s), repeat(1)))
1086 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1087
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001088 def test_product(self):
1089 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1090 self.assertRaises(TypeError, product, X(s))
1091 self.assertRaises(TypeError, product, N(s))
1092 self.assertRaises(ZeroDivisionError, product, E(s))
1093
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001094 def test_cycle(self):
1095 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1096 for g in (G, I, Ig, S, L, R):
1097 tgtlen = len(s) * 3
1098 expected = list(g(s))*3
1099 actual = list(islice(cycle(g(s)), tgtlen))
1100 self.assertEqual(actual, expected)
1101 self.assertRaises(TypeError, cycle, X(s))
1102 self.assertRaises(TypeError, list, cycle(N(s)))
1103 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1104
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001105 def test_groupby(self):
1106 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1107 for g in (G, I, Ig, S, L, R):
1108 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1109 self.assertRaises(TypeError, groupby, X(s))
1110 self.assertRaises(TypeError, list, groupby(N(s)))
1111 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1112
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001113 def test_ifilter(self):
1114 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1115 for g in (G, I, Ig, S, L, R):
1116 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
1117 self.assertRaises(TypeError, ifilter, isEven, X(s))
1118 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
1119 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
1120
1121 def test_ifilterfalse(self):
1122 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1123 for g in (G, I, Ig, S, L, R):
1124 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
1125 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
1126 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
1127 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
1128
1129 def test_izip(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 self.assertEqual(list(izip(g(s))), zip(g(s)))
1133 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
1134 self.assertRaises(TypeError, izip, X(s))
1135 self.assertRaises(TypeError, list, izip(N(s)))
1136 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
1137
Raymond Hettingerd36862c2007-02-21 05:20:38 +00001138 def test_iziplongest(self):
1139 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1140 for g in (G, I, Ig, S, L, R):
1141 self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
1142 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
1143 self.assertRaises(TypeError, izip_longest, X(s))
1144 self.assertRaises(TypeError, list, izip_longest(N(s)))
1145 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
1146
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001147 def test_imap(self):
1148 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1149 for g in (G, I, Ig, S, L, R):
1150 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
1151 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
1152 self.assertRaises(TypeError, imap, onearg, X(s))
1153 self.assertRaises(TypeError, list, imap(onearg, N(s)))
1154 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
1155
1156 def test_islice(self):
1157 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1158 for g in (G, I, Ig, S, L, R):
1159 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1160 self.assertRaises(TypeError, islice, X(s), 10)
1161 self.assertRaises(TypeError, list, islice(N(s), 10))
1162 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1163
1164 def test_starmap(self):
1165 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1166 for g in (G, I, Ig, S, L, R):
1167 ss = zip(s, s)
1168 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
1169 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
1170 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
1171 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1172
1173 def test_takewhile(self):
1174 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1175 for g in (G, I, Ig, S, L, R):
1176 tgt = []
1177 for elem in g(s):
1178 if not isEven(elem): break
1179 tgt.append(elem)
1180 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1181 self.assertRaises(TypeError, takewhile, isEven, X(s))
1182 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
1183 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1184
1185 def test_dropwhile(self):
1186 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1187 for g in (G, I, Ig, S, L, R):
1188 tgt = []
1189 for elem in g(s):
1190 if not tgt and isOdd(elem): continue
1191 tgt.append(elem)
1192 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1193 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
1194 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
1195 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1196
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001197 def test_tee(self):
1198 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1199 for g in (G, I, Ig, S, L, R):
1200 it1, it2 = tee(g(s))
1201 self.assertEqual(list(it1), list(g(s)))
1202 self.assertEqual(list(it2), list(g(s)))
1203 self.assertRaises(TypeError, tee, X(s))
1204 self.assertRaises(TypeError, list, tee(N(s))[0])
1205 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1206
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001207class LengthTransparency(unittest.TestCase):
1208
1209 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001210 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001211 self.assertEqual(len(repeat(None, 50)), 50)
1212 self.assertRaises(TypeError, len, repeat(None))
1213
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001214class RegressionTests(unittest.TestCase):
1215
1216 def test_sf_793826(self):
1217 # Fix Armin Rigo's successful efforts to wreak havoc
1218
1219 def mutatingtuple(tuple1, f, tuple2):
1220 # this builds a tuple t which is a copy of tuple1,
1221 # then calls f(t), then mutates t to be equal to tuple2
1222 # (needs len(tuple1) == len(tuple2)).
1223 def g(value, first=[1]):
1224 if first:
1225 del first[:]
1226 f(z.next())
1227 return value
1228 items = list(tuple2)
1229 items[1:1] = list(tuple1)
1230 gen = imap(g, items)
1231 z = izip(*[gen]*len(tuple1))
1232 z.next()
1233
1234 def f(t):
1235 global T
1236 T = t
1237 first[:] = list(T)
1238
1239 first = []
1240 mutatingtuple((1,2,3), f, (4,5,6))
1241 second = list(T)
1242 self.assertEqual(first, second)
1243
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001244
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001245 def test_sf_950057(self):
1246 # Make sure that chain() and cycle() catch exceptions immediately
1247 # rather than when shifting between input sources
1248
1249 def gen1():
1250 hist.append(0)
1251 yield 1
1252 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001253 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001254 hist.append(2)
1255
1256 def gen2(x):
1257 hist.append(3)
1258 yield 2
1259 hist.append(4)
1260 if x:
1261 raise StopIteration
1262
1263 hist = []
1264 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1265 self.assertEqual(hist, [0,1])
1266
1267 hist = []
1268 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1269 self.assertEqual(hist, [0,1])
1270
1271 hist = []
1272 self.assertRaises(AssertionError, list, cycle(gen1()))
1273 self.assertEqual(hist, [0,1])
1274
Georg Brandlb84c1372007-01-21 10:28:43 +00001275class SubclassWithKwargsTest(unittest.TestCase):
1276 def test_keywords_in_subclass(self):
1277 # count is not subclassable...
1278 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001279 starmap, islice, takewhile, dropwhile, cycle, compress):
Georg Brandlb84c1372007-01-21 10:28:43 +00001280 class Subclass(cls):
1281 def __init__(self, newarg=None, *args):
1282 cls.__init__(self, *args)
1283 try:
1284 Subclass(newarg=1)
1285 except TypeError, err:
1286 # we expect type errors because of wrong argument count
1287 self.failIf("does not take keyword arguments" in err.args[0])
1288
1289
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001290libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001291
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001292
1293>>> amounts = [120.15, 764.05, 823.14]
1294>>> for checknum, amount in izip(count(1200), amounts):
1295... print 'Check %d is for $%.2f' % (checknum, amount)
1296...
1297Check 1200 is for $120.15
1298Check 1201 is for $764.05
1299Check 1202 is for $823.14
1300
1301>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001302>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1303... print cube
1304...
13051
13068
130727
1308
1309>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001310>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001311... print name.title()
1312...
1313Alex
1314Laura
1315Martin
1316Walter
1317Samuele
1318
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001319>>> from operator import itemgetter
1320>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Armin Rigoa3f09272006-05-28 19:13:17 +00001321>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001322>>> for k, g in groupby(di, itemgetter(1)):
1323... print k, map(itemgetter(0), g)
1324...
13251 ['a', 'c', 'e']
13262 ['b', 'd', 'f']
13273 ['g']
1328
Raymond Hettinger734fb572004-01-20 20:04:40 +00001329# Find runs of consecutive numbers using groupby. The key to the solution
1330# is differencing with a range so that consecutive numbers all appear in
1331# same group.
1332>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
1333>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
1334... print map(operator.itemgetter(1), g)
1335...
1336[1]
1337[4, 5, 6]
1338[10]
1339[15, 16, 17, 18]
1340[22]
1341[25, 26, 27, 28]
1342
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001343>>> def take(n, iterable):
1344... "Return first n items of the iterable as a list"
1345... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001346
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001347>>> def enumerate(iterable, start=0):
1348... return izip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001349
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001350>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001351... "Return function(0), function(1), ..."
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001352... return imap(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001353
1354>>> def nth(iterable, n):
Raymond Hettingerd507afd2009-02-04 10:52:32 +00001355... "Returns the nth item or None"
1356... return next(islice(iterable, n, None), None)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001357
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001358>>> def quantify(iterable, pred=bool):
1359... "Count how many times the predicate is true"
1360... return sum(imap(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001361
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001362>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001363... "Returns the sequence elements and then returns None indefinitely"
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001364... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001365
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001366>>> def ncycles(iterable, n):
1367... "Returns the seqeuence elements n times"
1368... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001369
1370>>> def dotproduct(vec1, vec2):
1371... return sum(imap(operator.mul, vec1, vec2))
1372
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001373>>> def flatten(listOfLists):
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001374... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001375
1376>>> def repeatfunc(func, times=None, *args):
1377... "Repeat calls to func with specified arguments."
1378... " Example: repeatfunc(random.random)"
1379... if times is None:
1380... return starmap(func, repeat(args))
1381... else:
1382... return starmap(func, repeat(args, times))
1383
Raymond Hettingerd591f662003-10-26 15:34:50 +00001384>>> def pairwise(iterable):
1385... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1386... a, b = tee(iterable)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001387... for elem in b:
1388... break
Raymond Hettingerad983e72003-11-12 14:32:26 +00001389... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001390
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001391>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001392... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001393... args = [iter(iterable)] * n
Raymond Hettingerf080e6d2008-07-31 01:19:50 +00001394... return izip_longest(fillvalue=fillvalue, *args)
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001395
1396>>> def roundrobin(*iterables):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001397... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001398... # Recipe credited to George Sakkis
1399... pending = len(iterables)
1400... nexts = cycle(iter(it).next for it in iterables)
1401... while pending:
1402... try:
1403... for next in nexts:
1404... yield next()
1405... except StopIteration:
1406... pending -= 1
1407... nexts = cycle(islice(nexts, pending))
1408
1409>>> def powerset(iterable):
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001410... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1411... s = list(iterable)
1412... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001413
Raymond Hettinger44e15812009-01-02 21:26:45 +00001414>>> def unique_everseen(iterable, key=None):
1415... "List unique elements, preserving order. Remember all elements ever seen."
1416... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1417... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1418... seen = set()
1419... seen_add = seen.add
1420... if key is None:
1421... for element in iterable:
1422... if element not in seen:
1423... seen_add(element)
1424... yield element
1425... else:
1426... for element in iterable:
1427... k = key(element)
1428... if k not in seen:
1429... seen_add(k)
1430... yield element
1431
1432>>> def unique_justseen(iterable, key=None):
1433... "List unique elements, preserving order. Remember only the element just seen."
1434... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1435... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1436... return imap(next, imap(itemgetter(1), groupby(iterable, key)))
1437
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001438This is not part of the examples but it tests to make sure the definitions
1439perform as purported.
1440
Raymond Hettingera098b332003-09-08 23:58:40 +00001441>>> take(10, count())
1442[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1443
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001444>>> list(enumerate('abc'))
1445[(0, 'a'), (1, 'b'), (2, 'c')]
1446
1447>>> list(islice(tabulate(lambda x: 2*x), 4))
1448[0, 2, 4, 6]
1449
1450>>> nth('abcde', 3)
Raymond Hettingerd507afd2009-02-04 10:52:32 +00001451'd'
1452
1453>>> nth('abcde', 9) is None
1454True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001455
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001456>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000145750
1458
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001459>>> a = [[1, 2, 3], [4, 5, 6]]
1460>>> flatten(a)
1461[1, 2, 3, 4, 5, 6]
1462
1463>>> list(repeatfunc(pow, 5, 2, 3))
1464[8, 8, 8, 8, 8]
1465
1466>>> import random
1467>>> take(5, imap(int, repeatfunc(random.random)))
1468[0, 0, 0, 0, 0]
1469
Raymond Hettingerd591f662003-10-26 15:34:50 +00001470>>> list(pairwise('abcd'))
1471[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001472
Raymond Hettingerd591f662003-10-26 15:34:50 +00001473>>> list(pairwise([]))
1474[]
1475
1476>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001477[]
1478
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001479>>> list(islice(padnone('abc'), 0, 6))
1480['a', 'b', 'c', None, None, None]
1481
1482>>> list(ncycles('abc', 3))
1483['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1484
1485>>> dotproduct([1,2,3], [4,5,6])
148632
1487
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001488>>> list(grouper(3, 'abcdefg', 'x'))
1489[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1490
1491>>> list(roundrobin('abc', 'd', 'ef'))
1492['a', 'd', 'e', 'b', 'f', 'c']
1493
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001494>>> list(powerset([1,2,3]))
1495[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001496
Raymond Hettinger560f9a82009-01-27 13:26:35 +00001497>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1498True
1499
1500>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1501True
1502
Raymond Hettinger44e15812009-01-02 21:26:45 +00001503>>> list(unique_everseen('AAAABBBCCDAABBB'))
1504['A', 'B', 'C', 'D']
1505
1506>>> list(unique_everseen('ABBCcAD', str.lower))
1507['A', 'B', 'C', 'D']
1508
1509>>> list(unique_justseen('AAAABBBCCDAABBB'))
1510['A', 'B', 'C', 'D', 'A', 'B']
1511
1512>>> list(unique_justseen('ABBCcAD', str.lower))
1513['A', 'B', 'C', 'A', 'D']
1514
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001515"""
1516
1517__test__ = {'libreftest' : libreftest}
1518
1519def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001520 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Georg Brandlb84c1372007-01-21 10:28:43 +00001521 RegressionTests, LengthTransparency,
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001522 SubclassWithKwargsTest, TestExamples)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001523 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001524
1525 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001526 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001527 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001528 counts = [None] * 5
1529 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001530 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001531 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001532 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001533 print counts
1534
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001535 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001536 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001537
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001538if __name__ == "__main__":
1539 test_main(verbose=True)