blob: 23a87654c7a7eafe70ad1fbcf0e5e9920c6faaba [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
109 for n in range(7):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000110 values = [5*x-12 for x in range(n)]
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000111 for r in range(n+2):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000112 result = list(combinations(values, r))
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000113 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 +0000114 self.assertEqual(len(result), len(set(result))) # no repeats
115 self.assertEqual(result, sorted(result)) # lexicographic order
116 for c in result:
117 self.assertEqual(len(c), r) # r-length combinations
118 self.assertEqual(len(set(c)), r) # no duplicate elements
119 self.assertEqual(list(c), sorted(c)) # keep original ordering
120 self.assert_(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000121 self.assertEqual(list(c),
122 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Raymond Hettingerd553d852008-03-04 04:17:08 +0000123 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000124 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000125
126 # Test implementation detail: tuple re-use
127 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
128 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
129
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000130 def test_combinations_with_replacement(self):
131 cwr = combinations_with_replacement
132 self.assertRaises(TypeError, cwr, 'abc') # missing r argument
133 self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
134 self.assertRaises(TypeError, cwr, None) # pool is not iterable
135 self.assertRaises(ValueError, cwr, 'abc', -2) # r is negative
136 self.assertEqual(list(cwr('ABC', 2)),
137 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
138
139 def cwr1(iterable, r):
140 'Pure python version shown in the docs'
141 # number items returned: (n+r-1)! / r! / (n-1)! when n>0
142 pool = tuple(iterable)
143 n = len(pool)
144 if not n and r:
145 return
146 indices = [0] * r
147 yield tuple(pool[i] for i in indices)
148 while 1:
149 for i in reversed(range(r)):
150 if indices[i] != n - 1:
151 break
152 else:
153 return
154 indices[i:] = [indices[i] + 1] * (r - i)
155 yield tuple(pool[i] for i in indices)
156
157 def cwr2(iterable, r):
158 'Pure python version shown in the docs'
159 pool = tuple(iterable)
160 n = len(pool)
161 for indices in product(range(n), repeat=r):
162 if sorted(indices) == list(indices):
163 yield tuple(pool[i] for i in indices)
164
165 def numcombs(n, r):
166 if not n:
167 return 0 if r else 1
168 return fact(n+r-1) / fact(r)/ fact(n-1)
169
170 for n in range(7):
171 values = [5*x-12 for x in range(n)]
172 for r in range(n+2):
173 result = list(cwr(values, r))
174
175 self.assertEqual(len(result), numcombs(n, r)) # right number of combs
176 self.assertEqual(len(result), len(set(result))) # no repeats
177 self.assertEqual(result, sorted(result)) # lexicographic order
178
179 regular_combs = list(combinations(values, r)) # compare to combs without replacement
180 if n == 0 or r <= 1:
181 self.assertEquals(result, regular_combs) # cases that should be identical
182 else:
183 self.assert_(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
184
185 for c in result:
186 self.assertEqual(len(c), r) # r-length combinations
187 noruns = [k for k,v in groupby(c)] # combo without consecutive repeats
188 self.assertEqual(len(noruns), len(set(noruns))) # no repeats other than consecutive
189 self.assertEqual(list(c), sorted(c)) # keep original ordering
190 self.assert_(all(e in values for e in c)) # elements taken from input iterable
191 self.assertEqual(noruns,
192 [e for e in values if e in c]) # comb is a subsequence of the input iterable
193 self.assertEqual(result, list(cwr1(values, r))) # matches first pure python version
194 self.assertEqual(result, list(cwr2(values, r))) # matches second pure python version
195
196 # Test implementation detail: tuple re-use
197 self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
198 self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
199
Raymond Hettingerd553d852008-03-04 04:17:08 +0000200 def test_permutations(self):
201 self.assertRaises(TypeError, permutations) # too few arguments
202 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000203 self.assertRaises(TypeError, permutations, None) # pool is not iterable
204 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000205 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000206 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Raymond Hettingerd553d852008-03-04 04:17:08 +0000207 self.assertEqual(list(permutations(range(3), 2)),
208 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
209
210 def permutations1(iterable, r=None):
211 'Pure python version shown in the docs'
212 pool = tuple(iterable)
213 n = len(pool)
214 r = n if r is None else r
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000215 if r > n:
216 return
Raymond Hettingerd553d852008-03-04 04:17:08 +0000217 indices = range(n)
Raymond Hettingere70bb8d2008-03-23 00:55:46 +0000218 cycles = range(n, n-r, -1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000219 yield tuple(pool[i] for i in indices[:r])
220 while n:
221 for i in reversed(range(r)):
222 cycles[i] -= 1
223 if cycles[i] == 0:
224 indices[i:] = indices[i+1:] + indices[i:i+1]
225 cycles[i] = n - i
226 else:
227 j = cycles[i]
228 indices[i], indices[-j] = indices[-j], indices[i]
229 yield tuple(pool[i] for i in indices[:r])
230 break
231 else:
232 return
233
234 def permutations2(iterable, r=None):
235 'Pure python version shown in the docs'
236 pool = tuple(iterable)
237 n = len(pool)
238 r = n if r is None else r
239 for indices in product(range(n), repeat=r):
240 if len(set(indices)) == r:
241 yield tuple(pool[i] for i in indices)
242
243 for n in range(7):
244 values = [5*x-12 for x in range(n)]
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000245 for r in range(n+2):
Raymond Hettingerd553d852008-03-04 04:17:08 +0000246 result = list(permutations(values, r))
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000247 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 +0000248 self.assertEqual(len(result), len(set(result))) # no repeats
249 self.assertEqual(result, sorted(result)) # lexicographic order
250 for p in result:
251 self.assertEqual(len(p), r) # r-length permutations
252 self.assertEqual(len(set(p)), r) # no duplicate elements
253 self.assert_(all(e in values for e in p)) # elements taken from input iterable
254 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000255 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000256 if r == n:
257 self.assertEqual(result, list(permutations(values, None))) # test r as None
258 self.assertEqual(result, list(permutations(values))) # test default r
259
260 # Test implementation detail: tuple re-use
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000261 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000262 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000263
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000264 def test_compress(self):
265 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
266 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
267 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
268 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
269 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
270 n = 10000
271 data = chain.from_iterable(repeat(range(6), n))
272 selectors = chain.from_iterable(repeat((0, 1)))
273 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
274 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
275 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
276 self.assertRaises(TypeError, compress, range(6)) # too few args
277 self.assertRaises(TypeError, compress, range(6), None) # too many args
278
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000279 def test_count(self):
280 self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
281 self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000282 self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000283 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
284 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000285 self.assertRaises(TypeError, count, 2, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000286 self.assertRaises(TypeError, count, 'a')
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000287 self.assertEqual(list(islice(count(maxsize-5), 10)), range(maxsize-5, maxsize+5))
288 self.assertEqual(list(islice(count(-maxsize-5), 10)), range(-maxsize-5, -maxsize+5))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000289 c = count(3)
290 self.assertEqual(repr(c), 'count(3)')
291 c.next()
292 self.assertEqual(repr(c), 'count(4)')
Jack Diederich36234e82006-09-21 17:50:26 +0000293 c = count(-9)
294 self.assertEqual(repr(c), 'count(-9)')
295 c.next()
296 self.assertEqual(c.next(), -8)
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000297 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 +0000298 # Test repr (ignoring the L in longs)
299 r1 = repr(count(i)).replace('L', '')
300 r2 = 'count(%r)'.__mod__(i).replace('L', '')
301 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000302
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000303 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000304 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000305 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000306 self.assertRaises(TypeError, cycle)
307 self.assertRaises(TypeError, cycle, 5)
308 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000309
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000310 def test_groupby(self):
311 # Check whether it accepts arguments correctly
312 self.assertEqual([], list(groupby([])))
313 self.assertEqual([], list(groupby([], key=id)))
314 self.assertRaises(TypeError, list, groupby('abc', []))
315 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000316 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000317
318 # Check normal input
319 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
320 (2,15,22), (3,16,23), (3,17,23)]
321 dup = []
322 for k, g in groupby(s, lambda r:r[0]):
323 for elem in g:
324 self.assertEqual(k, elem[0])
325 dup.append(elem)
326 self.assertEqual(s, dup)
327
328 # Check nested case
329 dup = []
330 for k, g in groupby(s, lambda r:r[0]):
331 for ik, ig in groupby(g, lambda r:r[2]):
332 for elem in ig:
333 self.assertEqual(k, elem[0])
334 self.assertEqual(ik, elem[2])
335 dup.append(elem)
336 self.assertEqual(s, dup)
337
338 # Check case where inner iterator is not used
339 keys = [k for k, g in groupby(s, lambda r:r[0])]
340 expectedkeys = set([r[0] for r in s])
341 self.assertEqual(set(keys), expectedkeys)
342 self.assertEqual(len(keys), len(expectedkeys))
343
344 # Exercise pipes and filters style
345 s = 'abracadabra'
346 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000347 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000348 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
349 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000350 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000351 self.assertEqual(r, ['a', 'b', 'r'])
352 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000353 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000354 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
355 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000356 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000357 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
358
359 # iter.next failure
360 class ExpectedError(Exception):
361 pass
362 def delayed_raise(n=0):
363 for i in range(n):
364 yield 'yo'
365 raise ExpectedError
366 def gulp(iterable, keyp=None, func=list):
367 return [func(g) for k, g in groupby(iterable, keyp)]
368
369 # iter.next failure on outer object
370 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
371 # iter.next failure on inner object
372 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
373
374 # __cmp__ failure
375 class DummyCmp:
376 def __cmp__(self, dst):
377 raise ExpectedError
378 s = [DummyCmp(), DummyCmp(), None]
379
380 # __cmp__ failure on outer object
381 self.assertRaises(ExpectedError, gulp, s, func=id)
382 # __cmp__ failure on inner object
383 self.assertRaises(ExpectedError, gulp, s)
384
385 # keyfunc failure
386 def keyfunc(obj):
387 if keyfunc.skip > 0:
388 keyfunc.skip -= 1
389 return obj
390 else:
391 raise ExpectedError
392
393 # keyfunc failure on outer object
394 keyfunc.skip = 0
395 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
396 keyfunc.skip = 1
397 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
398
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000399 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000400 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000401 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000402 self.assertEqual(list(ifilter(bool, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000403 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000404 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000405 self.assertRaises(TypeError, ifilter, lambda x:x)
406 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000407 self.assertRaises(TypeError, ifilter, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000408 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000409
410 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000411 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
412 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000413 self.assertEqual(list(ifilterfalse(bool, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000414 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000415 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000416 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
417 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000418 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000419 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000420
421 def test_izip(self):
422 ans = [(x,y) for x, y in izip('abc',count())]
423 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000424 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
425 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
Raymond Hettinger02420702003-06-29 20:36:23 +0000426 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000427 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000428 self.assertEqual(list(izip()), zip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000429 self.assertRaises(TypeError, izip, 3)
430 self.assertRaises(TypeError, izip, range(3), 3)
431 # Check tuple re-use (implementation detail)
432 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
433 zip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000434 self.assertEqual([pair for pair in izip('abc', 'def')],
435 zip('abc', 'def'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000436 ids = map(id, izip('abc', 'def'))
437 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000438 ids = map(id, list(izip('abc', 'def')))
439 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000440
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000441 def test_iziplongest(self):
442 for args in [
443 ['abc', range(6)],
444 [range(6), 'abc'],
445 [range(1000), range(2000,2100), range(3000,3050)],
446 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
447 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
448 ]:
449 target = map(None, *args)
450 self.assertEqual(list(izip_longest(*args)), target)
451 self.assertEqual(list(izip_longest(*args, **{})), target)
452 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
453 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Tim Petersea5962f2007-03-12 18:07:52 +0000454
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000455 self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
456
457 self.assertEqual(list(izip_longest()), zip())
458 self.assertEqual(list(izip_longest([])), zip([]))
459 self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
Tim Petersea5962f2007-03-12 18:07:52 +0000460
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000461 self.assertEqual(list(izip_longest('abc', 'defg', **{})), map(None, 'abc', 'defg')) # empty keyword dict
462 self.assertRaises(TypeError, izip_longest, 3)
463 self.assertRaises(TypeError, izip_longest, range(3), 3)
464
465 for stmt in [
466 "izip_longest('abc', fv=1)",
Tim Petersea5962f2007-03-12 18:07:52 +0000467 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000468 ]:
469 try:
470 eval(stmt, globals(), locals())
471 except TypeError:
472 pass
473 else:
474 self.fail('Did not raise Type in: ' + stmt)
Tim Petersea5962f2007-03-12 18:07:52 +0000475
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000476 # Check tuple re-use (implementation detail)
477 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
478 zip('abc', 'def'))
479 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
480 zip('abc', 'def'))
481 ids = map(id, izip_longest('abc', 'def'))
482 self.assertEqual(min(ids), max(ids))
483 ids = map(id, list(izip_longest('abc', 'def')))
484 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
485
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000486 def test_product(self):
487 for args, result in [
Raymond Hettingerd553d852008-03-04 04:17:08 +0000488 ([], [()]), # zero iterables
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000489 (['ab'], [('a',), ('b',)]), # one iterable
490 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
491 ([range(0), range(2), range(3)], []), # first iterable with zero length
492 ([range(2), range(0), range(3)], []), # middle iterable with zero length
493 ([range(2), range(3), range(0)], []), # last iterable with zero length
494 ]:
495 self.assertEqual(list(product(*args)), result)
Raymond Hettinger08ff6822008-02-29 02:21:48 +0000496 for r in range(4):
497 self.assertEqual(list(product(*(args*r))),
498 list(product(*args, **dict(repeat=r))))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000499 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
500 self.assertRaises(TypeError, product, range(6), None)
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000501
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000502 def product1(*args, **kwds):
503 pools = map(tuple, args) * kwds.get('repeat', 1)
504 n = len(pools)
505 if n == 0:
506 yield ()
507 return
508 if any(len(pool) == 0 for pool in pools):
509 return
510 indices = [0] * n
511 yield tuple(pool[i] for pool, i in zip(pools, indices))
512 while 1:
513 for i in reversed(range(n)): # right to left
514 if indices[i] == len(pools[i]) - 1:
515 continue
516 indices[i] += 1
517 for j in range(i+1, n):
518 indices[j] = 0
519 yield tuple(pool[i] for pool, i in zip(pools, indices))
520 break
521 else:
522 return
523
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000524 def product2(*args, **kwds):
525 'Pure python version used in docs'
526 pools = map(tuple, args) * kwds.get('repeat', 1)
527 result = [[]]
528 for pool in pools:
529 result = [x+[y] for x in result for y in pool]
530 for prod in result:
531 yield tuple(prod)
532
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000533 argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
534 set('abcdefg'), range(11), tuple(range(13))]
535 for i in range(100):
536 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Raymond Hettingerd553d852008-03-04 04:17:08 +0000537 expected_len = prod(map(len, args))
538 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000539 self.assertEqual(list(product(*args)), list(product1(*args)))
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000540 self.assertEqual(list(product(*args)), list(product2(*args)))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000541 args = map(iter, args)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000542 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000543
Raymond Hettinger73d79632008-02-23 02:20:41 +0000544 # Test implementation detail: tuple re-use
545 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
546 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000547
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000548 def test_repeat(self):
549 self.assertEqual(zip(xrange(3),repeat('a')),
550 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000551 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000552 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000553 self.assertEqual(list(repeat('a', 0)), [])
554 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000555 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000556 self.assertRaises(TypeError, repeat, None, 3, 4)
557 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000558 r = repeat(1+0j)
559 self.assertEqual(repr(r), 'repeat((1+0j))')
560 r = repeat(1+0j, 5)
561 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
562 list(r)
563 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000564
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000565 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000566 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
567 [0**1, 1**2, 2**3])
568 self.assertEqual(list(imap(None, 'abc', range(5))),
569 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000570 self.assertEqual(list(imap(None, 'abc', count())),
571 [('a',0),('b',1),('c',2)])
572 self.assertEqual(take(2,imap(None, 'abc', count())),
573 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000574 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000575 self.assertRaises(TypeError, imap)
576 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000577 self.assertRaises(TypeError, imap(10, range(5)).next)
578 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
579 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000580
581 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000582 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
583 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000584 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
585 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000586 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger47317092008-01-17 03:02:14 +0000587 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
588 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000589 self.assertRaises(TypeError, starmap)
590 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
591 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
592 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
593 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000594
595 def test_islice(self):
596 for args in [ # islice(args) should agree with range(args)
597 (10, 20, 3),
598 (10, 3, 20),
599 (10, 20),
600 (10, 3),
601 (20,)
602 ]:
603 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
604
605 for args, tgtargs in [ # Stop when seqn is exhausted
606 ((10, 110, 3), ((10, 100, 3))),
607 ((10, 110), ((10, 100))),
608 ((110,), (100,))
609 ]:
610 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
611
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000612 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000613 self.assertEqual(list(islice(xrange(10), None)), range(10))
Raymond Hettingerb2594052004-12-05 09:25:51 +0000614 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
615 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000616 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
617 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
618
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000619 # Test number of items consumed SF #1171417
620 it = iter(range(10))
621 self.assertEqual(list(islice(it, 3)), range(3))
622 self.assertEqual(list(it), range(3, 10))
623
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000624 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000625 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000626 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
627 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
628 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
629 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
630 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000631 self.assertRaises(ValueError, islice, xrange(10), 'a')
632 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
633 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
634 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
635 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +0000636 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000637
638 def test_takewhile(self):
639 data = [1, 3, 5, 20, 2, 4, 6, 8]
640 underten = lambda x: x<10
641 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000642 self.assertEqual(list(takewhile(underten, [])), [])
643 self.assertRaises(TypeError, takewhile)
644 self.assertRaises(TypeError, takewhile, operator.pow)
645 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
646 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
647 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000648 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
649 self.assertEqual(list(t), [1, 1, 1])
650 self.assertRaises(StopIteration, t.next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000651
652 def test_dropwhile(self):
653 data = [1, 3, 5, 20, 2, 4, 6, 8]
654 underten = lambda x: x<10
655 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000656 self.assertEqual(list(dropwhile(underten, [])), [])
657 self.assertRaises(TypeError, dropwhile)
658 self.assertRaises(TypeError, dropwhile, operator.pow)
659 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
660 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
661 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000662
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000663 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000664 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000665 def irange(n):
666 for i in xrange(n):
667 yield i
668
669 a, b = tee([]) # test empty iterator
670 self.assertEqual(list(a), [])
671 self.assertEqual(list(b), [])
672
673 a, b = tee(irange(n)) # test 100% interleaved
674 self.assertEqual(zip(a,b), zip(range(n),range(n)))
675
676 a, b = tee(irange(n)) # test 0% interleaved
677 self.assertEqual(list(a), range(n))
678 self.assertEqual(list(b), range(n))
679
680 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000681 for i in xrange(100):
682 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000683 del a
684 self.assertEqual(list(b), range(n))
685
686 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000687 for i in xrange(100):
688 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000689 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000690 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000691
692 for j in xrange(5): # test randomly interleaved
693 order = [0]*n + [1]*n
694 random.shuffle(order)
695 lists = ([], [])
696 its = tee(irange(n))
697 for i in order:
698 value = its[i].next()
699 lists[i].append(value)
700 self.assertEqual(lists[0], range(n))
701 self.assertEqual(lists[1], range(n))
702
Raymond Hettingerad983e72003-11-12 14:32:26 +0000703 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000704 self.assertRaises(TypeError, tee)
705 self.assertRaises(TypeError, tee, 3)
706 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000707 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000708
Raymond Hettingerad983e72003-11-12 14:32:26 +0000709 # tee object should be instantiable
710 a, b = tee('abc')
711 c = type(a)('def')
712 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000713
Raymond Hettingerad983e72003-11-12 14:32:26 +0000714 # test long-lagged and multi-way split
715 a, b, c = tee(xrange(2000), 3)
716 for i in xrange(100):
717 self.assertEqual(a.next(), i)
718 self.assertEqual(list(b), range(2000))
719 self.assertEqual([c.next(), c.next()], range(2))
720 self.assertEqual(list(a), range(100,2000))
721 self.assertEqual(list(c), range(2,2000))
722
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000723 # test values of n
724 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Neal Norwitz69e88972006-09-02 02:58:13 +0000725 self.assertRaises(ValueError, tee, [], -1)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000726 for n in xrange(5):
727 result = tee('abc', n)
728 self.assertEqual(type(result), tuple)
729 self.assertEqual(len(result), n)
730 self.assertEqual(map(list, result), [list('abc')]*n)
731
Raymond Hettingerad983e72003-11-12 14:32:26 +0000732 # tee pass-through to copyable iterator
733 a, b = tee('abc')
734 c, d = tee(a)
735 self.assert_(a is c)
736
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000737 # test tee_new
738 t1, t2 = tee('abc')
739 tnew = type(t1)
740 self.assertRaises(TypeError, tnew)
741 self.assertRaises(TypeError, tnew, 10)
742 t3 = tnew(t1)
743 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000744
Raymond Hettingera9f60922004-10-17 16:40:14 +0000745 # test that tee objects are weak referencable
746 a, b = tee(xrange(10))
747 p = proxy(a)
748 self.assertEqual(getattr(p, '__class__'), type(b))
749 del a
750 self.assertRaises(ReferenceError, getattr, p, '__class__')
751
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000752 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000753 self.assertRaises(StopIteration, izip().next)
754
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000755 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000756 self.assertRaises(StopIteration, f([]).next)
757 self.assertRaises(StopIteration, f(StopNow()).next)
758
759 self.assertRaises(StopIteration, islice([], None).next)
760 self.assertRaises(StopIteration, islice(StopNow(), None).next)
761
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000762 p, q = tee([])
763 self.assertRaises(StopIteration, p.next)
764 self.assertRaises(StopIteration, q.next)
765 p, q = tee(StopNow())
766 self.assertRaises(StopIteration, p.next)
767 self.assertRaises(StopIteration, q.next)
768
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000769 self.assertRaises(StopIteration, repeat(None, 0).next)
770
771 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
772 self.assertRaises(StopIteration, f(lambda x:x, []).next)
773 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
774
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000775class TestExamples(unittest.TestCase):
776
777 def test_chain(self):
778 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
779
780 def test_chain_from_iterable(self):
781 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
782
783 def test_combinations(self):
784 self.assertEqual(list(combinations('ABCD', 2)),
785 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
786 self.assertEqual(list(combinations(range(4), 3)),
787 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
788
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000789 def test_combinations_with_replacement(self):
790 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
791 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
792
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000793 def test_compress(self):
794 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
795
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000796 def test_count(self):
797 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
798
799 def test_cycle(self):
800 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
801
802 def test_dropwhile(self):
803 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
804
805 def test_groupby(self):
806 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
807 list('ABCDAB'))
808 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
809 [list('AAAA'), list('BBB'), list('CC'), list('D')])
810
811 def test_ifilter(self):
812 self.assertEqual(list(ifilter(lambda x: x%2, range(10))), [1,3,5,7,9])
813
814 def test_ifilterfalse(self):
815 self.assertEqual(list(ifilterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
816
817 def test_imap(self):
818 self.assertEqual(list(imap(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
819
820 def test_islice(self):
821 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
822 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
823 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
824 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
825
826 def test_izip(self):
827 self.assertEqual(list(izip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
828
829 def test_izip_longest(self):
830 self.assertEqual(list(izip_longest('ABCD', 'xy', fillvalue='-')),
831 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
832
833 def test_permutations(self):
834 self.assertEqual(list(permutations('ABCD', 2)),
835 map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split()))
836 self.assertEqual(list(permutations(range(3))),
837 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
838
839 def test_product(self):
840 self.assertEqual(list(product('ABCD', 'xy')),
841 map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split()))
842 self.assertEqual(list(product(range(2), repeat=3)),
843 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
844 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
845
846 def test_repeat(self):
847 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
848
849 def test_stapmap(self):
850 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
851 [32, 9, 1000])
852
853 def test_takewhile(self):
854 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
855
856
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000857class TestGC(unittest.TestCase):
858
859 def makecycle(self, iterator, container):
860 container.append(iterator)
861 iterator.next()
862 del container, iterator
863
864 def test_chain(self):
865 a = []
866 self.makecycle(chain(a), a)
867
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000868 def test_chain_from_iterable(self):
869 a = []
870 self.makecycle(chain.from_iterable([a]), a)
871
872 def test_combinations(self):
873 a = []
874 self.makecycle(combinations([1,2,a,3], 3), a)
875
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000876 def test_combinations_with_replacement(self):
877 a = []
878 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
879
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000880 def test_compress(self):
881 a = []
882 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
883
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000884 def test_cycle(self):
885 a = []
886 self.makecycle(cycle([a]*2), a)
887
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000888 def test_dropwhile(self):
889 a = []
890 self.makecycle(dropwhile(bool, [0, a, a]), a)
891
892 def test_groupby(self):
893 a = []
894 self.makecycle(groupby([a]*2, lambda x:x), a)
895
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000896 def test_issue2246(self):
897 # Issue 2246 -- the _grouper iterator was not included in GC
898 n = 10
899 keyfunc = lambda x: x
900 for i, j in groupby(xrange(n), key=keyfunc):
901 keyfunc.__dict__.setdefault('x',[]).append(j)
902
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000903 def test_ifilter(self):
904 a = []
905 self.makecycle(ifilter(lambda x:True, [a]*2), a)
906
907 def test_ifilterfalse(self):
908 a = []
909 self.makecycle(ifilterfalse(lambda x:False, a), a)
910
911 def test_izip(self):
912 a = []
913 self.makecycle(izip([a]*2, [a]*3), a)
914
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000915 def test_izip_longest(self):
916 a = []
917 self.makecycle(izip_longest([a]*2, [a]*3), a)
918 b = [a, None]
919 self.makecycle(izip_longest([a]*2, [a]*3, fillvalue=b), a)
920
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000921 def test_imap(self):
922 a = []
923 self.makecycle(imap(lambda x:x, [a]*2), a)
924
925 def test_islice(self):
926 a = []
927 self.makecycle(islice([a]*2, None), a)
928
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000929 def test_permutations(self):
930 a = []
931 self.makecycle(permutations([1,2,a,3], 3), a)
932
933 def test_product(self):
934 a = []
935 self.makecycle(product([1,2,a,3], repeat=3), a)
936
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000937 def test_repeat(self):
938 a = []
939 self.makecycle(repeat(a), a)
940
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000941 def test_starmap(self):
942 a = []
943 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
944
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000945 def test_takewhile(self):
946 a = []
947 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
948
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000949def R(seqn):
950 'Regular generator'
951 for i in seqn:
952 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000953
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000954class G:
955 'Sequence using __getitem__'
956 def __init__(self, seqn):
957 self.seqn = seqn
958 def __getitem__(self, i):
959 return self.seqn[i]
960
961class I:
962 'Sequence using iterator protocol'
963 def __init__(self, seqn):
964 self.seqn = seqn
965 self.i = 0
966 def __iter__(self):
967 return self
968 def next(self):
969 if self.i >= len(self.seqn): raise StopIteration
970 v = self.seqn[self.i]
971 self.i += 1
972 return v
973
974class Ig:
975 'Sequence using iterator protocol defined with a generator'
976 def __init__(self, seqn):
977 self.seqn = seqn
978 self.i = 0
979 def __iter__(self):
980 for val in self.seqn:
981 yield val
982
983class X:
984 'Missing __getitem__ and __iter__'
985 def __init__(self, seqn):
986 self.seqn = seqn
987 self.i = 0
988 def next(self):
989 if self.i >= len(self.seqn): raise StopIteration
990 v = self.seqn[self.i]
991 self.i += 1
992 return v
993
994class N:
995 'Iterator missing next()'
996 def __init__(self, seqn):
997 self.seqn = seqn
998 self.i = 0
999 def __iter__(self):
1000 return self
1001
1002class E:
1003 'Test propagation of exceptions'
1004 def __init__(self, seqn):
1005 self.seqn = seqn
1006 self.i = 0
1007 def __iter__(self):
1008 return self
1009 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001010 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001011
1012class S:
1013 'Test immediate stop'
1014 def __init__(self, seqn):
1015 pass
1016 def __iter__(self):
1017 return self
1018 def next(self):
1019 raise StopIteration
1020
1021def L(seqn):
1022 'Test multiple tiers of iterators'
1023 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
1024
1025
1026class TestVariousIteratorArgs(unittest.TestCase):
1027
1028 def test_chain(self):
1029 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1030 for g in (G, I, Ig, S, L, R):
1031 self.assertEqual(list(chain(g(s))), list(g(s)))
1032 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001033 self.assertRaises(TypeError, list, chain(X(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001034 self.assertRaises(TypeError, list, chain(N(s)))
1035 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1036
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001037 def test_compress(self):
1038 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1039 n = len(s)
1040 for g in (G, I, Ig, S, L, R):
1041 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1042 self.assertRaises(TypeError, compress, X(s), repeat(1))
1043 self.assertRaises(TypeError, list, compress(N(s), repeat(1)))
1044 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1045
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001046 def test_product(self):
1047 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1048 self.assertRaises(TypeError, product, X(s))
1049 self.assertRaises(TypeError, product, N(s))
1050 self.assertRaises(ZeroDivisionError, product, E(s))
1051
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001052 def test_cycle(self):
1053 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1054 for g in (G, I, Ig, S, L, R):
1055 tgtlen = len(s) * 3
1056 expected = list(g(s))*3
1057 actual = list(islice(cycle(g(s)), tgtlen))
1058 self.assertEqual(actual, expected)
1059 self.assertRaises(TypeError, cycle, X(s))
1060 self.assertRaises(TypeError, list, cycle(N(s)))
1061 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1062
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001063 def test_groupby(self):
1064 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1065 for g in (G, I, Ig, S, L, R):
1066 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1067 self.assertRaises(TypeError, groupby, X(s))
1068 self.assertRaises(TypeError, list, groupby(N(s)))
1069 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1070
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001071 def test_ifilter(self):
1072 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1073 for g in (G, I, Ig, S, L, R):
1074 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
1075 self.assertRaises(TypeError, ifilter, isEven, X(s))
1076 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
1077 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
1078
1079 def test_ifilterfalse(self):
1080 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1081 for g in (G, I, Ig, S, L, R):
1082 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
1083 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
1084 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
1085 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
1086
1087 def test_izip(self):
1088 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1089 for g in (G, I, Ig, S, L, R):
1090 self.assertEqual(list(izip(g(s))), zip(g(s)))
1091 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
1092 self.assertRaises(TypeError, izip, X(s))
1093 self.assertRaises(TypeError, list, izip(N(s)))
1094 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
1095
Raymond Hettingerd36862c2007-02-21 05:20:38 +00001096 def test_iziplongest(self):
1097 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1098 for g in (G, I, Ig, S, L, R):
1099 self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
1100 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
1101 self.assertRaises(TypeError, izip_longest, X(s))
1102 self.assertRaises(TypeError, list, izip_longest(N(s)))
1103 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
1104
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001105 def test_imap(self):
1106 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1107 for g in (G, I, Ig, S, L, R):
1108 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
1109 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
1110 self.assertRaises(TypeError, imap, onearg, X(s))
1111 self.assertRaises(TypeError, list, imap(onearg, N(s)))
1112 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
1113
1114 def test_islice(self):
1115 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1116 for g in (G, I, Ig, S, L, R):
1117 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1118 self.assertRaises(TypeError, islice, X(s), 10)
1119 self.assertRaises(TypeError, list, islice(N(s), 10))
1120 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1121
1122 def test_starmap(self):
1123 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1124 for g in (G, I, Ig, S, L, R):
1125 ss = zip(s, s)
1126 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
1127 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
1128 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
1129 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1130
1131 def test_takewhile(self):
1132 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1133 for g in (G, I, Ig, S, L, R):
1134 tgt = []
1135 for elem in g(s):
1136 if not isEven(elem): break
1137 tgt.append(elem)
1138 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1139 self.assertRaises(TypeError, takewhile, isEven, X(s))
1140 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
1141 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1142
1143 def test_dropwhile(self):
1144 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1145 for g in (G, I, Ig, S, L, R):
1146 tgt = []
1147 for elem in g(s):
1148 if not tgt and isOdd(elem): continue
1149 tgt.append(elem)
1150 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1151 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
1152 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
1153 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1154
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001155 def test_tee(self):
1156 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1157 for g in (G, I, Ig, S, L, R):
1158 it1, it2 = tee(g(s))
1159 self.assertEqual(list(it1), list(g(s)))
1160 self.assertEqual(list(it2), list(g(s)))
1161 self.assertRaises(TypeError, tee, X(s))
1162 self.assertRaises(TypeError, list, tee(N(s))[0])
1163 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1164
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001165class LengthTransparency(unittest.TestCase):
1166
1167 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001168 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001169 self.assertEqual(len(repeat(None, 50)), 50)
1170 self.assertRaises(TypeError, len, repeat(None))
1171
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001172class RegressionTests(unittest.TestCase):
1173
1174 def test_sf_793826(self):
1175 # Fix Armin Rigo's successful efforts to wreak havoc
1176
1177 def mutatingtuple(tuple1, f, tuple2):
1178 # this builds a tuple t which is a copy of tuple1,
1179 # then calls f(t), then mutates t to be equal to tuple2
1180 # (needs len(tuple1) == len(tuple2)).
1181 def g(value, first=[1]):
1182 if first:
1183 del first[:]
1184 f(z.next())
1185 return value
1186 items = list(tuple2)
1187 items[1:1] = list(tuple1)
1188 gen = imap(g, items)
1189 z = izip(*[gen]*len(tuple1))
1190 z.next()
1191
1192 def f(t):
1193 global T
1194 T = t
1195 first[:] = list(T)
1196
1197 first = []
1198 mutatingtuple((1,2,3), f, (4,5,6))
1199 second = list(T)
1200 self.assertEqual(first, second)
1201
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001202
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001203 def test_sf_950057(self):
1204 # Make sure that chain() and cycle() catch exceptions immediately
1205 # rather than when shifting between input sources
1206
1207 def gen1():
1208 hist.append(0)
1209 yield 1
1210 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001211 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001212 hist.append(2)
1213
1214 def gen2(x):
1215 hist.append(3)
1216 yield 2
1217 hist.append(4)
1218 if x:
1219 raise StopIteration
1220
1221 hist = []
1222 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1223 self.assertEqual(hist, [0,1])
1224
1225 hist = []
1226 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1227 self.assertEqual(hist, [0,1])
1228
1229 hist = []
1230 self.assertRaises(AssertionError, list, cycle(gen1()))
1231 self.assertEqual(hist, [0,1])
1232
Georg Brandlb84c1372007-01-21 10:28:43 +00001233class SubclassWithKwargsTest(unittest.TestCase):
1234 def test_keywords_in_subclass(self):
1235 # count is not subclassable...
1236 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001237 starmap, islice, takewhile, dropwhile, cycle, compress):
Georg Brandlb84c1372007-01-21 10:28:43 +00001238 class Subclass(cls):
1239 def __init__(self, newarg=None, *args):
1240 cls.__init__(self, *args)
1241 try:
1242 Subclass(newarg=1)
1243 except TypeError, err:
1244 # we expect type errors because of wrong argument count
1245 self.failIf("does not take keyword arguments" in err.args[0])
1246
1247
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001248libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001249
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001250
1251>>> amounts = [120.15, 764.05, 823.14]
1252>>> for checknum, amount in izip(count(1200), amounts):
1253... print 'Check %d is for $%.2f' % (checknum, amount)
1254...
1255Check 1200 is for $120.15
1256Check 1201 is for $764.05
1257Check 1202 is for $823.14
1258
1259>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001260>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1261... print cube
1262...
12631
12648
126527
1266
1267>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001268>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001269... print name.title()
1270...
1271Alex
1272Laura
1273Martin
1274Walter
1275Samuele
1276
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001277>>> from operator import itemgetter
1278>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Armin Rigoa3f09272006-05-28 19:13:17 +00001279>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001280>>> for k, g in groupby(di, itemgetter(1)):
1281... print k, map(itemgetter(0), g)
1282...
12831 ['a', 'c', 'e']
12842 ['b', 'd', 'f']
12853 ['g']
1286
Raymond Hettinger734fb572004-01-20 20:04:40 +00001287# Find runs of consecutive numbers using groupby. The key to the solution
1288# is differencing with a range so that consecutive numbers all appear in
1289# same group.
1290>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
1291>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
1292... print map(operator.itemgetter(1), g)
1293...
1294[1]
1295[4, 5, 6]
1296[10]
1297[15, 16, 17, 18]
1298[22]
1299[25, 26, 27, 28]
1300
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001301>>> def take(n, iterable):
1302... "Return first n items of the iterable as a list"
1303... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001304
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001305>>> def enumerate(iterable, start=0):
1306... return izip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001307
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001308>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001309... "Return function(0), function(1), ..."
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001310... return imap(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001311
1312>>> def nth(iterable, n):
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001313... "Returns the nth item or empty list"
Raymond Hettinger60eca932003-02-09 06:40:58 +00001314... return list(islice(iterable, n, n+1))
1315
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001316>>> def quantify(iterable, pred=bool):
1317... "Count how many times the predicate is true"
1318... return sum(imap(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001319
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001320>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001321... "Returns the sequence elements and then returns None indefinitely"
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001322... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001323
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001324>>> def ncycles(iterable, n):
1325... "Returns the seqeuence elements n times"
1326... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001327
1328>>> def dotproduct(vec1, vec2):
1329... return sum(imap(operator.mul, vec1, vec2))
1330
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001331>>> def flatten(listOfLists):
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001332... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001333
1334>>> def repeatfunc(func, times=None, *args):
1335... "Repeat calls to func with specified arguments."
1336... " Example: repeatfunc(random.random)"
1337... if times is None:
1338... return starmap(func, repeat(args))
1339... else:
1340... return starmap(func, repeat(args, times))
1341
Raymond Hettingerd591f662003-10-26 15:34:50 +00001342>>> def pairwise(iterable):
1343... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1344... a, b = tee(iterable)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001345... for elem in b:
1346... break
Raymond Hettingerad983e72003-11-12 14:32:26 +00001347... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001348
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001349>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001350... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001351... args = [iter(iterable)] * n
Raymond Hettingerf080e6d2008-07-31 01:19:50 +00001352... return izip_longest(fillvalue=fillvalue, *args)
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001353
1354>>> def roundrobin(*iterables):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001355... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001356... # Recipe credited to George Sakkis
1357... pending = len(iterables)
1358... nexts = cycle(iter(it).next for it in iterables)
1359... while pending:
1360... try:
1361... for next in nexts:
1362... yield next()
1363... except StopIteration:
1364... pending -= 1
1365... nexts = cycle(islice(nexts, pending))
1366
1367>>> def powerset(iterable):
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001368... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1369... s = list(iterable)
1370... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001371
Raymond Hettinger44e15812009-01-02 21:26:45 +00001372>>> def unique_everseen(iterable, key=None):
1373... "List unique elements, preserving order. Remember all elements ever seen."
1374... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1375... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1376... seen = set()
1377... seen_add = seen.add
1378... if key is None:
1379... for element in iterable:
1380... if element not in seen:
1381... seen_add(element)
1382... yield element
1383... else:
1384... for element in iterable:
1385... k = key(element)
1386... if k not in seen:
1387... seen_add(k)
1388... yield element
1389
1390>>> def unique_justseen(iterable, key=None):
1391... "List unique elements, preserving order. Remember only the element just seen."
1392... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1393... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1394... return imap(next, imap(itemgetter(1), groupby(iterable, key)))
1395
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001396This is not part of the examples but it tests to make sure the definitions
1397perform as purported.
1398
Raymond Hettingera098b332003-09-08 23:58:40 +00001399>>> take(10, count())
1400[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1401
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001402>>> list(enumerate('abc'))
1403[(0, 'a'), (1, 'b'), (2, 'c')]
1404
1405>>> list(islice(tabulate(lambda x: 2*x), 4))
1406[0, 2, 4, 6]
1407
1408>>> nth('abcde', 3)
1409['d']
1410
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001411>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000141250
1413
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001414>>> a = [[1, 2, 3], [4, 5, 6]]
1415>>> flatten(a)
1416[1, 2, 3, 4, 5, 6]
1417
1418>>> list(repeatfunc(pow, 5, 2, 3))
1419[8, 8, 8, 8, 8]
1420
1421>>> import random
1422>>> take(5, imap(int, repeatfunc(random.random)))
1423[0, 0, 0, 0, 0]
1424
Raymond Hettingerd591f662003-10-26 15:34:50 +00001425>>> list(pairwise('abcd'))
1426[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001427
Raymond Hettingerd591f662003-10-26 15:34:50 +00001428>>> list(pairwise([]))
1429[]
1430
1431>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001432[]
1433
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001434>>> list(islice(padnone('abc'), 0, 6))
1435['a', 'b', 'c', None, None, None]
1436
1437>>> list(ncycles('abc', 3))
1438['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1439
1440>>> dotproduct([1,2,3], [4,5,6])
144132
1442
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001443>>> list(grouper(3, 'abcdefg', 'x'))
1444[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1445
1446>>> list(roundrobin('abc', 'd', 'ef'))
1447['a', 'd', 'e', 'b', 'f', 'c']
1448
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001449>>> list(powerset([1,2,3]))
1450[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001451
Raymond Hettinger44e15812009-01-02 21:26:45 +00001452>>> list(unique_everseen('AAAABBBCCDAABBB'))
1453['A', 'B', 'C', 'D']
1454
1455>>> list(unique_everseen('ABBCcAD', str.lower))
1456['A', 'B', 'C', 'D']
1457
1458>>> list(unique_justseen('AAAABBBCCDAABBB'))
1459['A', 'B', 'C', 'D', 'A', 'B']
1460
1461>>> list(unique_justseen('ABBCcAD', str.lower))
1462['A', 'B', 'C', 'A', 'D']
1463
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001464"""
1465
1466__test__ = {'libreftest' : libreftest}
1467
1468def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001469 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Georg Brandlb84c1372007-01-21 10:28:43 +00001470 RegressionTests, LengthTransparency,
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001471 SubclassWithKwargsTest, TestExamples)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001472 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001473
1474 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001475 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001476 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001477 counts = [None] * 5
1478 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001479 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001480 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001481 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001482 print counts
1483
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001484 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001485 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001486
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001487if __name__ == "__main__":
1488 test_main(verbose=True)