?  search:
summary | log | graphiclog1 | graphiclog2 | commit | commitdiff | tree | refs | edit | fork
blame | history | raw | HEAD
Neaten-up the itertools recipes.
[python.git] / Lib / test / test_itertools.py
blobf6174ab63d1da9f2af7da3c03a7987a2c1cbec79
1 import unittest
2 from test import test_support
3 from itertools import *
4 from weakref import proxy
5 import sys
6 import operator
7 import random
8 maxsize = test_support.MAX_Py_ssize_t
9 minsize = -maxsize-1
10
11 def onearg(x):
12 'Test function of one argument'
13 return 2*x
14
15 def errfunc(*args):
16 'Test function that raises an error'
17 raise ValueError
18
19 def gen3():
20 'Non-restartable source sequence'
21 for i in (0, 1, 2):
22 yield i
23
24 def isEven(x):
25 'Test predicate'
26 return x%2==0
27
28 def isOdd(x):
29 'Test predicate'
30 return x%2==1
31
32 class StopNow:
33 'Class emulating an empty iterable.'
34 def __iter__(self):
35 return self
36 def next(self):
37 raise StopIteration
38
39 def take(n, seq):
40 'Convenience function for partially consuming a long of infinite iterable'
41 return list(islice(seq, n))
42
43 def prod(iterable):
44 return reduce(operator.mul, iterable, 1)
45
46 def fact(n):
47 'Factorial'
48 return prod(range(1, n+1))
49
50 class TestBasicOps(unittest.TestCase):
51 def test_chain(self):
52
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))
65
66 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
73 def test_combinations(self):
74 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
75 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
76 self.assertRaises(TypeError, combinations, None) # pool is not iterable
77 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
78 self.assertRaises(ValueError, combinations, 'abc', 32) # r is too big
79 self.assertEqual(list(combinations(range(4), 3)),
80 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
81
82 def combinations1(iterable, r):
83 'Pure python version shown in the docs'
84 pool = tuple(iterable)
85 n = len(pool)
86 indices = range(r)
87 yield tuple(pool[i] for i in indices)
88 while 1:
89 for i in reversed(range(r)):
90 if indices[i] != i + n - r:
91 break
92 else:
93 return
94 indices[i] += 1
95 for j in range(i+1, r):
96 indices[j] = indices[j-1] + 1
97 yield tuple(pool[i] for i in indices)
98
99 def combinations2(iterable, r):
100 'Pure python version shown in the docs'
101 pool = tuple(iterable)
102 n = len(pool)
103 for indices in permutations(range(n), r):
104 if sorted(indices) == list(indices):
105 yield tuple(pool[i] for i in indices)
106
107 for n in range(7):
108 values = [5*x-12 for x in range(n)]
109 for r in range(n+1):
110 result = list(combinations(values, r))
111 self.assertEqual(len(result), fact(n) / fact(r) / fact(n-r)) # right number of combs
112 self.assertEqual(len(result), len(set(result))) # no repeats
113 self.assertEqual(result, sorted(result)) # lexicographic order
114 for c in result:
115 self.assertEqual(len(c), r) # r-length combinations
116 self.assertEqual(len(set(c)), r) # no duplicate elements
117 self.assertEqual(list(c), sorted(c)) # keep original ordering
118 self.assert_(all(e in values for e in c)) # elements taken from input iterable
119 self.assertEqual(list(c),
120 [e for e in values if e in c]) # comb is a subsequence of the input iterable
121 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
122 self.assertEqual(result, list(combinations2(values, r))) # matches first pure python version
123
124 # Test implementation detail: tuple re-use
125 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
126 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
127
128 def test_permutations(self):
129 self.assertRaises(TypeError, permutations) # too few arguments
130 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
131 self.assertRaises(TypeError, permutations, None) # pool is not iterable
132 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
133 self.assertRaises(ValueError, permutations, 'abc', 32) # r is too big
134 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
135 self.assertEqual(list(permutations(range(3), 2)),
136 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
137
138 def permutations1(iterable, r=None):
139 'Pure python version shown in the docs'
140 pool = tuple(iterable)
141 n = len(pool)
142 r = n if r is None else r
143 indices = range(n)
144 cycles = range(n, n-r, -1)
145 yield tuple(pool[i] for i in indices[:r])
146 while n:
147 for i in reversed(range(r)):
148 cycles[i] -= 1
149 if cycles[i] == 0:
150 indices[i:] = indices[i+1:] + indices[i:i+1]
151 cycles[i] = n - i
152 else:
153 j = cycles[i]
154 indices[i], indices[-j] = indices[-j], indices[i]
155 yield tuple(pool[i] for i in indices[:r])
156 break
157 else:
158 return
159
160 def permutations2(iterable, r=None):
161 'Pure python version shown in the docs'
162 pool = tuple(iterable)
163 n = len(pool)
164 r = n if r is None else r
165 for indices in product(range(n), repeat=r):
166 if len(set(indices)) == r:
167 yield tuple(pool[i] for i in indices)
168
169 for n in range(7):
170 values = [5*x-12 for x in range(n)]
171 for r in range(n+1):
172 result = list(permutations(values, r))
173 self.assertEqual(len(result), fact(n) / fact(n-r)) # right number of perms
174 self.assertEqual(len(result), len(set(result))) # no repeats
175 self.assertEqual(result, sorted(result)) # lexicographic order
176 for p in result:
177 self.assertEqual(len(p), r) # r-length permutations
178 self.assertEqual(len(set(p)), r) # no duplicate elements
179 self.assert_(all(e in values for e in p)) # elements taken from input iterable
180 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
181 self.assertEqual(result, list(permutations2(values, r))) # matches first pure python version
182 if r == n:
183 self.assertEqual(result, list(permutations(values, None))) # test r as None
184 self.assertEqual(result, list(permutations(values))) # test default r
185
186 # Test implementation detail: tuple re-use
187 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
188 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
189
190 def test_count(self):
191 self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
192 self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
193 self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
194 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
195 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
196 self.assertRaises(TypeError, count, 2, 3)
197 self.assertRaises(TypeError, count, 'a')
198 self.assertEqual(list(islice(count(maxsize-5), 10)), range(maxsize-5, maxsize+5))
199 self.assertEqual(list(islice(count(-maxsize-5), 10)), range(-maxsize-5, -maxsize+5))
200 c = count(3)
201 self.assertEqual(repr(c), 'count(3)')
202 c.next()
203 self.assertEqual(repr(c), 'count(4)')
204 c = count(-9)
205 self.assertEqual(repr(c), 'count(-9)')
206 c.next()
207 self.assertEqual(c.next(), -8)
208 for i in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 10, sys.maxint-5, sys.maxint+5):
209 # Test repr (ignoring the L in longs)
210 r1 = repr(count(i)).replace('L', '')
211 r2 = 'count(%r)'.__mod__(i).replace('L', '')
212 self.assertEqual(r1, r2)
213
214 def test_cycle(self):
215 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
216 self.assertEqual(list(cycle('')), [])
217 self.assertRaises(TypeError, cycle)
218 self.assertRaises(TypeError, cycle, 5)
219 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
220
221 def test_groupby(self):
222 # Check whether it accepts arguments correctly
223 self.assertEqual([], list(groupby([])))
224 self.assertEqual([], list(groupby([], key=id)))
225 self.assertRaises(TypeError, list, groupby('abc', []))
226 self.assertRaises(TypeError, groupby, None)
227 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
228
229 # Check normal input
230 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
231 (2,15,22), (3,16,23), (3,17,23)]
232 dup = []
233 for k, g in groupby(s, lambda r:r[0]):
234 for elem in g:
235 self.assertEqual(k, elem[0])
236 dup.append(elem)
237 self.assertEqual(s, dup)
238
239 # Check nested case
240 dup = []
241 for k, g in groupby(s, lambda r:r[0]):
242 for ik, ig in groupby(g, lambda r:r[2]):
243 for elem in ig:
244 self.assertEqual(k, elem[0])
245 self.assertEqual(ik, elem[2])
246 dup.append(elem)
247 self.assertEqual(s, dup)
248
249 # Check case where inner iterator is not used
250 keys = [k for k, g in groupby(s, lambda r:r[0])]
251 expectedkeys = set([r[0] for r in s])
252 self.assertEqual(set(keys), expectedkeys)
253 self.assertEqual(len(keys), len(expectedkeys))
254
255 # Exercise pipes and filters style
256 s = 'abracadabra'
257 # sort s | uniq
258 r = [k for k, g in groupby(sorted(s))]
259 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
260 # sort s | uniq -d
261 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
262 self.assertEqual(r, ['a', 'b', 'r'])
263 # sort s | uniq -c
264 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
265 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
266 # sort s | uniq -c | sort -rn | head -3
267 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
268 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
269
270 # iter.next failure
271 class ExpectedError(Exception):
272 pass
273 def delayed_raise(n=0):
274 for i in range(n):
275 yield 'yo'
276 raise ExpectedError
277 def gulp(iterable, keyp=None, func=list):
278 return [func(g) for k, g in groupby(iterable, keyp)]
279
280 # iter.next failure on outer object
281 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
282 # iter.next failure on inner object
283 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
284
285 # __cmp__ failure
286 class DummyCmp:
287 def __cmp__(self, dst):
288 raise ExpectedError
289 s = [DummyCmp(), DummyCmp(), None]
290
291 # __cmp__ failure on outer object
292 self.assertRaises(ExpectedError, gulp, s, func=id)
293 # __cmp__ failure on inner object
294 self.assertRaises(ExpectedError, gulp, s)
295
296 # keyfunc failure
297 def keyfunc(obj):
298 if keyfunc.skip > 0:
299 keyfunc.skip -= 1
300 return obj
301 else:
302 raise ExpectedError
303
304 # keyfunc failure on outer object
305 keyfunc.skip = 0
306 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
307 keyfunc.skip = 1
308 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
309
310 def test_ifilter(self):
311 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
312 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
313 self.assertEqual(list(ifilter(bool, [0,1,0,2,0])), [1,2])
314 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
315 self.assertRaises(TypeError, ifilter)
316 self.assertRaises(TypeError, ifilter, lambda x:x)
317 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
318 self.assertRaises(TypeError, ifilter, isEven, 3)
319 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
320
321 def test_ifilterfalse(self):
322 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
323 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
324 self.assertEqual(list(ifilterfalse(bool, [0,1,0,2,0])), [0,0,0])
325 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
326 self.assertRaises(TypeError, ifilterfalse)
327 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
328 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
329 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
330 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
331
332 def test_izip(self):
333 ans = [(x,y) for x, y in izip('abc',count())]
334 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
335 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
336 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
337 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
338 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
339 self.assertEqual(list(izip()), zip())
340 self.assertRaises(TypeError, izip, 3)
341 self.assertRaises(TypeError, izip, range(3), 3)
342 # Check tuple re-use (implementation detail)
343 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
344 zip('abc', 'def'))
345 self.assertEqual([pair for pair in izip('abc', 'def')],
346 zip('abc', 'def'))
347 ids = map(id, izip('abc', 'def'))
348 self.assertEqual(min(ids), max(ids))
349 ids = map(id, list(izip('abc', 'def')))
350 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
351
352 def test_iziplongest(self):
353 for args in [
354 ['abc', range(6)],
355 [range(6), 'abc'],
356 [range(1000), range(2000,2100), range(3000,3050)],
357 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
358 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
359 ]:
360 target = map(None, *args)
361 self.assertEqual(list(izip_longest(*args)), target)
362 self.assertEqual(list(izip_longest(*args, **{})), target)
363 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
364 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
365
366 self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
367
368 self.assertEqual(list(izip_longest()), zip())
369 self.assertEqual(list(izip_longest([])), zip([]))
370 self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
371
372 self.assertEqual(list(izip_longest('abc', 'defg', **{})), map(None, 'abc', 'defg')) # empty keyword dict
373 self.assertRaises(TypeError, izip_longest, 3)
374 self.assertRaises(TypeError, izip_longest, range(3), 3)
375
376 for stmt in [
377 "izip_longest('abc', fv=1)",
378 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
379 ]:
380 try:
381 eval(stmt, globals(), locals())
382 except TypeError:
383 pass
384 else:
385 self.fail('Did not raise Type in: ' + stmt)
386
387 # Check tuple re-use (implementation detail)
388 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
389 zip('abc', 'def'))
390 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
391 zip('abc', 'def'))
392 ids = map(id, izip_longest('abc', 'def'))
393 self.assertEqual(min(ids), max(ids))
394 ids = map(id, list(izip_longest('abc', 'def')))
395 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
396
397 def test_product(self):
398 for args, result in [
399 ([], [()]), # zero iterables
400 (['ab'], [('a',), ('b',)]), # one iterable
401 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
402 ([range(0), range(2), range(3)], []), # first iterable with zero length
403 ([range(2), range(0), range(3)], []), # middle iterable with zero length
404 ([range(2), range(3), range(0)], []), # last iterable with zero length
405 ]:
406 self.assertEqual(list(product(*args)), result)
407 for r in range(4):
408 self.assertEqual(list(product(*(args*r))),
409 list(product(*args, **dict(repeat=r))))
410 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
411 self.assertRaises(TypeError, product, range(6), None)
412
413 def product1(*args, **kwds):
414 pools = map(tuple, args) * kwds.get('repeat', 1)
415 n = len(pools)
416 if n == 0:
417 yield ()
418 return
419 if any(len(pool) == 0 for pool in pools):
420 return
421 indices = [0] * n
422 yield tuple(pool[i] for pool, i in zip(pools, indices))
423 while 1:
424 for i in reversed(range(n)): # right to left
425 if indices[i] == len(pools[i]) - 1:
426 continue
427 indices[i] += 1
428 for j in range(i+1, n):
429 indices[j] = 0
430 yield tuple(pool[i] for pool, i in zip(pools, indices))
431 break
432 else:
433 return
434
435 def product2(*args, **kwds):
436 'Pure python version used in docs'
437 pools = map(tuple, args) * kwds.get('repeat', 1)
438 result = [[]]
439 for pool in pools:
440 result = [x+[y] for x in result for y in pool]
441 for prod in result:
442 yield tuple(prod)
443
444 argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
445 set('abcdefg'), range(11), tuple(range(13))]
446 for i in range(100):
447 args = [random.choice(argtypes) for j in range(random.randrange(5))]
448 expected_len = prod(map(len, args))
449 self.assertEqual(len(list(product(*args))), expected_len)
450 self.assertEqual(list(product(*args)), list(product1(*args)))
451 self.assertEqual(list(product(*args)), list(product2(*args)))
452 args = map(iter, args)
453 self.assertEqual(len(list(product(*args))), expected_len)
454
455 # Test implementation detail: tuple re-use
456 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
457 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
458
459 def test_repeat(self):
460 self.assertEqual(zip(xrange(3),repeat('a')),
461 [(0, 'a'), (1, 'a'), (2, 'a')])
462 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
463 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
464 self.assertEqual(list(repeat('a', 0)), [])
465 self.assertEqual(list(repeat('a', -3)), [])
466 self.assertRaises(TypeError, repeat)
467 self.assertRaises(TypeError, repeat, None, 3, 4)
468 self.assertRaises(TypeError, repeat, None, 'a')
469 r = repeat(1+0j)
470 self.assertEqual(repr(r), 'repeat((1+0j))')
471 r = repeat(1+0j, 5)
472 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
473 list(r)
474 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
475
476 def test_imap(self):
477 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
478 [0**1, 1**2, 2**3])
479 self.assertEqual(list(imap(None, 'abc', range(5))),
480 [('a',0),('b',1),('c',2)])
481 self.assertEqual(list(imap(None, 'abc', count())),
482 [('a',0),('b',1),('c',2)])
483 self.assertEqual(take(2,imap(None, 'abc', count())),
484 [('a',0),('b',1)])
485 self.assertEqual(list(imap(operator.pow, [])), [])
486 self.assertRaises(TypeError, imap)
487 self.assertRaises(TypeError, imap, operator.neg)
488 self.assertRaises(TypeError, imap(10, range(5)).next)
489 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
490 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
491
492 def test_starmap(self):
493 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
494 [0**1, 1**2, 2**3])
495 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
496 [0**1, 1**2, 2**3])
497 self.assertEqual(list(starmap(operator.pow, [])), [])
498 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
499 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
500 self.assertRaises(TypeError, starmap)
501 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
502 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
503 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
504 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
505
506 def test_islice(self):
507 for args in [ # islice(args) should agree with range(args)
508 (10, 20, 3),
509 (10, 3, 20),
510 (10, 20),
511 (10, 3),
512 (20,)
513 ]:
514 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
515
516 for args, tgtargs in [ # Stop when seqn is exhausted
517 ((10, 110, 3), ((10, 100, 3))),
518 ((10, 110), ((10, 100))),
519 ((110,), (100,))
520 ]:
521 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
522
523 # Test stop=None
524 self.assertEqual(list(islice(xrange(10), None)), range(10))
525 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
526 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
527 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
528 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
529
530 # Test number of items consumed SF #1171417
531 it = iter(range(10))
532 self.assertEqual(list(islice(it, 3)), range(3))
533 self.assertEqual(list(it), range(3, 10))
534
535 # Test invalid arguments
536 self.assertRaises(TypeError, islice, xrange(10))
537 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
538 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
539 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
540 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
541 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
542 self.assertRaises(ValueError, islice, xrange(10), 'a')
543 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
544 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
545 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
546 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
547 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
548
549 def test_takewhile(self):
550 data = [1, 3, 5, 20, 2, 4, 6, 8]
551 underten = lambda x: x<10
552 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
553 self.assertEqual(list(takewhile(underten, [])), [])
554 self.assertRaises(TypeError, takewhile)
555 self.assertRaises(TypeError, takewhile, operator.pow)
556 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
557 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
558 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
559 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
560 self.assertEqual(list(t), [1, 1, 1])
561 self.assertRaises(StopIteration, t.next)
562
563 def test_dropwhile(self):
564 data = [1, 3, 5, 20, 2, 4, 6, 8]
565 underten = lambda x: x<10
566 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
567 self.assertEqual(list(dropwhile(underten, [])), [])
568 self.assertRaises(TypeError, dropwhile)
569 self.assertRaises(TypeError, dropwhile, operator.pow)
570 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
571 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
572 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
573
574 def test_tee(self):
575 n = 200
576 def irange(n):
577 for i in xrange(n):
578 yield i
579
580 a, b = tee([]) # test empty iterator
581 self.assertEqual(list(a), [])
582 self.assertEqual(list(b), [])
583
584 a, b = tee(irange(n)) # test 100% interleaved
585 self.assertEqual(zip(a,b), zip(range(n),range(n)))
586
587 a, b = tee(irange(n)) # test 0% interleaved
588 self.assertEqual(list(a), range(n))
589 self.assertEqual(list(b), range(n))
590
591 a, b = tee(irange(n)) # test dealloc of leading iterator
592 for i in xrange(100):
593 self.assertEqual(a.next(), i)
594 del a
595 self.assertEqual(list(b), range(n))
596
597 a, b = tee(irange(n)) # test dealloc of trailing iterator
598 for i in xrange(100):
599 self.assertEqual(a.next(), i)
600 del b
601 self.assertEqual(list(a), range(100, n))
602
603 for j in xrange(5): # test randomly interleaved
604 order = [0]*n + [1]*n
605 random.shuffle(order)
606 lists = ([], [])
607 its = tee(irange(n))
608 for i in order:
609 value = its[i].next()
610 lists[i].append(value)
611 self.assertEqual(lists[0], range(n))
612 self.assertEqual(lists[1], range(n))
613
614 # test argument format checking
615 self.assertRaises(TypeError, tee)
616 self.assertRaises(TypeError, tee, 3)
617 self.assertRaises(TypeError, tee, [1,2], 'x')
618 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
619
620 # tee object should be instantiable
621 a, b = tee('abc')
622 c = type(a)('def')
623 self.assertEqual(list(c), list('def'))
624
625 # test long-lagged and multi-way split
626 a, b, c = tee(xrange(2000), 3)
627 for i in xrange(100):
628 self.assertEqual(a.next(), i)
629 self.assertEqual(list(b), range(2000))
630 self.assertEqual([c.next(), c.next()], range(2))
631 self.assertEqual(list(a), range(100,2000))
632 self.assertEqual(list(c), range(2,2000))
633
634 # test values of n
635 self.assertRaises(TypeError, tee, 'abc', 'invalid')
636 self.assertRaises(ValueError, tee, [], -1)
637 for n in xrange(5):
638 result = tee('abc', n)
639 self.assertEqual(type(result), tuple)
640 self.assertEqual(len(result), n)
641 self.assertEqual(map(list, result), [list('abc')]*n)
642
643 # tee pass-through to copyable iterator
644 a, b = tee('abc')
645 c, d = tee(a)
646 self.assert_(a is c)
647
648 # test tee_new
649 t1, t2 = tee('abc')
650 tnew = type(t1)
651 self.assertRaises(TypeError, tnew)
652 self.assertRaises(TypeError, tnew, 10)
653 t3 = tnew(t1)
654 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
655
656 # test that tee objects are weak referencable
657 a, b = tee(xrange(10))
658 p = proxy(a)
659 self.assertEqual(getattr(p, '__class__'), type(b))
660 del a
661 self.assertRaises(ReferenceError, getattr, p, '__class__')
662
663 def test_StopIteration(self):
664 self.assertRaises(StopIteration, izip().next)
665
666 for f in (chain, cycle, izip, groupby):
667 self.assertRaises(StopIteration, f([]).next)
668 self.assertRaises(StopIteration, f(StopNow()).next)
669
670 self.assertRaises(StopIteration, islice([], None).next)
671 self.assertRaises(StopIteration, islice(StopNow(), None).next)
672
673 p, q = tee([])
674 self.assertRaises(StopIteration, p.next)
675 self.assertRaises(StopIteration, q.next)
676 p, q = tee(StopNow())
677 self.assertRaises(StopIteration, p.next)
678 self.assertRaises(StopIteration, q.next)
679
680 self.assertRaises(StopIteration, repeat(None, 0).next)
681
682 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
683 self.assertRaises(StopIteration, f(lambda x:x, []).next)
684 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
685
686 class TestExamples(unittest.TestCase):
687
688 def test_chain(self):
689 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
690
691 def test_chain_from_iterable(self):
692 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
693
694 def test_combinations(self):
695 self.assertEqual(list(combinations('ABCD', 2)),
696 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
697 self.assertEqual(list(combinations(range(4), 3)),
698 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
699
700 def test_count(self):
701 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
702
703 def test_cycle(self):
704 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
705
706 def test_dropwhile(self):
707 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
708
709 def test_groupby(self):
710 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
711 list('ABCDAB'))
712 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
713 [list('AAAA'), list('BBB'), list('CC'), list('D')])
714
715 def test_ifilter(self):
716 self.assertEqual(list(ifilter(lambda x: x%2, range(10))), [1,3,5,7,9])
717
718 def test_ifilterfalse(self):
719 self.assertEqual(list(ifilterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
720
721 def test_imap(self):
722 self.assertEqual(list(imap(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
723
724 def test_islice(self):
725 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
726 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
727 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
728 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
729
730 def test_izip(self):
731 self.assertEqual(list(izip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
732
733 def test_izip_longest(self):
734 self.assertEqual(list(izip_longest('ABCD', 'xy', fillvalue='-')),
735 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
736
737 def test_permutations(self):
738 self.assertEqual(list(permutations('ABCD', 2)),
739 map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split()))
740 self.assertEqual(list(permutations(range(3))),
741 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
742
743 def test_product(self):
744 self.assertEqual(list(product('ABCD', 'xy')),
745 map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split()))
746 self.assertEqual(list(product(range(2), repeat=3)),
747 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
748 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
749
750 def test_repeat(self):
751 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
752
753 def test_stapmap(self):
754 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
755 [32, 9, 1000])
756
757 def test_takewhile(self):
758 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
759
760
761 class TestGC(unittest.TestCase):
762
763 def makecycle(self, iterator, container):
764 container.append(iterator)
765 iterator.next()
766 del container, iterator
767
768 def test_chain(self):
769 a = []
770 self.makecycle(chain(a), a)
771
772 def test_chain_from_iterable(self):
773 a = []
774 self.makecycle(chain.from_iterable([a]), a)
775
776 def test_combinations(self):
777 a = []
778 self.makecycle(combinations([1,2,a,3], 3), a)
779
780 def test_cycle(self):
781 a = []
782 self.makecycle(cycle([a]*2), a)
783
784 def test_dropwhile(self):
785 a = []
786 self.makecycle(dropwhile(bool, [0, a, a]), a)
787
788 def test_groupby(self):
789 a = []
790 self.makecycle(groupby([a]*2, lambda x:x), a)
791
792 def test_issue2246(self):
793 # Issue 2246 -- the _grouper iterator was not included in GC
794 n = 10
795 keyfunc = lambda x: x
796 for i, j in groupby(xrange(n), key=keyfunc):
797 keyfunc.__dict__.setdefault('x',[]).append(j)
798
799 def test_ifilter(self):
800 a = []
801 self.makecycle(ifilter(lambda x:True, [a]*2), a)
802
803 def test_ifilterfalse(self):
804 a = []
805 self.makecycle(ifilterfalse(lambda x:False, a), a)
806
807 def test_izip(self):
808 a = []
809 self.makecycle(izip([a]*2, [a]*3), a)
810
811 def test_izip_longest(self):
812 a = []
813 self.makecycle(izip_longest([a]*2, [a]*3), a)
814 b = [a, None]
815 self.makecycle(izip_longest([a]*2, [a]*3, fillvalue=b), a)
816
817 def test_imap(self):
818 a = []
819 self.makecycle(imap(lambda x:x, [a]*2), a)
820
821 def test_islice(self):
822 a = []
823 self.makecycle(islice([a]*2, None), a)
824
825 def test_permutations(self):
826 a = []
827 self.makecycle(permutations([1,2,a,3], 3), a)
828
829 def test_product(self):
830 a = []
831 self.makecycle(product([1,2,a,3], repeat=3), a)
832
833 def test_repeat(self):
834 a = []
835 self.makecycle(repeat(a), a)
836
837 def test_starmap(self):
838 a = []
839 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
840
841 def test_takewhile(self):
842 a = []
843 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
844
845 def R(seqn):
846 'Regular generator'
847 for i in seqn:
848 yield i
849
850 class G:
851 'Sequence using __getitem__'
852 def __init__(self, seqn):
853 self.seqn = seqn
854 def __getitem__(self, i):
855 return self.seqn[i]
856
857 class I:
858 'Sequence using iterator protocol'
859 def __init__(self, seqn):
860 self.seqn = seqn
861 self.i = 0
862 def __iter__(self):
863 return self
864 def next(self):
865 if self.i >= len(self.seqn): raise StopIteration
866 v = self.seqn[self.i]
867 self.i += 1
868 return v
869
870 class Ig:
871 'Sequence using iterator protocol defined with a generator'
872 def __init__(self, seqn):
873 self.seqn = seqn
874 self.i = 0
875 def __iter__(self):
876 for val in self.seqn:
877 yield val
878
879 class X:
880 'Missing __getitem__ and __iter__'
881 def __init__(self, seqn):
882 self.seqn = seqn
883 self.i = 0
884 def next(self):
885 if self.i >= len(self.seqn): raise StopIteration
886 v = self.seqn[self.i]
887 self.i += 1
888 return v
889
890 class N:
891 'Iterator missing next()'
892 def __init__(self, seqn):
893 self.seqn = seqn
894 self.i = 0
895 def __iter__(self):
896 return self
897
898 class E:
899 'Test propagation of exceptions'
900 def __init__(self, seqn):
901 self.seqn = seqn
902 self.i = 0
903 def __iter__(self):
904 return self
905 def next(self):
906 3 // 0
907
908 class S:
909 'Test immediate stop'
910 def __init__(self, seqn):
911 pass
912 def __iter__(self):
913 return self
914 def next(self):
915 raise StopIteration
916
917 def L(seqn):
918 'Test multiple tiers of iterators'
919 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
920
921
922 class TestVariousIteratorArgs(unittest.TestCase):
923
924 def test_chain(self):
925 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
926 for g in (G, I, Ig, S, L, R):
927 self.assertEqual(list(chain(g(s))), list(g(s)))
928 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
929 self.assertRaises(TypeError, list, chain(X(s)))
930 self.assertRaises(TypeError, list, chain(N(s)))
931 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
932
933 def test_product(self):
934 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
935 self.assertRaises(TypeError, product, X(s))
936 self.assertRaises(TypeError, product, N(s))
937 self.assertRaises(ZeroDivisionError, product, E(s))
938
939 def test_cycle(self):
940 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
941 for g in (G, I, Ig, S, L, R):
942 tgtlen = len(s) * 3
943 expected = list(g(s))*3
944 actual = list(islice(cycle(g(s)), tgtlen))
945 self.assertEqual(actual, expected)
946 self.assertRaises(TypeError, cycle, X(s))
947 self.assertRaises(TypeError, list, cycle(N(s)))
948 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
949
950 def test_groupby(self):
951 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
952 for g in (G, I, Ig, S, L, R):
953 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
954 self.assertRaises(TypeError, groupby, X(s))
955 self.assertRaises(TypeError, list, groupby(N(s)))
956 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
957
958 def test_ifilter(self):
959 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
960 for g in (G, I, Ig, S, L, R):
961 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
962 self.assertRaises(TypeError, ifilter, isEven, X(s))
963 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
964 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
965
966 def test_ifilterfalse(self):
967 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
968 for g in (G, I, Ig, S, L, R):
969 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
970 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
971 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
972 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
973
974 def test_izip(self):
975 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
976 for g in (G, I, Ig, S, L, R):
977 self.assertEqual(list(izip(g(s))), zip(g(s)))
978 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
979 self.assertRaises(TypeError, izip, X(s))
980 self.assertRaises(TypeError, list, izip(N(s)))
981 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
982
983 def test_iziplongest(self):
984 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
985 for g in (G, I, Ig, S, L, R):
986 self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
987 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
988 self.assertRaises(TypeError, izip_longest, X(s))
989 self.assertRaises(TypeError, list, izip_longest(N(s)))
990 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
991
992 def test_imap(self):
993 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
994 for g in (G, I, Ig, S, L, R):
995 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
996 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
997 self.assertRaises(TypeError, imap, onearg, X(s))
998 self.assertRaises(TypeError, list, imap(onearg, N(s)))
999 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
1000
1001 def test_islice(self):
1002 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1003 for g in (G, I, Ig, S, L, R):
1004 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1005 self.assertRaises(TypeError, islice, X(s), 10)
1006 self.assertRaises(TypeError, list, islice(N(s), 10))
1007 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1008
1009 def test_starmap(self):
1010 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1011 for g in (G, I, Ig, S, L, R):
1012 ss = zip(s, s)
1013 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
1014 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
1015 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
1016 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1017
1018 def test_takewhile(self):
1019 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1020 for g in (G, I, Ig, S, L, R):
1021 tgt = []
1022 for elem in g(s):
1023 if not isEven(elem): break
1024 tgt.append(elem)
1025 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1026 self.assertRaises(TypeError, takewhile, isEven, X(s))
1027 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
1028 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1029
1030 def test_dropwhile(self):
1031 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1032 for g in (G, I, Ig, S, L, R):
1033 tgt = []
1034 for elem in g(s):
1035 if not tgt and isOdd(elem): continue
1036 tgt.append(elem)
1037 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1038 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
1039 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
1040 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1041
1042 def test_tee(self):
1043 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1044 for g in (G, I, Ig, S, L, R):
1045 it1, it2 = tee(g(s))
1046 self.assertEqual(list(it1), list(g(s)))
1047 self.assertEqual(list(it2), list(g(s)))
1048 self.assertRaises(TypeError, tee, X(s))
1049 self.assertRaises(TypeError, list, tee(N(s))[0])
1050 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1051
1052 class LengthTransparency(unittest.TestCase):
1053
1054 def test_repeat(self):
1055 from test.test_iterlen import len
1056 self.assertEqual(len(repeat(None, 50)), 50)
1057 self.assertRaises(TypeError, len, repeat(None))
1058
1059 class RegressionTests(unittest.TestCase):
1060
1061 def test_sf_793826(self):
1062 # Fix Armin Rigo's successful efforts to wreak havoc
1063
1064 def mutatingtuple(tuple1, f, tuple2):
1065 # this builds a tuple t which is a copy of tuple1,
1066 # then calls f(t), then mutates t to be equal to tuple2
1067 # (needs len(tuple1) == len(tuple2)).
1068 def g(value, first=[1]):
1069 if first:
1070 del first[:]
1071 f(z.next())
1072 return value
1073 items = list(tuple2)
1074 items[1:1] = list(tuple1)
1075 gen = imap(g, items)
1076 z = izip(*[gen]*len(tuple1))
1077 z.next()
1078
1079 def f(t):
1080 global T
1081 T = t
1082 first[:] = list(T)
1083
1084 first = []
1085 mutatingtuple((1,2,3), f, (4,5,6))
1086 second = list(T)
1087 self.assertEqual(first, second)
1088
1089
1090 def test_sf_950057(self):
1091 # Make sure that chain() and cycle() catch exceptions immediately
1092 # rather than when shifting between input sources
1093
1094 def gen1():
1095 hist.append(0)
1096 yield 1
1097 hist.append(1)
1098 raise AssertionError
1099 hist.append(2)
1100
1101 def gen2(x):
1102 hist.append(3)
1103 yield 2
1104 hist.append(4)
1105 if x:
1106 raise StopIteration
1107
1108 hist = []
1109 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1110 self.assertEqual(hist, [0,1])
1111
1112 hist = []
1113 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1114 self.assertEqual(hist, [0,1])
1115
1116 hist = []
1117 self.assertRaises(AssertionError, list, cycle(gen1()))
1118 self.assertEqual(hist, [0,1])
1119
1120 class SubclassWithKwargsTest(unittest.TestCase):
1121 def test_keywords_in_subclass(self):
1122 # count is not subclassable...
1123 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
1124 starmap, islice, takewhile, dropwhile, cycle):
1125 class Subclass(cls):
1126 def __init__(self, newarg=None, *args):
1127 cls.__init__(self, *args)
1128 try:
1129 Subclass(newarg=1)
1130 except TypeError, err:
1131 # we expect type errors because of wrong argument count
1132 self.failIf("does not take keyword arguments" in err.args[0])
1133
1134
1135 libreftest = """ Doctest for examples in the library reference: libitertools.tex
1136
1137
1138 >>> amounts = [120.15, 764.05, 823.14]
1139 >>> for checknum, amount in izip(count(1200), amounts):
1140 ... print 'Check %d is for $%.2f' % (checknum, amount)
1141 ...
1142 Check 1200 is for $120.15
1143 Check 1201 is for $764.05
1144 Check 1202 is for $823.14
1145
1146 >>> import operator
1147 >>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1148 ... print cube
1149 ...
1150 1
1151 8
1152 27
1153
1154 >>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
1155 >>> for name in islice(reportlines, 3, None, 2):
1156 ... print name.title()
1157 ...
1158 Alex
1159 Laura
1160 Martin
1161 Walter
1162 Samuele
1163
1164 >>> from operator import itemgetter
1165 >>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
1166 >>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
1167 >>> for k, g in groupby(di, itemgetter(1)):
1168 ... print k, map(itemgetter(0), g)
1169 ...
1170 1 ['a', 'c', 'e']
1171 2 ['b', 'd', 'f']
1172 3 ['g']
1173
1174 # Find runs of consecutive numbers using groupby. The key to the solution
1175 # is differencing with a range so that consecutive numbers all appear in
1176 # same group.
1177 >>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
1178 >>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
1179 ... print map(operator.itemgetter(1), g)
1180 ...
1181 [1]
1182 [4, 5, 6]
1183 [10]
1184 [15, 16, 17, 18]
1185 [22]
1186 [25, 26, 27, 28]
1187
1188 >>> def take(n, iterable):
1189 ... "Return first n items of the iterable as a list"
1190 ... return list(islice(iterable, n))
1191
1192 >>> def enumerate(iterable, start=0):
1193 ... return izip(count(start), iterable)
1194
1195 >>> def tabulate(function, start=0):
1196 ... "Return function(0), function(1), ..."
1197 ... return imap(function, count(start))
1198
1199 >>> def nth(iterable, n):
1200 ... "Returns the nth item or empty list"
1201 ... return list(islice(iterable, n, n+1))
1202
1203 >>> def quantify(iterable, pred=bool):
1204 ... "Count how many times the predicate is true"
1205 ... return sum(imap(pred, iterable))
1206
1207 >>> def padnone(iterable):
1208 ... "Returns the sequence elements and then returns None indefinitely"
1209 ... return chain(iterable, repeat(None))
1210
1211 >>> def ncycles(iterable, n):
1212 ... "Returns the seqeuence elements n times"
1213 ... return chain(*repeat(iterable, n))
1214
1215 >>> def dotproduct(vec1, vec2):
1216 ... return sum(imap(operator.mul, vec1, vec2))
1217
1218 >>> def flatten(listOfLists):
1219 ... return list(chain.from_iterable(listOfLists))
1220
1221 >>> def repeatfunc(func, times=None, *args):
1222 ... "Repeat calls to func with specified arguments."
1223 ... " Example: repeatfunc(random.random)"
1224 ... if times is None:
1225 ... return starmap(func, repeat(args))
1226 ... else:
1227 ... return starmap(func, repeat(args, times))
1228
1229 >>> def pairwise(iterable):
1230 ... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1231 ... a, b = tee(iterable)
1232 ... for elem in b:
1233 ... break
1234 ... return izip(a, b)
1235
1236 >>> def grouper(n, iterable, fillvalue=None):
1237 ... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
1238 ... args = [iter(iterable)] * n
1239 ... kwds = dict(fillvalue=fillvalue)
1240 ... return izip_longest(*args, **kwds)
1241
1242 >>> def roundrobin(*iterables):
1243 ... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
1244 ... # Recipe credited to George Sakkis
1245 ... pending = len(iterables)
1246 ... nexts = cycle(iter(it).next for it in iterables)
1247 ... while pending:
1248 ... try:
1249 ... for next in nexts:
1250 ... yield next()
1251 ... except StopIteration:
1252 ... pending -= 1
1253 ... nexts = cycle(islice(nexts, pending))
1254
1255 >>> def powerset(iterable):
1256 ... "powerset('ab') --> set([]), set(['a']), set(['b']), set(['a', 'b'])"
1257 ... # Recipe credited to Eric Raymond
1258 ... pairs = [(2**i, x) for i, x in enumerate(iterable)]
1259 ... for n in xrange(2**len(pairs)):
1260 ... yield set(x for m, x in pairs if m&n)
1261
1262 >>> def compress(data, selectors):
1263 ... "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
1264 ... return (d for d, s in izip(data, selectors) if s)
1265
1266 >>> def combinations_with_replacement(iterable, r):
1267 ... "combinations_with_replacement('ABC', 3) --> AA AB AC BB BC CC"
1268 ... pool = tuple(iterable)
1269 ... n = len(pool)
1270 ... indices = [0] * r
1271 ... yield tuple(pool[i] for i in indices)
1272 ... while 1:
1273 ... for i in reversed(range(r)):
1274 ... if indices[i] != n - 1:
1275 ... break
1276 ... else:
1277 ... return
1278 ... indices[i:] = [indices[i] + 1] * (r - i)
1279 ... yield tuple(pool[i] for i in indices)
1280
1281 This is not part of the examples but it tests to make sure the definitions
1282 perform as purported.
1283
1284 >>> take(10, count())
1285 [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1286
1287 >>> list(enumerate('abc'))
1288 [(0, 'a'), (1, 'b'), (2, 'c')]
1289
1290 >>> list(islice(tabulate(lambda x: 2*x), 4))
1291 [0, 2, 4, 6]
1292
1293 >>> nth('abcde', 3)
1294 ['d']
1295
1296 >>> quantify(xrange(99), lambda x: x%2==0)
1297 50
1298
1299 >>> a = [[1, 2, 3], [4, 5, 6]]
1300 >>> flatten(a)
1301 [1, 2, 3, 4, 5, 6]
1302
1303 >>> list(repeatfunc(pow, 5, 2, 3))
1304 [8, 8, 8, 8, 8]
1305
1306 >>> import random
1307 >>> take(5, imap(int, repeatfunc(random.random)))
1308 [0, 0, 0, 0, 0]
1309
1310 >>> list(pairwise('abcd'))
1311 [('a', 'b'), ('b', 'c'), ('c', 'd')]
1312
1313 >>> list(pairwise([]))
1314 []
1315
1316 >>> list(pairwise('a'))
1317 []
1318
1319 >>> list(islice(padnone('abc'), 0, 6))
1320 ['a', 'b', 'c', None, None, None]
1321
1322 >>> list(ncycles('abc', 3))
1323 ['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1324
1325 >>> dotproduct([1,2,3], [4,5,6])
1326 32
1327
1328 >>> list(grouper(3, 'abcdefg', 'x'))
1329 [('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1330
1331 >>> list(roundrobin('abc', 'd', 'ef'))
1332 ['a', 'd', 'e', 'b', 'f', 'c']
1333
1334 >>> map(sorted, powerset('ab'))
1335 [[], ['a'], ['b'], ['a', 'b']]
1336
1337 >>> list(compress('abcdef', [1,0,1,0,1,1]))
1338 ['a', 'c', 'e', 'f']
1339
1340 >>> list(combinations_with_replacement('abc', 2))
1341 [('a', 'a'), ('a', 'b'), ('a', 'c'), ('b', 'b'), ('b', 'c'), ('c', 'c')]
1342
1343 """
1344
1345 __test__ = {'libreftest' : libreftest}
1346
1347 def test_main(verbose=None):
1348 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
1349 RegressionTests, LengthTransparency,
1350 SubclassWithKwargsTest, TestExamples)
1351 test_support.run_unittest(*test_classes)
1352
1353 # verify reference counting
1354 if verbose and hasattr(sys, "gettotalrefcount"):
1355 import gc
1356 counts = [None] * 5
1357 for i in xrange(len(counts)):
1358 test_support.run_unittest(*test_classes)
1359 gc.collect()
1360 counts[i] = sys.gettotalrefcount()
1361 print counts
1362
1363 # doctest the examples in the library reference
1364 test_support.run_doctest(sys.modules[__name__], verbose)
1365
1366 if __name__ == "__main__":
1367 test_main(verbose=True)
python git mirror