Issue #5788: `datetime.timedelta` objects get a new `total_seconds()` method returning
[python.git] / Lib / test / test_itertools.py
blobf8bdb8bdec9fe713d0c28c5ac9a2aaca479c186b
1 import unittest
2 from test import test_support
3 from itertools import *
4 from weakref import proxy
5 from decimal import Decimal
6 from fractions import Fraction
7 import sys
8 import operator
9 import random
10 maxsize = test_support.MAX_Py_ssize_t
11 minsize = -maxsize-1
13 def onearg(x):
14 'Test function of one argument'
15 return 2*x
17 def errfunc(*args):
18 'Test function that raises an error'
19 raise ValueError
21 def gen3():
22 'Non-restartable source sequence'
23 for i in (0, 1, 2):
24 yield i
26 def isEven(x):
27 'Test predicate'
28 return x%2==0
30 def isOdd(x):
31 'Test predicate'
32 return x%2==1
34 class StopNow:
35 'Class emulating an empty iterable.'
36 def __iter__(self):
37 return self
38 def next(self):
39 raise StopIteration
41 def take(n, seq):
42 'Convenience function for partially consuming a long of infinite iterable'
43 return list(islice(seq, n))
45 def prod(iterable):
46 return reduce(operator.mul, iterable, 1)
48 def fact(n):
49 'Factorial'
50 return prod(range(1, n+1))
52 class TestBasicOps(unittest.TestCase):
53 def test_chain(self):
55 def chain2(*iterables):
56 'Pure python version in the docs'
57 for it in iterables:
58 for element in it:
59 yield element
61 for c in (chain, chain2):
62 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
63 self.assertEqual(list(c('abc')), list('abc'))
64 self.assertEqual(list(c('')), [])
65 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
66 self.assertRaises(TypeError, list,c(2, 3))
68 def test_chain_from_iterable(self):
69 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
70 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
71 self.assertEqual(list(chain.from_iterable([''])), [])
72 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
73 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
75 def test_combinations(self):
76 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
77 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
78 self.assertRaises(TypeError, combinations, None) # pool is not iterable
79 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
80 self.assertEqual(list(combinations('abc', 32)), []) # r > n
81 self.assertEqual(list(combinations(range(4), 3)),
82 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
84 def combinations1(iterable, r):
85 'Pure python version shown in the docs'
86 pool = tuple(iterable)
87 n = len(pool)
88 if r > n:
89 return
90 indices = range(r)
91 yield tuple(pool[i] for i in indices)
92 while 1:
93 for i in reversed(range(r)):
94 if indices[i] != i + n - r:
95 break
96 else:
97 return
98 indices[i] += 1
99 for j in range(i+1, r):
100 indices[j] = indices[j-1] + 1
101 yield tuple(pool[i] for i in indices)
103 def combinations2(iterable, r):
104 'Pure python version shown in the docs'
105 pool = tuple(iterable)
106 n = len(pool)
107 for indices in permutations(range(n), r):
108 if sorted(indices) == list(indices):
109 yield tuple(pool[i] for i in indices)
111 def combinations3(iterable, r):
112 'Pure python version from cwr()'
113 pool = tuple(iterable)
114 n = len(pool)
115 for indices in combinations_with_replacement(range(n), r):
116 if len(set(indices)) == r:
117 yield tuple(pool[i] for i in indices)
119 for n in range(7):
120 values = [5*x-12 for x in range(n)]
121 for r in range(n+2):
122 result = list(combinations(values, r))
123 self.assertEqual(len(result), 0 if r>n else fact(n) / fact(r) / fact(n-r)) # right number of combs
124 self.assertEqual(len(result), len(set(result))) # no repeats
125 self.assertEqual(result, sorted(result)) # lexicographic order
126 for c in result:
127 self.assertEqual(len(c), r) # r-length combinations
128 self.assertEqual(len(set(c)), r) # no duplicate elements
129 self.assertEqual(list(c), sorted(c)) # keep original ordering
130 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
131 self.assertEqual(list(c),
132 [e for e in values if e in c]) # comb is a subsequence of the input iterable
133 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
134 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
135 self.assertEqual(result, list(combinations3(values, r))) # matches second pure python version
137 # Test implementation detail: tuple re-use
138 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
139 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
141 def test_combinations_with_replacement(self):
142 cwr = combinations_with_replacement
143 self.assertRaises(TypeError, cwr, 'abc') # missing r argument
144 self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
145 self.assertRaises(TypeError, cwr, None) # pool is not iterable
146 self.assertRaises(ValueError, cwr, 'abc', -2) # r is negative
147 self.assertEqual(list(cwr('ABC', 2)),
148 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
150 def cwr1(iterable, r):
151 'Pure python version shown in the docs'
152 # number items returned: (n+r-1)! / r! / (n-1)! when n>0
153 pool = tuple(iterable)
154 n = len(pool)
155 if not n and r:
156 return
157 indices = [0] * r
158 yield tuple(pool[i] for i in indices)
159 while 1:
160 for i in reversed(range(r)):
161 if indices[i] != n - 1:
162 break
163 else:
164 return
165 indices[i:] = [indices[i] + 1] * (r - i)
166 yield tuple(pool[i] for i in indices)
168 def cwr2(iterable, r):
169 'Pure python version shown in the docs'
170 pool = tuple(iterable)
171 n = len(pool)
172 for indices in product(range(n), repeat=r):
173 if sorted(indices) == list(indices):
174 yield tuple(pool[i] for i in indices)
176 def numcombs(n, r):
177 if not n:
178 return 0 if r else 1
179 return fact(n+r-1) / fact(r)/ fact(n-1)
181 for n in range(7):
182 values = [5*x-12 for x in range(n)]
183 for r in range(n+2):
184 result = list(cwr(values, r))
186 self.assertEqual(len(result), numcombs(n, r)) # right number of combs
187 self.assertEqual(len(result), len(set(result))) # no repeats
188 self.assertEqual(result, sorted(result)) # lexicographic order
190 regular_combs = list(combinations(values, r)) # compare to combs without replacement
191 if n == 0 or r <= 1:
192 self.assertEquals(result, regular_combs) # cases that should be identical
193 else:
194 self.assertTrue(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
196 for c in result:
197 self.assertEqual(len(c), r) # r-length combinations
198 noruns = [k for k,v in groupby(c)] # combo without consecutive repeats
199 self.assertEqual(len(noruns), len(set(noruns))) # no repeats other than consecutive
200 self.assertEqual(list(c), sorted(c)) # keep original ordering
201 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
202 self.assertEqual(noruns,
203 [e for e in values if e in c]) # comb is a subsequence of the input iterable
204 self.assertEqual(result, list(cwr1(values, r))) # matches first pure python version
205 self.assertEqual(result, list(cwr2(values, r))) # matches second pure python version
207 # Test implementation detail: tuple re-use
208 self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
209 self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
211 def test_permutations(self):
212 self.assertRaises(TypeError, permutations) # too few arguments
213 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
214 self.assertRaises(TypeError, permutations, None) # pool is not iterable
215 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
216 self.assertEqual(list(permutations('abc', 32)), []) # r > n
217 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
218 self.assertEqual(list(permutations(range(3), 2)),
219 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
221 def permutations1(iterable, r=None):
222 'Pure python version shown in the docs'
223 pool = tuple(iterable)
224 n = len(pool)
225 r = n if r is None else r
226 if r > n:
227 return
228 indices = range(n)
229 cycles = range(n, n-r, -1)
230 yield tuple(pool[i] for i in indices[:r])
231 while n:
232 for i in reversed(range(r)):
233 cycles[i] -= 1
234 if cycles[i] == 0:
235 indices[i:] = indices[i+1:] + indices[i:i+1]
236 cycles[i] = n - i
237 else:
238 j = cycles[i]
239 indices[i], indices[-j] = indices[-j], indices[i]
240 yield tuple(pool[i] for i in indices[:r])
241 break
242 else:
243 return
245 def permutations2(iterable, r=None):
246 'Pure python version shown in the docs'
247 pool = tuple(iterable)
248 n = len(pool)
249 r = n if r is None else r
250 for indices in product(range(n), repeat=r):
251 if len(set(indices)) == r:
252 yield tuple(pool[i] for i in indices)
254 for n in range(7):
255 values = [5*x-12 for x in range(n)]
256 for r in range(n+2):
257 result = list(permutations(values, r))
258 self.assertEqual(len(result), 0 if r>n else fact(n) / fact(n-r)) # right number of perms
259 self.assertEqual(len(result), len(set(result))) # no repeats
260 self.assertEqual(result, sorted(result)) # lexicographic order
261 for p in result:
262 self.assertEqual(len(p), r) # r-length permutations
263 self.assertEqual(len(set(p)), r) # no duplicate elements
264 self.assertTrue(all(e in values for e in p)) # elements taken from input iterable
265 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
266 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
267 if r == n:
268 self.assertEqual(result, list(permutations(values, None))) # test r as None
269 self.assertEqual(result, list(permutations(values))) # test default r
271 # Test implementation detail: tuple re-use
272 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
273 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
275 def test_combinatorics(self):
276 # Test relationships between product(), permutations(),
277 # combinations() and combinations_with_replacement().
279 for n in range(6):
280 s = 'ABCDEFG'[:n]
281 for r in range(8):
282 prod = list(product(s, repeat=r))
283 cwr = list(combinations_with_replacement(s, r))
284 perm = list(permutations(s, r))
285 comb = list(combinations(s, r))
287 # Check size
288 self.assertEquals(len(prod), n**r)
289 self.assertEquals(len(cwr), (fact(n+r-1) / fact(r)/ fact(n-1)) if n else (not r))
290 self.assertEquals(len(perm), 0 if r>n else fact(n) / fact(n-r))
291 self.assertEquals(len(comb), 0 if r>n else fact(n) / fact(r) / fact(n-r))
293 # Check lexicographic order without repeated tuples
294 self.assertEquals(prod, sorted(set(prod)))
295 self.assertEquals(cwr, sorted(set(cwr)))
296 self.assertEquals(perm, sorted(set(perm)))
297 self.assertEquals(comb, sorted(set(comb)))
299 # Check interrelationships
300 self.assertEquals(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
301 self.assertEquals(perm, [t for t in prod if len(set(t))==r]) # perm: prods with no dups
302 self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
303 self.assertEqual(comb, [t for t in cwr if len(set(t))==r]) # comb: cwrs without dups
304 self.assertEqual(comb, filter(set(cwr).__contains__, perm)) # comb: perm that is a cwr
305 self.assertEqual(comb, filter(set(perm).__contains__, cwr)) # comb: cwr that is a perm
306 self.assertEqual(comb, sorted(set(cwr) & set(perm))) # comb: both a cwr and a perm
308 def test_compress(self):
309 self.assertEqual(list(compress(data='ABCDEF', selectors=[1,0,1,0,1,1])), list('ACEF'))
310 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
311 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
312 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
313 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
314 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
315 n = 10000
316 data = chain.from_iterable(repeat(range(6), n))
317 selectors = chain.from_iterable(repeat((0, 1)))
318 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
319 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
320 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
321 self.assertRaises(TypeError, compress, range(6)) # too few args
322 self.assertRaises(TypeError, compress, range(6), None) # too many args
324 def test_count(self):
325 self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
326 self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
327 self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
328 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
329 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
330 self.assertRaises(TypeError, count, 2, 3, 4)
331 self.assertRaises(TypeError, count, 'a')
332 self.assertEqual(list(islice(count(maxsize-5), 10)), range(maxsize-5, maxsize+5))
333 self.assertEqual(list(islice(count(-maxsize-5), 10)), range(-maxsize-5, -maxsize+5))
334 c = count(3)
335 self.assertEqual(repr(c), 'count(3)')
336 c.next()
337 self.assertEqual(repr(c), 'count(4)')
338 c = count(-9)
339 self.assertEqual(repr(c), 'count(-9)')
340 c.next()
341 self.assertEqual(repr(count(10.25)), 'count(10.25)')
342 self.assertEqual(c.next(), -8)
343 for i in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 10, sys.maxint-5, sys.maxint+5):
344 # Test repr (ignoring the L in longs)
345 r1 = repr(count(i)).replace('L', '')
346 r2 = 'count(%r)'.__mod__(i).replace('L', '')
347 self.assertEqual(r1, r2)
349 def test_count_with_stride(self):
350 self.assertEqual(zip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
351 self.assertEqual(zip('abc',count(start=2,step=3)),
352 [('a', 2), ('b', 5), ('c', 8)])
353 self.assertEqual(zip('abc',count(step=-1)),
354 [('a', 0), ('b', -1), ('c', -2)])
355 self.assertEqual(zip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
356 self.assertEqual(zip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
357 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
358 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
359 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
360 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
361 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
362 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
363 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
364 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
365 c = count(3, 5)
366 self.assertEqual(repr(c), 'count(3, 5)')
367 c.next()
368 self.assertEqual(repr(c), 'count(8, 5)')
369 c = count(-9, 0)
370 self.assertEqual(repr(c), 'count(-9, 0)')
371 c.next()
372 self.assertEqual(repr(c), 'count(-9, 0)')
373 c = count(-9, -3)
374 self.assertEqual(repr(c), 'count(-9, -3)')
375 c.next()
376 self.assertEqual(repr(c), 'count(-12, -3)')
377 self.assertEqual(repr(c), 'count(-12, -3)')
378 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
379 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
380 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
381 for i in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 10, sys.maxint-5, sys.maxint+5):
382 for j in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 1, 10, sys.maxint-5, sys.maxint+5):
383 # Test repr (ignoring the L in longs)
384 r1 = repr(count(i, j)).replace('L', '')
385 if j == 1:
386 r2 = ('count(%r)' % i).replace('L', '')
387 else:
388 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
389 self.assertEqual(r1, r2)
391 def test_cycle(self):
392 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
393 self.assertEqual(list(cycle('')), [])
394 self.assertRaises(TypeError, cycle)
395 self.assertRaises(TypeError, cycle, 5)
396 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
398 def test_groupby(self):
399 # Check whether it accepts arguments correctly
400 self.assertEqual([], list(groupby([])))
401 self.assertEqual([], list(groupby([], key=id)))
402 self.assertRaises(TypeError, list, groupby('abc', []))
403 self.assertRaises(TypeError, groupby, None)
404 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
406 # Check normal input
407 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
408 (2,15,22), (3,16,23), (3,17,23)]
409 dup = []
410 for k, g in groupby(s, lambda r:r[0]):
411 for elem in g:
412 self.assertEqual(k, elem[0])
413 dup.append(elem)
414 self.assertEqual(s, dup)
416 # Check nested case
417 dup = []
418 for k, g in groupby(s, lambda r:r[0]):
419 for ik, ig in groupby(g, lambda r:r[2]):
420 for elem in ig:
421 self.assertEqual(k, elem[0])
422 self.assertEqual(ik, elem[2])
423 dup.append(elem)
424 self.assertEqual(s, dup)
426 # Check case where inner iterator is not used
427 keys = [k for k, g in groupby(s, lambda r:r[0])]
428 expectedkeys = set([r[0] for r in s])
429 self.assertEqual(set(keys), expectedkeys)
430 self.assertEqual(len(keys), len(expectedkeys))
432 # Exercise pipes and filters style
433 s = 'abracadabra'
434 # sort s | uniq
435 r = [k for k, g in groupby(sorted(s))]
436 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
437 # sort s | uniq -d
438 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
439 self.assertEqual(r, ['a', 'b', 'r'])
440 # sort s | uniq -c
441 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
442 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
443 # sort s | uniq -c | sort -rn | head -3
444 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
445 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
447 # iter.next failure
448 class ExpectedError(Exception):
449 pass
450 def delayed_raise(n=0):
451 for i in range(n):
452 yield 'yo'
453 raise ExpectedError
454 def gulp(iterable, keyp=None, func=list):
455 return [func(g) for k, g in groupby(iterable, keyp)]
457 # iter.next failure on outer object
458 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
459 # iter.next failure on inner object
460 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
462 # __cmp__ failure
463 class DummyCmp:
464 def __cmp__(self, dst):
465 raise ExpectedError
466 s = [DummyCmp(), DummyCmp(), None]
468 # __cmp__ failure on outer object
469 self.assertRaises(ExpectedError, gulp, s, func=id)
470 # __cmp__ failure on inner object
471 self.assertRaises(ExpectedError, gulp, s)
473 # keyfunc failure
474 def keyfunc(obj):
475 if keyfunc.skip > 0:
476 keyfunc.skip -= 1
477 return obj
478 else:
479 raise ExpectedError
481 # keyfunc failure on outer object
482 keyfunc.skip = 0
483 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
484 keyfunc.skip = 1
485 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
487 def test_ifilter(self):
488 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
489 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
490 self.assertEqual(list(ifilter(bool, [0,1,0,2,0])), [1,2])
491 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
492 self.assertRaises(TypeError, ifilter)
493 self.assertRaises(TypeError, ifilter, lambda x:x)
494 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
495 self.assertRaises(TypeError, ifilter, isEven, 3)
496 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
498 def test_ifilterfalse(self):
499 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
500 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
501 self.assertEqual(list(ifilterfalse(bool, [0,1,0,2,0])), [0,0,0])
502 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
503 self.assertRaises(TypeError, ifilterfalse)
504 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
505 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
506 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
507 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
509 def test_izip(self):
510 ans = [(x,y) for x, y in izip('abc',count())]
511 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
512 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
513 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
514 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
515 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
516 self.assertEqual(list(izip()), zip())
517 self.assertRaises(TypeError, izip, 3)
518 self.assertRaises(TypeError, izip, range(3), 3)
519 # Check tuple re-use (implementation detail)
520 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
521 zip('abc', 'def'))
522 self.assertEqual([pair for pair in izip('abc', 'def')],
523 zip('abc', 'def'))
524 ids = map(id, izip('abc', 'def'))
525 self.assertEqual(min(ids), max(ids))
526 ids = map(id, list(izip('abc', 'def')))
527 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
529 def test_iziplongest(self):
530 for args in [
531 ['abc', range(6)],
532 [range(6), 'abc'],
533 [range(1000), range(2000,2100), range(3000,3050)],
534 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
535 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
537 target = map(None, *args)
538 self.assertEqual(list(izip_longest(*args)), target)
539 self.assertEqual(list(izip_longest(*args, **{})), target)
540 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
541 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
543 self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
545 self.assertEqual(list(izip_longest()), zip())
546 self.assertEqual(list(izip_longest([])), zip([]))
547 self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
549 self.assertEqual(list(izip_longest('abc', 'defg', **{})), map(None, 'abc', 'defg')) # empty keyword dict
550 self.assertRaises(TypeError, izip_longest, 3)
551 self.assertRaises(TypeError, izip_longest, range(3), 3)
553 for stmt in [
554 "izip_longest('abc', fv=1)",
555 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
557 try:
558 eval(stmt, globals(), locals())
559 except TypeError:
560 pass
561 else:
562 self.fail('Did not raise Type in: ' + stmt)
564 # Check tuple re-use (implementation detail)
565 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
566 zip('abc', 'def'))
567 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
568 zip('abc', 'def'))
569 ids = map(id, izip_longest('abc', 'def'))
570 self.assertEqual(min(ids), max(ids))
571 ids = map(id, list(izip_longest('abc', 'def')))
572 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
574 def test_bug_7244(self):
576 class Repeater(object):
577 # this class is similar to itertools.repeat
578 def __init__(self, o, t, e):
579 self.o = o
580 self.t = int(t)
581 self.e = e
582 def __iter__(self): # its iterator is itself
583 return self
584 def next(self):
585 if self.t > 0:
586 self.t -= 1
587 return self.o
588 else:
589 raise self.e
591 # Formerly this code in would fail in debug mode
592 # with Undetected Error and Stop Iteration
593 r1 = Repeater(1, 3, StopIteration)
594 r2 = Repeater(2, 4, StopIteration)
595 def run(r1, r2):
596 result = []
597 for i, j in izip_longest(r1, r2, fillvalue=0):
598 with test_support.captured_output('stdout'):
599 print (i, j)
600 result.append((i, j))
601 return result
602 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
604 # Formerly, the RuntimeError would be lost
605 # and StopIteration would stop as expected
606 r1 = Repeater(1, 3, RuntimeError)
607 r2 = Repeater(2, 4, StopIteration)
608 it = izip_longest(r1, r2, fillvalue=0)
609 self.assertEqual(next(it), (1, 2))
610 self.assertEqual(next(it), (1, 2))
611 self.assertEqual(next(it), (1, 2))
612 self.assertRaises(RuntimeError, next, it)
614 def test_product(self):
615 for args, result in [
616 ([], [()]), # zero iterables
617 (['ab'], [('a',), ('b',)]), # one iterable
618 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
619 ([range(0), range(2), range(3)], []), # first iterable with zero length
620 ([range(2), range(0), range(3)], []), # middle iterable with zero length
621 ([range(2), range(3), range(0)], []), # last iterable with zero length
623 self.assertEqual(list(product(*args)), result)
624 for r in range(4):
625 self.assertEqual(list(product(*(args*r))),
626 list(product(*args, **dict(repeat=r))))
627 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
628 self.assertRaises(TypeError, product, range(6), None)
630 def product1(*args, **kwds):
631 pools = map(tuple, args) * kwds.get('repeat', 1)
632 n = len(pools)
633 if n == 0:
634 yield ()
635 return
636 if any(len(pool) == 0 for pool in pools):
637 return
638 indices = [0] * n
639 yield tuple(pool[i] for pool, i in zip(pools, indices))
640 while 1:
641 for i in reversed(range(n)): # right to left
642 if indices[i] == len(pools[i]) - 1:
643 continue
644 indices[i] += 1
645 for j in range(i+1, n):
646 indices[j] = 0
647 yield tuple(pool[i] for pool, i in zip(pools, indices))
648 break
649 else:
650 return
652 def product2(*args, **kwds):
653 'Pure python version used in docs'
654 pools = map(tuple, args) * kwds.get('repeat', 1)
655 result = [[]]
656 for pool in pools:
657 result = [x+[y] for x in result for y in pool]
658 for prod in result:
659 yield tuple(prod)
661 argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
662 set('abcdefg'), range(11), tuple(range(13))]
663 for i in range(100):
664 args = [random.choice(argtypes) for j in range(random.randrange(5))]
665 expected_len = prod(map(len, args))
666 self.assertEqual(len(list(product(*args))), expected_len)
667 self.assertEqual(list(product(*args)), list(product1(*args)))
668 self.assertEqual(list(product(*args)), list(product2(*args)))
669 args = map(iter, args)
670 self.assertEqual(len(list(product(*args))), expected_len)
672 # Test implementation detail: tuple re-use
673 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
674 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
676 def test_repeat(self):
677 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
678 self.assertEqual(zip(xrange(3),repeat('a')),
679 [(0, 'a'), (1, 'a'), (2, 'a')])
680 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
681 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
682 self.assertEqual(list(repeat('a', 0)), [])
683 self.assertEqual(list(repeat('a', -3)), [])
684 self.assertRaises(TypeError, repeat)
685 self.assertRaises(TypeError, repeat, None, 3, 4)
686 self.assertRaises(TypeError, repeat, None, 'a')
687 r = repeat(1+0j)
688 self.assertEqual(repr(r), 'repeat((1+0j))')
689 r = repeat(1+0j, 5)
690 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
691 list(r)
692 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
694 def test_imap(self):
695 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
696 [0**1, 1**2, 2**3])
697 self.assertEqual(list(imap(None, 'abc', range(5))),
698 [('a',0),('b',1),('c',2)])
699 self.assertEqual(list(imap(None, 'abc', count())),
700 [('a',0),('b',1),('c',2)])
701 self.assertEqual(take(2,imap(None, 'abc', count())),
702 [('a',0),('b',1)])
703 self.assertEqual(list(imap(operator.pow, [])), [])
704 self.assertRaises(TypeError, imap)
705 self.assertRaises(TypeError, imap, operator.neg)
706 self.assertRaises(TypeError, imap(10, range(5)).next)
707 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
708 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
710 def test_starmap(self):
711 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
712 [0**1, 1**2, 2**3])
713 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
714 [0**1, 1**2, 2**3])
715 self.assertEqual(list(starmap(operator.pow, [])), [])
716 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
717 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
718 self.assertRaises(TypeError, starmap)
719 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
720 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
721 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
722 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
724 def test_islice(self):
725 for args in [ # islice(args) should agree with range(args)
726 (10, 20, 3),
727 (10, 3, 20),
728 (10, 20),
729 (10, 3),
730 (20,)
732 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
734 for args, tgtargs in [ # Stop when seqn is exhausted
735 ((10, 110, 3), ((10, 100, 3))),
736 ((10, 110), ((10, 100))),
737 ((110,), (100,))
739 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
741 # Test stop=None
742 self.assertEqual(list(islice(xrange(10), None)), range(10))
743 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
744 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
745 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
746 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
748 # Test number of items consumed SF #1171417
749 it = iter(range(10))
750 self.assertEqual(list(islice(it, 3)), range(3))
751 self.assertEqual(list(it), range(3, 10))
753 # Test invalid arguments
754 self.assertRaises(TypeError, islice, xrange(10))
755 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
756 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
757 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
758 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
759 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
760 self.assertRaises(ValueError, islice, xrange(10), 'a')
761 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
762 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
763 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
764 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
765 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
767 def test_takewhile(self):
768 data = [1, 3, 5, 20, 2, 4, 6, 8]
769 underten = lambda x: x<10
770 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
771 self.assertEqual(list(takewhile(underten, [])), [])
772 self.assertRaises(TypeError, takewhile)
773 self.assertRaises(TypeError, takewhile, operator.pow)
774 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
775 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
776 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
777 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
778 self.assertEqual(list(t), [1, 1, 1])
779 self.assertRaises(StopIteration, t.next)
781 def test_dropwhile(self):
782 data = [1, 3, 5, 20, 2, 4, 6, 8]
783 underten = lambda x: x<10
784 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
785 self.assertEqual(list(dropwhile(underten, [])), [])
786 self.assertRaises(TypeError, dropwhile)
787 self.assertRaises(TypeError, dropwhile, operator.pow)
788 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
789 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
790 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
792 def test_tee(self):
793 n = 200
794 def irange(n):
795 for i in xrange(n):
796 yield i
798 a, b = tee([]) # test empty iterator
799 self.assertEqual(list(a), [])
800 self.assertEqual(list(b), [])
802 a, b = tee(irange(n)) # test 100% interleaved
803 self.assertEqual(zip(a,b), zip(range(n),range(n)))
805 a, b = tee(irange(n)) # test 0% interleaved
806 self.assertEqual(list(a), range(n))
807 self.assertEqual(list(b), range(n))
809 a, b = tee(irange(n)) # test dealloc of leading iterator
810 for i in xrange(100):
811 self.assertEqual(a.next(), i)
812 del a
813 self.assertEqual(list(b), range(n))
815 a, b = tee(irange(n)) # test dealloc of trailing iterator
816 for i in xrange(100):
817 self.assertEqual(a.next(), i)
818 del b
819 self.assertEqual(list(a), range(100, n))
821 for j in xrange(5): # test randomly interleaved
822 order = [0]*n + [1]*n
823 random.shuffle(order)
824 lists = ([], [])
825 its = tee(irange(n))
826 for i in order:
827 value = its[i].next()
828 lists[i].append(value)
829 self.assertEqual(lists[0], range(n))
830 self.assertEqual(lists[1], range(n))
832 # test argument format checking
833 self.assertRaises(TypeError, tee)
834 self.assertRaises(TypeError, tee, 3)
835 self.assertRaises(TypeError, tee, [1,2], 'x')
836 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
838 # tee object should be instantiable
839 a, b = tee('abc')
840 c = type(a)('def')
841 self.assertEqual(list(c), list('def'))
843 # test long-lagged and multi-way split
844 a, b, c = tee(xrange(2000), 3)
845 for i in xrange(100):
846 self.assertEqual(a.next(), i)
847 self.assertEqual(list(b), range(2000))
848 self.assertEqual([c.next(), c.next()], range(2))
849 self.assertEqual(list(a), range(100,2000))
850 self.assertEqual(list(c), range(2,2000))
852 # test values of n
853 self.assertRaises(TypeError, tee, 'abc', 'invalid')
854 self.assertRaises(ValueError, tee, [], -1)
855 for n in xrange(5):
856 result = tee('abc', n)
857 self.assertEqual(type(result), tuple)
858 self.assertEqual(len(result), n)
859 self.assertEqual(map(list, result), [list('abc')]*n)
861 # tee pass-through to copyable iterator
862 a, b = tee('abc')
863 c, d = tee(a)
864 self.assertTrue(a is c)
866 # test tee_new
867 t1, t2 = tee('abc')
868 tnew = type(t1)
869 self.assertRaises(TypeError, tnew)
870 self.assertRaises(TypeError, tnew, 10)
871 t3 = tnew(t1)
872 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
874 # test that tee objects are weak referencable
875 a, b = tee(xrange(10))
876 p = proxy(a)
877 self.assertEqual(getattr(p, '__class__'), type(b))
878 del a
879 self.assertRaises(ReferenceError, getattr, p, '__class__')
881 def test_StopIteration(self):
882 self.assertRaises(StopIteration, izip().next)
884 for f in (chain, cycle, izip, groupby):
885 self.assertRaises(StopIteration, f([]).next)
886 self.assertRaises(StopIteration, f(StopNow()).next)
888 self.assertRaises(StopIteration, islice([], None).next)
889 self.assertRaises(StopIteration, islice(StopNow(), None).next)
891 p, q = tee([])
892 self.assertRaises(StopIteration, p.next)
893 self.assertRaises(StopIteration, q.next)
894 p, q = tee(StopNow())
895 self.assertRaises(StopIteration, p.next)
896 self.assertRaises(StopIteration, q.next)
898 self.assertRaises(StopIteration, repeat(None, 0).next)
900 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
901 self.assertRaises(StopIteration, f(lambda x:x, []).next)
902 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
904 class TestExamples(unittest.TestCase):
906 def test_chain(self):
907 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
909 def test_chain_from_iterable(self):
910 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
912 def test_combinations(self):
913 self.assertEqual(list(combinations('ABCD', 2)),
914 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
915 self.assertEqual(list(combinations(range(4), 3)),
916 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
918 def test_combinations_with_replacement(self):
919 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
920 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
922 def test_compress(self):
923 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
925 def test_count(self):
926 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
928 def test_cycle(self):
929 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
931 def test_dropwhile(self):
932 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
934 def test_groupby(self):
935 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
936 list('ABCDAB'))
937 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
938 [list('AAAA'), list('BBB'), list('CC'), list('D')])
940 def test_ifilter(self):
941 self.assertEqual(list(ifilter(lambda x: x%2, range(10))), [1,3,5,7,9])
943 def test_ifilterfalse(self):
944 self.assertEqual(list(ifilterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
946 def test_imap(self):
947 self.assertEqual(list(imap(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
949 def test_islice(self):
950 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
951 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
952 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
953 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
955 def test_izip(self):
956 self.assertEqual(list(izip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
958 def test_izip_longest(self):
959 self.assertEqual(list(izip_longest('ABCD', 'xy', fillvalue='-')),
960 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
962 def test_permutations(self):
963 self.assertEqual(list(permutations('ABCD', 2)),
964 map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split()))
965 self.assertEqual(list(permutations(range(3))),
966 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
968 def test_product(self):
969 self.assertEqual(list(product('ABCD', 'xy')),
970 map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split()))
971 self.assertEqual(list(product(range(2), repeat=3)),
972 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
973 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
975 def test_repeat(self):
976 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
978 def test_stapmap(self):
979 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
980 [32, 9, 1000])
982 def test_takewhile(self):
983 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
986 class TestGC(unittest.TestCase):
988 def makecycle(self, iterator, container):
989 container.append(iterator)
990 iterator.next()
991 del container, iterator
993 def test_chain(self):
994 a = []
995 self.makecycle(chain(a), a)
997 def test_chain_from_iterable(self):
998 a = []
999 self.makecycle(chain.from_iterable([a]), a)
1001 def test_combinations(self):
1002 a = []
1003 self.makecycle(combinations([1,2,a,3], 3), a)
1005 def test_combinations_with_replacement(self):
1006 a = []
1007 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1009 def test_compress(self):
1010 a = []
1011 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1013 def test_count(self):
1014 a = []
1015 Int = type('Int', (int,), dict(x=a))
1016 self.makecycle(count(Int(0), Int(1)), a)
1018 def test_cycle(self):
1019 a = []
1020 self.makecycle(cycle([a]*2), a)
1022 def test_dropwhile(self):
1023 a = []
1024 self.makecycle(dropwhile(bool, [0, a, a]), a)
1026 def test_groupby(self):
1027 a = []
1028 self.makecycle(groupby([a]*2, lambda x:x), a)
1030 def test_issue2246(self):
1031 # Issue 2246 -- the _grouper iterator was not included in GC
1032 n = 10
1033 keyfunc = lambda x: x
1034 for i, j in groupby(xrange(n), key=keyfunc):
1035 keyfunc.__dict__.setdefault('x',[]).append(j)
1037 def test_ifilter(self):
1038 a = []
1039 self.makecycle(ifilter(lambda x:True, [a]*2), a)
1041 def test_ifilterfalse(self):
1042 a = []
1043 self.makecycle(ifilterfalse(lambda x:False, a), a)
1045 def test_izip(self):
1046 a = []
1047 self.makecycle(izip([a]*2, [a]*3), a)
1049 def test_izip_longest(self):
1050 a = []
1051 self.makecycle(izip_longest([a]*2, [a]*3), a)
1052 b = [a, None]
1053 self.makecycle(izip_longest([a]*2, [a]*3, fillvalue=b), a)
1055 def test_imap(self):
1056 a = []
1057 self.makecycle(imap(lambda x:x, [a]*2), a)
1059 def test_islice(self):
1060 a = []
1061 self.makecycle(islice([a]*2, None), a)
1063 def test_permutations(self):
1064 a = []
1065 self.makecycle(permutations([1,2,a,3], 3), a)
1067 def test_product(self):
1068 a = []
1069 self.makecycle(product([1,2,a,3], repeat=3), a)
1071 def test_repeat(self):
1072 a = []
1073 self.makecycle(repeat(a), a)
1075 def test_starmap(self):
1076 a = []
1077 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1079 def test_takewhile(self):
1080 a = []
1081 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1083 def R(seqn):
1084 'Regular generator'
1085 for i in seqn:
1086 yield i
1088 class G:
1089 'Sequence using __getitem__'
1090 def __init__(self, seqn):
1091 self.seqn = seqn
1092 def __getitem__(self, i):
1093 return self.seqn[i]
1095 class I:
1096 'Sequence using iterator protocol'
1097 def __init__(self, seqn):
1098 self.seqn = seqn
1099 self.i = 0
1100 def __iter__(self):
1101 return self
1102 def next(self):
1103 if self.i >= len(self.seqn): raise StopIteration
1104 v = self.seqn[self.i]
1105 self.i += 1
1106 return v
1108 class Ig:
1109 'Sequence using iterator protocol defined with a generator'
1110 def __init__(self, seqn):
1111 self.seqn = seqn
1112 self.i = 0
1113 def __iter__(self):
1114 for val in self.seqn:
1115 yield val
1117 class X:
1118 'Missing __getitem__ and __iter__'
1119 def __init__(self, seqn):
1120 self.seqn = seqn
1121 self.i = 0
1122 def next(self):
1123 if self.i >= len(self.seqn): raise StopIteration
1124 v = self.seqn[self.i]
1125 self.i += 1
1126 return v
1128 class N:
1129 'Iterator missing next()'
1130 def __init__(self, seqn):
1131 self.seqn = seqn
1132 self.i = 0
1133 def __iter__(self):
1134 return self
1136 class E:
1137 'Test propagation of exceptions'
1138 def __init__(self, seqn):
1139 self.seqn = seqn
1140 self.i = 0
1141 def __iter__(self):
1142 return self
1143 def next(self):
1144 3 // 0
1146 class S:
1147 'Test immediate stop'
1148 def __init__(self, seqn):
1149 pass
1150 def __iter__(self):
1151 return self
1152 def next(self):
1153 raise StopIteration
1155 def L(seqn):
1156 'Test multiple tiers of iterators'
1157 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
1160 class TestVariousIteratorArgs(unittest.TestCase):
1162 def test_chain(self):
1163 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1164 for g in (G, I, Ig, S, L, R):
1165 self.assertEqual(list(chain(g(s))), list(g(s)))
1166 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
1167 self.assertRaises(TypeError, list, chain(X(s)))
1168 self.assertRaises(TypeError, list, chain(N(s)))
1169 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1171 def test_compress(self):
1172 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1173 n = len(s)
1174 for g in (G, I, Ig, S, L, R):
1175 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1176 self.assertRaises(TypeError, compress, X(s), repeat(1))
1177 self.assertRaises(TypeError, list, compress(N(s), repeat(1)))
1178 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1180 def test_product(self):
1181 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1182 self.assertRaises(TypeError, product, X(s))
1183 self.assertRaises(TypeError, product, N(s))
1184 self.assertRaises(ZeroDivisionError, product, E(s))
1186 def test_cycle(self):
1187 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1188 for g in (G, I, Ig, S, L, R):
1189 tgtlen = len(s) * 3
1190 expected = list(g(s))*3
1191 actual = list(islice(cycle(g(s)), tgtlen))
1192 self.assertEqual(actual, expected)
1193 self.assertRaises(TypeError, cycle, X(s))
1194 self.assertRaises(TypeError, list, cycle(N(s)))
1195 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1197 def test_groupby(self):
1198 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1199 for g in (G, I, Ig, S, L, R):
1200 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1201 self.assertRaises(TypeError, groupby, X(s))
1202 self.assertRaises(TypeError, list, groupby(N(s)))
1203 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1205 def test_ifilter(self):
1206 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1207 for g in (G, I, Ig, S, L, R):
1208 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
1209 self.assertRaises(TypeError, ifilter, isEven, X(s))
1210 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
1211 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
1213 def test_ifilterfalse(self):
1214 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1215 for g in (G, I, Ig, S, L, R):
1216 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
1217 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
1218 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
1219 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
1221 def test_izip(self):
1222 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1223 for g in (G, I, Ig, S, L, R):
1224 self.assertEqual(list(izip(g(s))), zip(g(s)))
1225 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
1226 self.assertRaises(TypeError, izip, X(s))
1227 self.assertRaises(TypeError, list, izip(N(s)))
1228 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
1230 def test_iziplongest(self):
1231 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1232 for g in (G, I, Ig, S, L, R):
1233 self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
1234 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
1235 self.assertRaises(TypeError, izip_longest, X(s))
1236 self.assertRaises(TypeError, list, izip_longest(N(s)))
1237 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
1239 def test_imap(self):
1240 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1241 for g in (G, I, Ig, S, L, R):
1242 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
1243 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
1244 self.assertRaises(TypeError, imap, onearg, X(s))
1245 self.assertRaises(TypeError, list, imap(onearg, N(s)))
1246 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
1248 def test_islice(self):
1249 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1250 for g in (G, I, Ig, S, L, R):
1251 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1252 self.assertRaises(TypeError, islice, X(s), 10)
1253 self.assertRaises(TypeError, list, islice(N(s), 10))
1254 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1256 def test_starmap(self):
1257 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1258 for g in (G, I, Ig, S, L, R):
1259 ss = zip(s, s)
1260 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
1261 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
1262 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
1263 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1265 def test_takewhile(self):
1266 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1267 for g in (G, I, Ig, S, L, R):
1268 tgt = []
1269 for elem in g(s):
1270 if not isEven(elem): break
1271 tgt.append(elem)
1272 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1273 self.assertRaises(TypeError, takewhile, isEven, X(s))
1274 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
1275 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1277 def test_dropwhile(self):
1278 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1279 for g in (G, I, Ig, S, L, R):
1280 tgt = []
1281 for elem in g(s):
1282 if not tgt and isOdd(elem): continue
1283 tgt.append(elem)
1284 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1285 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
1286 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
1287 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1289 def test_tee(self):
1290 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1291 for g in (G, I, Ig, S, L, R):
1292 it1, it2 = tee(g(s))
1293 self.assertEqual(list(it1), list(g(s)))
1294 self.assertEqual(list(it2), list(g(s)))
1295 self.assertRaises(TypeError, tee, X(s))
1296 self.assertRaises(TypeError, list, tee(N(s))[0])
1297 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1299 class LengthTransparency(unittest.TestCase):
1301 def test_repeat(self):
1302 from test.test_iterlen import len
1303 self.assertEqual(len(repeat(None, 50)), 50)
1304 self.assertRaises(TypeError, len, repeat(None))
1306 class RegressionTests(unittest.TestCase):
1308 def test_sf_793826(self):
1309 # Fix Armin Rigo's successful efforts to wreak havoc
1311 def mutatingtuple(tuple1, f, tuple2):
1312 # this builds a tuple t which is a copy of tuple1,
1313 # then calls f(t), then mutates t to be equal to tuple2
1314 # (needs len(tuple1) == len(tuple2)).
1315 def g(value, first=[1]):
1316 if first:
1317 del first[:]
1318 f(z.next())
1319 return value
1320 items = list(tuple2)
1321 items[1:1] = list(tuple1)
1322 gen = imap(g, items)
1323 z = izip(*[gen]*len(tuple1))
1324 z.next()
1326 def f(t):
1327 global T
1328 T = t
1329 first[:] = list(T)
1331 first = []
1332 mutatingtuple((1,2,3), f, (4,5,6))
1333 second = list(T)
1334 self.assertEqual(first, second)
1337 def test_sf_950057(self):
1338 # Make sure that chain() and cycle() catch exceptions immediately
1339 # rather than when shifting between input sources
1341 def gen1():
1342 hist.append(0)
1343 yield 1
1344 hist.append(1)
1345 raise AssertionError
1346 hist.append(2)
1348 def gen2(x):
1349 hist.append(3)
1350 yield 2
1351 hist.append(4)
1352 if x:
1353 raise StopIteration
1355 hist = []
1356 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1357 self.assertEqual(hist, [0,1])
1359 hist = []
1360 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1361 self.assertEqual(hist, [0,1])
1363 hist = []
1364 self.assertRaises(AssertionError, list, cycle(gen1()))
1365 self.assertEqual(hist, [0,1])
1367 class SubclassWithKwargsTest(unittest.TestCase):
1368 def test_keywords_in_subclass(self):
1369 # count is not subclassable...
1370 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
1371 starmap, islice, takewhile, dropwhile, cycle, compress):
1372 class Subclass(cls):
1373 def __init__(self, newarg=None, *args):
1374 cls.__init__(self, *args)
1375 try:
1376 Subclass(newarg=1)
1377 except TypeError, err:
1378 # we expect type errors because of wrong argument count
1379 self.assertFalse("does not take keyword arguments" in err.args[0])
1382 libreftest = """ Doctest for examples in the library reference: libitertools.tex
1385 >>> amounts = [120.15, 764.05, 823.14]
1386 >>> for checknum, amount in izip(count(1200), amounts):
1387 ... print 'Check %d is for $%.2f' % (checknum, amount)
1389 Check 1200 is for $120.15
1390 Check 1201 is for $764.05
1391 Check 1202 is for $823.14
1393 >>> import operator
1394 >>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1395 ... print cube
1401 >>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
1402 >>> for name in islice(reportlines, 3, None, 2):
1403 ... print name.title()
1405 Alex
1406 Laura
1407 Martin
1408 Walter
1409 Samuele
1411 >>> from operator import itemgetter
1412 >>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
1413 >>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
1414 >>> for k, g in groupby(di, itemgetter(1)):
1415 ... print k, map(itemgetter(0), g)
1417 1 ['a', 'c', 'e']
1418 2 ['b', 'd', 'f']
1419 3 ['g']
1421 # Find runs of consecutive numbers using groupby. The key to the solution
1422 # is differencing with a range so that consecutive numbers all appear in
1423 # same group.
1424 >>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
1425 >>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
1426 ... print map(operator.itemgetter(1), g)
1429 [4, 5, 6]
1430 [10]
1431 [15, 16, 17, 18]
1432 [22]
1433 [25, 26, 27, 28]
1435 >>> def take(n, iterable):
1436 ... "Return first n items of the iterable as a list"
1437 ... return list(islice(iterable, n))
1439 >>> def enumerate(iterable, start=0):
1440 ... return izip(count(start), iterable)
1442 >>> def tabulate(function, start=0):
1443 ... "Return function(0), function(1), ..."
1444 ... return imap(function, count(start))
1446 >>> def nth(iterable, n, default=None):
1447 ... "Returns the nth item or a default value"
1448 ... return next(islice(iterable, n, None), default)
1450 >>> def quantify(iterable, pred=bool):
1451 ... "Count how many times the predicate is true"
1452 ... return sum(imap(pred, iterable))
1454 >>> def padnone(iterable):
1455 ... "Returns the sequence elements and then returns None indefinitely"
1456 ... return chain(iterable, repeat(None))
1458 >>> def ncycles(iterable, n):
1459 ... "Returns the seqeuence elements n times"
1460 ... return chain(*repeat(iterable, n))
1462 >>> def dotproduct(vec1, vec2):
1463 ... return sum(imap(operator.mul, vec1, vec2))
1465 >>> def flatten(listOfLists):
1466 ... return list(chain.from_iterable(listOfLists))
1468 >>> def repeatfunc(func, times=None, *args):
1469 ... "Repeat calls to func with specified arguments."
1470 ... " Example: repeatfunc(random.random)"
1471 ... if times is None:
1472 ... return starmap(func, repeat(args))
1473 ... else:
1474 ... return starmap(func, repeat(args, times))
1476 >>> def pairwise(iterable):
1477 ... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1478 ... a, b = tee(iterable)
1479 ... for elem in b:
1480 ... break
1481 ... return izip(a, b)
1483 >>> def grouper(n, iterable, fillvalue=None):
1484 ... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
1485 ... args = [iter(iterable)] * n
1486 ... return izip_longest(fillvalue=fillvalue, *args)
1488 >>> def roundrobin(*iterables):
1489 ... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
1490 ... # Recipe credited to George Sakkis
1491 ... pending = len(iterables)
1492 ... nexts = cycle(iter(it).next for it in iterables)
1493 ... while pending:
1494 ... try:
1495 ... for next in nexts:
1496 ... yield next()
1497 ... except StopIteration:
1498 ... pending -= 1
1499 ... nexts = cycle(islice(nexts, pending))
1501 >>> def powerset(iterable):
1502 ... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1503 ... s = list(iterable)
1504 ... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
1506 >>> def unique_everseen(iterable, key=None):
1507 ... "List unique elements, preserving order. Remember all elements ever seen."
1508 ... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1509 ... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1510 ... seen = set()
1511 ... seen_add = seen.add
1512 ... if key is None:
1513 ... for element in iterable:
1514 ... if element not in seen:
1515 ... seen_add(element)
1516 ... yield element
1517 ... else:
1518 ... for element in iterable:
1519 ... k = key(element)
1520 ... if k not in seen:
1521 ... seen_add(k)
1522 ... yield element
1524 >>> def unique_justseen(iterable, key=None):
1525 ... "List unique elements, preserving order. Remember only the element just seen."
1526 ... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1527 ... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1528 ... return imap(next, imap(itemgetter(1), groupby(iterable, key)))
1530 This is not part of the examples but it tests to make sure the definitions
1531 perform as purported.
1533 >>> take(10, count())
1534 [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1536 >>> list(enumerate('abc'))
1537 [(0, 'a'), (1, 'b'), (2, 'c')]
1539 >>> list(islice(tabulate(lambda x: 2*x), 4))
1540 [0, 2, 4, 6]
1542 >>> nth('abcde', 3)
1545 >>> nth('abcde', 9) is None
1546 True
1548 >>> quantify(xrange(99), lambda x: x%2==0)
1551 >>> a = [[1, 2, 3], [4, 5, 6]]
1552 >>> flatten(a)
1553 [1, 2, 3, 4, 5, 6]
1555 >>> list(repeatfunc(pow, 5, 2, 3))
1556 [8, 8, 8, 8, 8]
1558 >>> import random
1559 >>> take(5, imap(int, repeatfunc(random.random)))
1560 [0, 0, 0, 0, 0]
1562 >>> list(pairwise('abcd'))
1563 [('a', 'b'), ('b', 'c'), ('c', 'd')]
1565 >>> list(pairwise([]))
1568 >>> list(pairwise('a'))
1571 >>> list(islice(padnone('abc'), 0, 6))
1572 ['a', 'b', 'c', None, None, None]
1574 >>> list(ncycles('abc', 3))
1575 ['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1577 >>> dotproduct([1,2,3], [4,5,6])
1580 >>> list(grouper(3, 'abcdefg', 'x'))
1581 [('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1583 >>> list(roundrobin('abc', 'd', 'ef'))
1584 ['a', 'd', 'e', 'b', 'f', 'c']
1586 >>> list(powerset([1,2,3]))
1587 [(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
1589 >>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1590 True
1592 >>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1593 True
1595 >>> list(unique_everseen('AAAABBBCCDAABBB'))
1596 ['A', 'B', 'C', 'D']
1598 >>> list(unique_everseen('ABBCcAD', str.lower))
1599 ['A', 'B', 'C', 'D']
1601 >>> list(unique_justseen('AAAABBBCCDAABBB'))
1602 ['A', 'B', 'C', 'D', 'A', 'B']
1604 >>> list(unique_justseen('ABBCcAD', str.lower))
1605 ['A', 'B', 'C', 'A', 'D']
1609 __test__ = {'libreftest' : libreftest}
1611 def test_main(verbose=None):
1612 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
1613 RegressionTests, LengthTransparency,
1614 SubclassWithKwargsTest, TestExamples)
1615 test_support.run_unittest(*test_classes)
1617 # verify reference counting
1618 if verbose and hasattr(sys, "gettotalrefcount"):
1619 import gc
1620 counts = [None] * 5
1621 for i in xrange(len(counts)):
1622 test_support.run_unittest(*test_classes)
1623 gc.collect()
1624 counts[i] = sys.gettotalrefcount()
1625 print counts
1627 # doctest the examples in the library reference
1628 test_support.run_doctest(sys.modules[__name__], verbose)
1630 if __name__ == "__main__":
1631 test_main(verbose=True)