• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1import unittest
2from test import test_support
3from itertools import *
4import weakref
5from decimal import Decimal
6from fractions import Fraction
7import sys
8import operator
9import random
10import copy
11import pickle
12from functools import reduce
13maxsize = test_support.MAX_Py_ssize_t
14minsize = -maxsize-1
15
16def onearg(x):
17    'Test function of one argument'
18    return 2*x
19
20def errfunc(*args):
21    'Test function that raises an error'
22    raise ValueError
23
24def gen3():
25    'Non-restartable source sequence'
26    for i in (0, 1, 2):
27        yield i
28
29def isEven(x):
30    'Test predicate'
31    return x%2==0
32
33def isOdd(x):
34    'Test predicate'
35    return x%2==1
36
37class StopNow:
38    'Class emulating an empty iterable.'
39    def __iter__(self):
40        return self
41    def next(self):
42        raise StopIteration
43
44def take(n, seq):
45    'Convenience function for partially consuming a long of infinite iterable'
46    return list(islice(seq, n))
47
48def prod(iterable):
49    return reduce(operator.mul, iterable, 1)
50
51def fact(n):
52    'Factorial'
53    return prod(range(1, n+1))
54
55class TestBasicOps(unittest.TestCase):
56    def test_chain(self):
57
58        def chain2(*iterables):
59            'Pure python version in the docs'
60            for it in iterables:
61                for element in it:
62                    yield element
63
64        for c in (chain, chain2):
65            self.assertEqual(list(c('abc', 'def')), list('abcdef'))
66            self.assertEqual(list(c('abc')), list('abc'))
67            self.assertEqual(list(c('')), [])
68            self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
69            self.assertRaises(TypeError, list,c(2, 3))
70
71    def test_chain_from_iterable(self):
72        self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
73        self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
74        self.assertEqual(list(chain.from_iterable([''])), [])
75        self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
76        self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
77
78    def test_combinations(self):
79        self.assertRaises(TypeError, combinations, 'abc')       # missing r argument
80        self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
81        self.assertRaises(TypeError, combinations, None)        # pool is not iterable
82        self.assertRaises(ValueError, combinations, 'abc', -2)  # r is negative
83        self.assertEqual(list(combinations('abc', 32)), [])     # r > n
84        self.assertEqual(list(combinations(range(4), 3)),
85                                           [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
86
87        def combinations1(iterable, r):
88            'Pure python version shown in the docs'
89            pool = tuple(iterable)
90            n = len(pool)
91            if r > n:
92                return
93            indices = range(r)
94            yield tuple(pool[i] for i in indices)
95            while 1:
96                for i in reversed(range(r)):
97                    if indices[i] != i + n - r:
98                        break
99                else:
100                    return
101                indices[i] += 1
102                for j in range(i+1, r):
103                    indices[j] = indices[j-1] + 1
104                yield tuple(pool[i] for i in indices)
105
106        def combinations2(iterable, r):
107            'Pure python version shown in the docs'
108            pool = tuple(iterable)
109            n = len(pool)
110            for indices in permutations(range(n), r):
111                if sorted(indices) == list(indices):
112                    yield tuple(pool[i] for i in indices)
113
114        def combinations3(iterable, r):
115            'Pure python version from cwr()'
116            pool = tuple(iterable)
117            n = len(pool)
118            for indices in combinations_with_replacement(range(n), r):
119                if len(set(indices)) == r:
120                    yield tuple(pool[i] for i in indices)
121
122        for n in range(7):
123            values = [5*x-12 for x in range(n)]
124            for r in range(n+2):
125                result = list(combinations(values, r))
126                self.assertEqual(len(result), 0 if r>n else fact(n) // fact(r) // fact(n-r)) # right number of combs
127                self.assertEqual(len(result), len(set(result)))         # no repeats
128                self.assertEqual(result, sorted(result))                # lexicographic order
129                for c in result:
130                    self.assertEqual(len(c), r)                         # r-length combinations
131                    self.assertEqual(len(set(c)), r)                    # no duplicate elements
132                    self.assertEqual(list(c), sorted(c))                # keep original ordering
133                    self.assertTrue(all(e in values for e in c))           # elements taken from input iterable
134                    self.assertEqual(list(c),
135                                     [e for e in values if e in c])      # comb is a subsequence of the input iterable
136                self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
137                self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
138                self.assertEqual(result, list(combinations3(values, r))) # matches second pure python version
139
140    @test_support.bigaddrspacetest
141    def test_combinations_overflow(self):
142        with self.assertRaises((OverflowError, MemoryError)):
143            combinations("AA", 2**29)
144
145    @test_support.impl_detail("tuple reuse is specific to CPython")
146    def test_combinations_tuple_reuse(self):
147        self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
148        self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
149
150    def test_combinations_with_replacement(self):
151        cwr = combinations_with_replacement
152        self.assertRaises(TypeError, cwr, 'abc')       # missing r argument
153        self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
154        self.assertRaises(TypeError, cwr, None)        # pool is not iterable
155        self.assertRaises(ValueError, cwr, 'abc', -2)  # r is negative
156        self.assertEqual(list(cwr('ABC', 2)),
157                         [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
158
159        def cwr1(iterable, r):
160            'Pure python version shown in the docs'
161            # number items returned:  (n+r-1)! / r! / (n-1)! when n>0
162            pool = tuple(iterable)
163            n = len(pool)
164            if not n and r:
165                return
166            indices = [0] * r
167            yield tuple(pool[i] for i in indices)
168            while 1:
169                for i in reversed(range(r)):
170                    if indices[i] != n - 1:
171                        break
172                else:
173                    return
174                indices[i:] = [indices[i] + 1] * (r - i)
175                yield tuple(pool[i] for i in indices)
176
177        def cwr2(iterable, r):
178            'Pure python version shown in the docs'
179            pool = tuple(iterable)
180            n = len(pool)
181            for indices in product(range(n), repeat=r):
182                if sorted(indices) == list(indices):
183                    yield tuple(pool[i] for i in indices)
184
185        def numcombs(n, r):
186            if not n:
187                return 0 if r else 1
188            return fact(n+r-1) // fact(r) // fact(n-1)
189
190        for n in range(7):
191            values = [5*x-12 for x in range(n)]
192            for r in range(n+2):
193                result = list(cwr(values, r))
194
195                self.assertEqual(len(result), numcombs(n, r))           # right number of combs
196                self.assertEqual(len(result), len(set(result)))         # no repeats
197                self.assertEqual(result, sorted(result))                # lexicographic order
198
199                regular_combs = list(combinations(values, r))           # compare to combs without replacement
200                if n == 0 or r <= 1:
201                    self.assertEqual(result, regular_combs)            # cases that should be identical
202                else:
203                    self.assertTrue(set(result) >= set(regular_combs))     # rest should be supersets of regular combs
204
205                for c in result:
206                    self.assertEqual(len(c), r)                         # r-length combinations
207                    noruns = [k for k,v in groupby(c)]                  # combo without consecutive repeats
208                    self.assertEqual(len(noruns), len(set(noruns)))     # no repeats other than consecutive
209                    self.assertEqual(list(c), sorted(c))                # keep original ordering
210                    self.assertTrue(all(e in values for e in c))           # elements taken from input iterable
211                    self.assertEqual(noruns,
212                                     [e for e in values if e in c])     # comb is a subsequence of the input iterable
213                self.assertEqual(result, list(cwr1(values, r)))         # matches first pure python version
214                self.assertEqual(result, list(cwr2(values, r)))         # matches second pure python version
215
216    @test_support.bigaddrspacetest
217    def test_combinations_with_replacement_overflow(self):
218        with self.assertRaises((OverflowError, MemoryError)):
219            combinations_with_replacement("AA", 2**30)
220
221    @test_support.impl_detail("tuple reuse is specific to CPython")
222    def test_combinations_with_replacement_tuple_reuse(self):
223        cwr = combinations_with_replacement
224        self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
225        self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
226
227    def test_permutations(self):
228        self.assertRaises(TypeError, permutations)              # too few arguments
229        self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
230        self.assertRaises(TypeError, permutations, None)        # pool is not iterable
231        self.assertRaises(ValueError, permutations, 'abc', -2)  # r is negative
232        self.assertEqual(list(permutations('abc', 32)), [])     # r > n
233        self.assertRaises(TypeError, permutations, 'abc', 's')  # r is not an int or None
234        self.assertEqual(list(permutations(range(3), 2)),
235                                           [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
236
237        def permutations1(iterable, r=None):
238            'Pure python version shown in the docs'
239            pool = tuple(iterable)
240            n = len(pool)
241            r = n if r is None else r
242            if r > n:
243                return
244            indices = range(n)
245            cycles = range(n, n-r, -1)
246            yield tuple(pool[i] for i in indices[:r])
247            while n:
248                for i in reversed(range(r)):
249                    cycles[i] -= 1
250                    if cycles[i] == 0:
251                        indices[i:] = indices[i+1:] + indices[i:i+1]
252                        cycles[i] = n - i
253                    else:
254                        j = cycles[i]
255                        indices[i], indices[-j] = indices[-j], indices[i]
256                        yield tuple(pool[i] for i in indices[:r])
257                        break
258                else:
259                    return
260
261        def permutations2(iterable, r=None):
262            'Pure python version shown in the docs'
263            pool = tuple(iterable)
264            n = len(pool)
265            r = n if r is None else r
266            for indices in product(range(n), repeat=r):
267                if len(set(indices)) == r:
268                    yield tuple(pool[i] for i in indices)
269
270        for n in range(7):
271            values = [5*x-12 for x in range(n)]
272            for r in range(n+2):
273                result = list(permutations(values, r))
274                self.assertEqual(len(result), 0 if r>n else fact(n) // fact(n-r))      # right number of perms
275                self.assertEqual(len(result), len(set(result)))         # no repeats
276                self.assertEqual(result, sorted(result))                # lexicographic order
277                for p in result:
278                    self.assertEqual(len(p), r)                         # r-length permutations
279                    self.assertEqual(len(set(p)), r)                    # no duplicate elements
280                    self.assertTrue(all(e in values for e in p))           # elements taken from input iterable
281                self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
282                self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
283                if r == n:
284                    self.assertEqual(result, list(permutations(values, None))) # test r as None
285                    self.assertEqual(result, list(permutations(values)))       # test default r
286
287    @test_support.bigaddrspacetest
288    def test_permutations_overflow(self):
289        with self.assertRaises((OverflowError, MemoryError)):
290            permutations("A", 2**30)
291
292    @test_support.impl_detail("tuple reuse is specific to CPython")
293    def test_permutations_tuple_reuse(self):
294        self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
295        self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
296
297    def test_combinatorics(self):
298        # Test relationships between product(), permutations(),
299        # combinations() and combinations_with_replacement().
300
301        for n in range(6):
302            s = 'ABCDEFG'[:n]
303            for r in range(8):
304                prod = list(product(s, repeat=r))
305                cwr = list(combinations_with_replacement(s, r))
306                perm = list(permutations(s, r))
307                comb = list(combinations(s, r))
308
309                # Check size
310                self.assertEqual(len(prod), n**r)
311                self.assertEqual(len(cwr), (fact(n+r-1) // fact(r) // fact(n-1)) if n else (not r))
312                self.assertEqual(len(perm), 0 if r>n else fact(n) // fact(n-r))
313                self.assertEqual(len(comb), 0 if r>n else fact(n) // fact(r) // fact(n-r))
314
315                # Check lexicographic order without repeated tuples
316                self.assertEqual(prod, sorted(set(prod)))
317                self.assertEqual(cwr, sorted(set(cwr)))
318                self.assertEqual(perm, sorted(set(perm)))
319                self.assertEqual(comb, sorted(set(comb)))
320
321                # Check interrelationships
322                self.assertEqual(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
323                self.assertEqual(perm, [t for t in prod if len(set(t))==r])    # perm: prods with no dups
324                self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
325                self.assertEqual(comb, [t for t in cwr if len(set(t))==r])      # comb: cwrs without dups
326                self.assertEqual(comb, filter(set(cwr).__contains__, perm))     # comb: perm that is a cwr
327                self.assertEqual(comb, filter(set(perm).__contains__, cwr))     # comb: cwr that is a perm
328                self.assertEqual(comb, sorted(set(cwr) & set(perm)))            # comb: both a cwr and a perm
329
330    def test_compress(self):
331        self.assertEqual(list(compress(data='ABCDEF', selectors=[1,0,1,0,1,1])), list('ACEF'))
332        self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
333        self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
334        self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
335        self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
336        self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
337        n = 10000
338        data = chain.from_iterable(repeat(range(6), n))
339        selectors = chain.from_iterable(repeat((0, 1)))
340        self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
341        self.assertRaises(TypeError, compress, None, range(6))      # 1st arg not iterable
342        self.assertRaises(TypeError, compress, range(6), None)      # 2nd arg not iterable
343        self.assertRaises(TypeError, compress, range(6))            # too few args
344        self.assertRaises(TypeError, compress, range(6), None)      # too many args
345
346    def test_count(self):
347        self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
348        self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
349        self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
350        self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
351        self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
352        self.assertRaises(TypeError, count, 2, 3, 4)
353        self.assertRaises(TypeError, count, 'a')
354        self.assertEqual(take(10, count(maxsize-5)), range(maxsize-5, maxsize+5))
355        self.assertEqual(take(10, count(-maxsize-5)), range(-maxsize-5, -maxsize+5))
356        self.assertEqual(take(3, count(3.25)), [3.25, 4.25, 5.25])
357        self.assertEqual(take(3, count(3.25-4j)), [3.25-4j, 4.25-4j, 5.25-4j])
358        self.assertEqual(take(3, count(Decimal('1.1'))),
359                         [Decimal('1.1'), Decimal('2.1'), Decimal('3.1')])
360        self.assertEqual(take(3, count(Fraction(2, 3))),
361                         [Fraction(2, 3), Fraction(5, 3), Fraction(8, 3)])
362        BIGINT = 1<<1000
363        self.assertEqual(take(3, count(BIGINT)), [BIGINT, BIGINT+1, BIGINT+2])
364        c = count(3)
365        self.assertEqual(repr(c), 'count(3)')
366        c.next()
367        self.assertEqual(repr(c), 'count(4)')
368        c = count(-9)
369        self.assertEqual(repr(c), 'count(-9)')
370        c.next()
371        self.assertEqual(next(c), -8)
372        self.assertEqual(repr(count(10.25)), 'count(10.25)')
373        self.assertEqual(repr(count(10.0)), 'count(10.0)')
374        self.assertEqual(type(next(count(10.0))), float)
375        for i in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 10, sys.maxint-5, sys.maxint+5):
376            # Test repr (ignoring the L in longs)
377            r1 = repr(count(i)).replace('L', '')
378            r2 = 'count(%r)'.__mod__(i).replace('L', '')
379            self.assertEqual(r1, r2)
380
381        # check copy, deepcopy, pickle
382        for value in -3, 3, sys.maxint-5, sys.maxint+5:
383            c = count(value)
384            self.assertEqual(next(copy.copy(c)), value)
385            self.assertEqual(next(copy.deepcopy(c)), value)
386            for proto in range(pickle.HIGHEST_PROTOCOL + 1):
387                self.assertEqual(next(pickle.loads(pickle.dumps(c, proto))), value)
388
389    def test_count_with_stride(self):
390        self.assertEqual(zip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
391        self.assertEqual(zip('abc',count(start=2,step=3)),
392                         [('a', 2), ('b', 5), ('c', 8)])
393        self.assertEqual(zip('abc',count(step=-1)),
394                         [('a', 0), ('b', -1), ('c', -2)])
395        self.assertRaises(TypeError, count, 'a', 'b')
396        self.assertEqual(zip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
397        self.assertEqual(zip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
398        self.assertEqual(zip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
399        self.assertEqual(zip('abc',count(2,1L)), [('a', 2), ('b', 3), ('c', 4)])
400        self.assertEqual(zip('abc',count(2,3L)), [('a', 2), ('b', 5), ('c', 8)])
401        self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
402        self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
403        self.assertEqual(take(3, count(10, maxsize+5)),
404                         range(10, 10+3*(maxsize+5), maxsize+5))
405        self.assertEqual(take(3, count(2, 1.25)), [2, 3.25, 4.5])
406        self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
407        self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
408                         [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
409        self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
410                         [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
411        BIGINT = 1<<1000
412        self.assertEqual(take(3, count(step=BIGINT)), [0, BIGINT, 2*BIGINT])
413        self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
414        c = count(3, 5)
415        self.assertEqual(repr(c), 'count(3, 5)')
416        c.next()
417        self.assertEqual(repr(c), 'count(8, 5)')
418        c = count(-9, 0)
419        self.assertEqual(repr(c), 'count(-9, 0)')
420        c.next()
421        self.assertEqual(repr(c), 'count(-9, 0)')
422        c = count(-9, -3)
423        self.assertEqual(repr(c), 'count(-9, -3)')
424        c.next()
425        self.assertEqual(repr(c), 'count(-12, -3)')
426        self.assertEqual(repr(c), 'count(-12, -3)')
427        self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
428        self.assertEqual(repr(count(10.5, 1)), 'count(10.5)')           # suppress step=1 when it's an int
429        self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)')   # do show float values lilke 1.0
430        self.assertEqual(repr(count(10, 1.00)), 'count(10, 1.0)')
431        c = count(10, 1.0)
432        self.assertEqual(type(next(c)), int)
433        self.assertEqual(type(next(c)), float)
434        for i in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 10, sys.maxint-5, sys.maxint+5):
435            for j in  (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 1, 10, sys.maxint-5, sys.maxint+5):
436                # Test repr (ignoring the L in longs)
437                r1 = repr(count(i, j)).replace('L', '')
438                if j == 1:
439                    r2 = ('count(%r)' % i).replace('L', '')
440                else:
441                    r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
442                self.assertEqual(r1, r2)
443
444    def test_cycle(self):
445        self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
446        self.assertEqual(list(cycle('')), [])
447        self.assertRaises(TypeError, cycle)
448        self.assertRaises(TypeError, cycle, 5)
449        self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
450
451    def test_groupby(self):
452        # Check whether it accepts arguments correctly
453        self.assertEqual([], list(groupby([])))
454        self.assertEqual([], list(groupby([], key=id)))
455        self.assertRaises(TypeError, list, groupby('abc', []))
456        self.assertRaises(TypeError, groupby, None)
457        self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
458
459        # Check normal input
460        s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
461             (2,15,22), (3,16,23), (3,17,23)]
462        dup = []
463        for k, g in groupby(s, lambda r:r[0]):
464            for elem in g:
465                self.assertEqual(k, elem[0])
466                dup.append(elem)
467        self.assertEqual(s, dup)
468
469        # Check nested case
470        dup = []
471        for k, g in groupby(s, lambda r:r[0]):
472            for ik, ig in groupby(g, lambda r:r[2]):
473                for elem in ig:
474                    self.assertEqual(k, elem[0])
475                    self.assertEqual(ik, elem[2])
476                    dup.append(elem)
477        self.assertEqual(s, dup)
478
479        # Check case where inner iterator is not used
480        keys = [k for k, g in groupby(s, lambda r:r[0])]
481        expectedkeys = set([r[0] for r in s])
482        self.assertEqual(set(keys), expectedkeys)
483        self.assertEqual(len(keys), len(expectedkeys))
484
485        # Exercise pipes and filters style
486        s = 'abracadabra'
487        # sort s | uniq
488        r = [k for k, g in groupby(sorted(s))]
489        self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
490        # sort s | uniq -d
491        r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
492        self.assertEqual(r, ['a', 'b', 'r'])
493        # sort s | uniq -c
494        r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
495        self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
496        # sort s | uniq -c | sort -rn | head -3
497        r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
498        self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
499
500        # iter.next failure
501        class ExpectedError(Exception):
502            pass
503        def delayed_raise(n=0):
504            for i in range(n):
505                yield 'yo'
506            raise ExpectedError
507        def gulp(iterable, keyp=None, func=list):
508            return [func(g) for k, g in groupby(iterable, keyp)]
509
510        # iter.next failure on outer object
511        self.assertRaises(ExpectedError, gulp, delayed_raise(0))
512        # iter.next failure on inner object
513        self.assertRaises(ExpectedError, gulp, delayed_raise(1))
514
515        # __cmp__ failure
516        class DummyCmp:
517            def __cmp__(self, dst):
518                raise ExpectedError
519        s = [DummyCmp(), DummyCmp(), None]
520
521        # __cmp__ failure on outer object
522        self.assertRaises(ExpectedError, gulp, s, func=id)
523        # __cmp__ failure on inner object
524        self.assertRaises(ExpectedError, gulp, s)
525
526        # keyfunc failure
527        def keyfunc(obj):
528            if keyfunc.skip > 0:
529                keyfunc.skip -= 1
530                return obj
531            else:
532                raise ExpectedError
533
534        # keyfunc failure on outer object
535        keyfunc.skip = 0
536        self.assertRaises(ExpectedError, gulp, [None], keyfunc)
537        keyfunc.skip = 1
538        self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
539
540    def test_ifilter(self):
541        self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
542        self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
543        self.assertEqual(list(ifilter(bool, [0,1,0,2,0])), [1,2])
544        self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
545        self.assertRaises(TypeError, ifilter)
546        self.assertRaises(TypeError, ifilter, lambda x:x)
547        self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
548        self.assertRaises(TypeError, ifilter, isEven, 3)
549        self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
550
551    def test_ifilterfalse(self):
552        self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
553        self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
554        self.assertEqual(list(ifilterfalse(bool, [0,1,0,2,0])), [0,0,0])
555        self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
556        self.assertRaises(TypeError, ifilterfalse)
557        self.assertRaises(TypeError, ifilterfalse, lambda x:x)
558        self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
559        self.assertRaises(TypeError, ifilterfalse, isEven, 3)
560        self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
561
562    def test_izip(self):
563        ans = [(x,y) for x, y in izip('abc',count())]
564        self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
565        self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
566        self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
567        self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
568        self.assertEqual(list(izip('abcdef')), zip('abcdef'))
569        self.assertEqual(list(izip()), zip())
570        self.assertRaises(TypeError, izip, 3)
571        self.assertRaises(TypeError, izip, range(3), 3)
572        self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
573                         zip('abc', 'def'))
574        self.assertEqual([pair for pair in izip('abc', 'def')],
575                         zip('abc', 'def'))
576
577    @test_support.impl_detail("tuple reuse is specific to CPython")
578    def test_izip_tuple_reuse(self):
579        ids = map(id, izip('abc', 'def'))
580        self.assertEqual(min(ids), max(ids))
581        ids = map(id, list(izip('abc', 'def')))
582        self.assertEqual(len(dict.fromkeys(ids)), len(ids))
583
584    def test_iziplongest(self):
585        for args in [
586                ['abc', range(6)],
587                [range(6), 'abc'],
588                [range(1000), range(2000,2100), range(3000,3050)],
589                [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
590                [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
591            ]:
592            # target = map(None, *args) <- this raises a py3k warning
593            # this is the replacement:
594            target = [tuple([arg[i] if i < len(arg) else None for arg in args])
595                      for i in range(max(map(len, args)))]
596            self.assertEqual(list(izip_longest(*args)), target)
597            self.assertEqual(list(izip_longest(*args, **{})), target)
598            target = [tuple((e is None and 'X' or e) for e in t) for t in target]   # Replace None fills with 'X'
599            self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
600
601        self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
602
603        self.assertEqual(list(izip_longest()), zip())
604        self.assertEqual(list(izip_longest([])), zip([]))
605        self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
606
607        self.assertEqual(list(izip_longest('abc', 'defg', **{})),
608                         zip(list('abc') + [None], 'defg'))  # empty keyword dict
609        self.assertRaises(TypeError, izip_longest, 3)
610        self.assertRaises(TypeError, izip_longest, range(3), 3)
611
612        for stmt in [
613            "izip_longest('abc', fv=1)",
614            "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
615        ]:
616            try:
617                eval(stmt, globals(), locals())
618            except TypeError:
619                pass
620            else:
621                self.fail('Did not raise Type in:  ' + stmt)
622
623        self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
624                         zip('abc', 'def'))
625        self.assertEqual([pair for pair in izip_longest('abc', 'def')],
626                         zip('abc', 'def'))
627
628    @test_support.impl_detail("tuple reuse is specific to CPython")
629    def test_izip_longest_tuple_reuse(self):
630        ids = map(id, izip_longest('abc', 'def'))
631        self.assertEqual(min(ids), max(ids))
632        ids = map(id, list(izip_longest('abc', 'def')))
633        self.assertEqual(len(dict.fromkeys(ids)), len(ids))
634
635    def test_bug_7244(self):
636
637        class Repeater(object):
638            # this class is similar to itertools.repeat
639            def __init__(self, o, t, e):
640                self.o = o
641                self.t = int(t)
642                self.e = e
643            def __iter__(self): # its iterator is itself
644                return self
645            def next(self):
646                if self.t > 0:
647                    self.t -= 1
648                    return self.o
649                else:
650                    raise self.e
651
652        # Formerly this code in would fail in debug mode
653        # with Undetected Error and Stop Iteration
654        r1 = Repeater(1, 3, StopIteration)
655        r2 = Repeater(2, 4, StopIteration)
656        def run(r1, r2):
657            result = []
658            for i, j in izip_longest(r1, r2, fillvalue=0):
659                with test_support.captured_output('stdout'):
660                    print (i, j)
661                result.append((i, j))
662            return result
663        self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
664
665        # Formerly, the RuntimeError would be lost
666        # and StopIteration would stop as expected
667        r1 = Repeater(1, 3, RuntimeError)
668        r2 = Repeater(2, 4, StopIteration)
669        it = izip_longest(r1, r2, fillvalue=0)
670        self.assertEqual(next(it), (1, 2))
671        self.assertEqual(next(it), (1, 2))
672        self.assertEqual(next(it), (1, 2))
673        self.assertRaises(RuntimeError, next, it)
674
675    def test_product(self):
676        for args, result in [
677            ([], [()]),                     # zero iterables
678            (['ab'], [('a',), ('b',)]),     # one iterable
679            ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]),     # two iterables
680            ([range(0), range(2), range(3)], []),           # first iterable with zero length
681            ([range(2), range(0), range(3)], []),           # middle iterable with zero length
682            ([range(2), range(3), range(0)], []),           # last iterable with zero length
683            ]:
684            self.assertEqual(list(product(*args)), result)
685            for r in range(4):
686                self.assertEqual(list(product(*(args*r))),
687                                 list(product(*args, **dict(repeat=r))))
688        self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
689        self.assertRaises(TypeError, product, range(6), None)
690
691        def product1(*args, **kwds):
692            pools = map(tuple, args) * kwds.get('repeat', 1)
693            n = len(pools)
694            if n == 0:
695                yield ()
696                return
697            if any(len(pool) == 0 for pool in pools):
698                return
699            indices = [0] * n
700            yield tuple(pool[i] for pool, i in zip(pools, indices))
701            while 1:
702                for i in reversed(range(n)):  # right to left
703                    if indices[i] == len(pools[i]) - 1:
704                        continue
705                    indices[i] += 1
706                    for j in range(i+1, n):
707                        indices[j] = 0
708                    yield tuple(pool[i] for pool, i in zip(pools, indices))
709                    break
710                else:
711                    return
712
713        def product2(*args, **kwds):
714            'Pure python version used in docs'
715            pools = map(tuple, args) * kwds.get('repeat', 1)
716            result = [[]]
717            for pool in pools:
718                result = [x+[y] for x in result for y in pool]
719            for prod in result:
720                yield tuple(prod)
721
722        argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
723                    set('abcdefg'), range(11), tuple(range(13))]
724        for i in range(100):
725            args = [random.choice(argtypes) for j in range(random.randrange(5))]
726            expected_len = prod(map(len, args))
727            self.assertEqual(len(list(product(*args))), expected_len)
728            self.assertEqual(list(product(*args)), list(product1(*args)))
729            self.assertEqual(list(product(*args)), list(product2(*args)))
730            args = map(iter, args)
731            self.assertEqual(len(list(product(*args))), expected_len)
732
733    @test_support.bigaddrspacetest
734    def test_product_overflow(self):
735        with self.assertRaises((OverflowError, MemoryError)):
736            product(*(['ab']*2**5), repeat=2**25)
737
738    @test_support.impl_detail("tuple reuse is specific to CPython")
739    def test_product_tuple_reuse(self):
740        self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
741        self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
742
743    def test_repeat(self):
744        self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
745        self.assertEqual(list(repeat(object='a', times=0)), [])
746        self.assertEqual(list(repeat(object='a', times=-1)), [])
747        self.assertEqual(list(repeat(object='a', times=-2)), [])
748        self.assertEqual(zip(xrange(3),repeat('a')),
749                         [(0, 'a'), (1, 'a'), (2, 'a')])
750        self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
751        self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
752        self.assertEqual(list(repeat('a', 0)), [])
753        self.assertEqual(list(repeat('a', -3)), [])
754        self.assertRaises(TypeError, repeat)
755        self.assertRaises(TypeError, repeat, None, 3, 4)
756        self.assertRaises(TypeError, repeat, None, 'a')
757        r = repeat(1+0j)
758        self.assertEqual(repr(r), 'repeat((1+0j))')
759        r = repeat(1+0j, 5)
760        self.assertEqual(repr(r), 'repeat((1+0j), 5)')
761        list(r)
762        self.assertEqual(repr(r), 'repeat((1+0j), 0)')
763
764    def test_repeat_with_negative_times(self):
765        self.assertEqual(repr(repeat('a', -1)), "repeat('a', 0)")
766        self.assertEqual(repr(repeat('a', -2)), "repeat('a', 0)")
767        self.assertEqual(repr(repeat('a', times=-1)), "repeat('a', 0)")
768        self.assertEqual(repr(repeat('a', times=-2)), "repeat('a', 0)")
769
770    def test_imap(self):
771        self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
772                         [0**1, 1**2, 2**3])
773        self.assertEqual(list(imap(None, 'abc', range(5))),
774                         [('a',0),('b',1),('c',2)])
775        self.assertEqual(list(imap(None, 'abc', count())),
776                         [('a',0),('b',1),('c',2)])
777        self.assertEqual(take(2,imap(None, 'abc', count())),
778                         [('a',0),('b',1)])
779        self.assertEqual(list(imap(operator.pow, [])), [])
780        self.assertRaises(TypeError, imap)
781        self.assertRaises(TypeError, imap, operator.neg)
782        self.assertRaises(TypeError, imap(10, range(5)).next)
783        self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
784        self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
785
786    def test_starmap(self):
787        self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
788                         [0**1, 1**2, 2**3])
789        self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
790                         [0**1, 1**2, 2**3])
791        self.assertEqual(list(starmap(operator.pow, [])), [])
792        self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
793        self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
794        self.assertRaises(TypeError, starmap)
795        self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
796        self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
797        self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
798        self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
799
800    def test_islice(self):
801        for args in [          # islice(args) should agree with range(args)
802                (10, 20, 3),
803                (10, 3, 20),
804                (10, 20),
805                (10, 10),
806                (10, 3),
807                (20,)
808                ]:
809            self.assertEqual(list(islice(xrange(100), *args)), range(*args))
810
811        for args, tgtargs in [  # Stop when seqn is exhausted
812                ((10, 110, 3), ((10, 100, 3))),
813                ((10, 110), ((10, 100))),
814                ((110,), (100,))
815                ]:
816            self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
817
818        # Test stop=None
819        self.assertEqual(list(islice(xrange(10), None)), range(10))
820        self.assertEqual(list(islice(xrange(10), None, None)), range(10))
821        self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
822        self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
823        self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
824
825        # Test number of items consumed     SF #1171417
826        it = iter(range(10))
827        self.assertEqual(list(islice(it, 3)), range(3))
828        self.assertEqual(list(it), range(3, 10))
829
830        it = iter(range(10))
831        self.assertEqual(list(islice(it, 3, 3)), [])
832        self.assertEqual(list(it), range(3, 10))
833
834        # Test invalid arguments
835        self.assertRaises(TypeError, islice, xrange(10))
836        self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
837        self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
838        self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
839        self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
840        self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
841        self.assertRaises(ValueError, islice, xrange(10), 'a')
842        self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
843        self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
844        self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
845        self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
846        self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
847
848        # Issue #10323:  Less islice in a predictable state
849        c = count()
850        self.assertEqual(list(islice(c, 1, 3, 50)), [1])
851        self.assertEqual(next(c), 3)
852
853        # Issue #21321: check source iterator is not referenced
854        # from islice() after the latter has been exhausted
855        it = (x for x in (1, 2))
856        wr = weakref.ref(it)
857        it = islice(it, 1)
858        self.assertIsNotNone(wr())
859        list(it) # exhaust the iterator
860        test_support.gc_collect()
861        self.assertIsNone(wr())
862
863    def test_takewhile(self):
864        data = [1, 3, 5, 20, 2, 4, 6, 8]
865        underten = lambda x: x<10
866        self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
867        self.assertEqual(list(takewhile(underten, [])), [])
868        self.assertRaises(TypeError, takewhile)
869        self.assertRaises(TypeError, takewhile, operator.pow)
870        self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
871        self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
872        self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
873        t = takewhile(bool, [1, 1, 1, 0, 0, 0])
874        self.assertEqual(list(t), [1, 1, 1])
875        self.assertRaises(StopIteration, t.next)
876
877    def test_dropwhile(self):
878        data = [1, 3, 5, 20, 2, 4, 6, 8]
879        underten = lambda x: x<10
880        self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
881        self.assertEqual(list(dropwhile(underten, [])), [])
882        self.assertRaises(TypeError, dropwhile)
883        self.assertRaises(TypeError, dropwhile, operator.pow)
884        self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
885        self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
886        self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
887
888    def test_tee(self):
889        n = 200
890        def irange(n):
891            for i in xrange(n):
892                yield i
893
894        a, b = tee([])        # test empty iterator
895        self.assertEqual(list(a), [])
896        self.assertEqual(list(b), [])
897
898        a, b = tee(irange(n)) # test 100% interleaved
899        self.assertEqual(zip(a,b), zip(range(n),range(n)))
900
901        a, b = tee(irange(n)) # test 0% interleaved
902        self.assertEqual(list(a), range(n))
903        self.assertEqual(list(b), range(n))
904
905        a, b = tee(irange(n)) # test dealloc of leading iterator
906        for i in xrange(100):
907            self.assertEqual(a.next(), i)
908        del a
909        self.assertEqual(list(b), range(n))
910
911        a, b = tee(irange(n)) # test dealloc of trailing iterator
912        for i in xrange(100):
913            self.assertEqual(a.next(), i)
914        del b
915        self.assertEqual(list(a), range(100, n))
916
917        for j in xrange(5):   # test randomly interleaved
918            order = [0]*n + [1]*n
919            random.shuffle(order)
920            lists = ([], [])
921            its = tee(irange(n))
922            for i in order:
923                value = its[i].next()
924                lists[i].append(value)
925            self.assertEqual(lists[0], range(n))
926            self.assertEqual(lists[1], range(n))
927
928        # test argument format checking
929        self.assertRaises(TypeError, tee)
930        self.assertRaises(TypeError, tee, 3)
931        self.assertRaises(TypeError, tee, [1,2], 'x')
932        self.assertRaises(TypeError, tee, [1,2], 3, 'x')
933
934        # tee object should be instantiable
935        a, b = tee('abc')
936        c = type(a)('def')
937        self.assertEqual(list(c), list('def'))
938
939        # test long-lagged and multi-way split
940        a, b, c = tee(xrange(2000), 3)
941        for i in xrange(100):
942            self.assertEqual(a.next(), i)
943        self.assertEqual(list(b), range(2000))
944        self.assertEqual([c.next(), c.next()], range(2))
945        self.assertEqual(list(a), range(100,2000))
946        self.assertEqual(list(c), range(2,2000))
947
948        # test values of n
949        self.assertRaises(TypeError, tee, 'abc', 'invalid')
950        self.assertRaises(ValueError, tee, [], -1)
951        for n in xrange(5):
952            result = tee('abc', n)
953            self.assertEqual(type(result), tuple)
954            self.assertEqual(len(result), n)
955            self.assertEqual(map(list, result), [list('abc')]*n)
956
957        # tee pass-through to copyable iterator
958        a, b = tee('abc')
959        c, d = tee(a)
960        self.assertTrue(a is c)
961
962        # test tee_new
963        t1, t2 = tee('abc')
964        tnew = type(t1)
965        self.assertRaises(TypeError, tnew)
966        self.assertRaises(TypeError, tnew, 10)
967        t3 = tnew(t1)
968        self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
969
970        # test that tee objects are weak referencable
971        a, b = tee(xrange(10))
972        p = weakref.proxy(a)
973        self.assertEqual(getattr(p, '__class__'), type(b))
974        del a
975        self.assertRaises(ReferenceError, getattr, p, '__class__')
976
977    # Issue 13454: Crash when deleting backward iterator from tee()
978    def test_tee_del_backward(self):
979        forward, backward = tee(repeat(None, 20000000))
980        try:
981            any(forward)  # exhaust the iterator
982            del backward
983        except:
984            del forward, backward
985            raise
986
987    def test_StopIteration(self):
988        self.assertRaises(StopIteration, izip().next)
989
990        for f in (chain, cycle, izip, groupby):
991            self.assertRaises(StopIteration, f([]).next)
992            self.assertRaises(StopIteration, f(StopNow()).next)
993
994        self.assertRaises(StopIteration, islice([], None).next)
995        self.assertRaises(StopIteration, islice(StopNow(), None).next)
996
997        p, q = tee([])
998        self.assertRaises(StopIteration, p.next)
999        self.assertRaises(StopIteration, q.next)
1000        p, q = tee(StopNow())
1001        self.assertRaises(StopIteration, p.next)
1002        self.assertRaises(StopIteration, q.next)
1003
1004        self.assertRaises(StopIteration, repeat(None, 0).next)
1005
1006        for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
1007            self.assertRaises(StopIteration, f(lambda x:x, []).next)
1008            self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
1009
1010class TestExamples(unittest.TestCase):
1011
1012    def test_chain(self):
1013        self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
1014
1015    def test_chain_from_iterable(self):
1016        self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
1017
1018    def test_combinations(self):
1019        self.assertEqual(list(combinations('ABCD', 2)),
1020                         [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
1021        self.assertEqual(list(combinations(range(4), 3)),
1022                         [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
1023
1024    def test_combinations_with_replacement(self):
1025        self.assertEqual(list(combinations_with_replacement('ABC', 2)),
1026                         [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
1027
1028    def test_compress(self):
1029        self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
1030
1031    def test_count(self):
1032        self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
1033
1034    def test_cycle(self):
1035        self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
1036
1037    def test_dropwhile(self):
1038        self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
1039
1040    def test_groupby(self):
1041        self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
1042                         list('ABCDAB'))
1043        self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
1044                         [list('AAAA'), list('BBB'), list('CC'), list('D')])
1045
1046    def test_ifilter(self):
1047        self.assertEqual(list(ifilter(lambda x: x%2, range(10))), [1,3,5,7,9])
1048
1049    def test_ifilterfalse(self):
1050        self.assertEqual(list(ifilterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
1051
1052    def test_imap(self):
1053        self.assertEqual(list(imap(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
1054
1055    def test_islice(self):
1056        self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
1057        self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
1058        self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
1059        self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1060
1061    def test_izip(self):
1062        self.assertEqual(list(izip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
1063
1064    def test_izip_longest(self):
1065        self.assertEqual(list(izip_longest('ABCD', 'xy', fillvalue='-')),
1066                         [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
1067
1068    def test_permutations(self):
1069        self.assertEqual(list(permutations('ABCD', 2)),
1070                         map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split()))
1071        self.assertEqual(list(permutations(range(3))),
1072                         [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
1073
1074    def test_product(self):
1075        self.assertEqual(list(product('ABCD', 'xy')),
1076                         map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split()))
1077        self.assertEqual(list(product(range(2), repeat=3)),
1078                        [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
1079                         (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
1080
1081    def test_repeat(self):
1082        self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
1083
1084    def test_stapmap(self):
1085        self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
1086                         [32, 9, 1000])
1087
1088    def test_takewhile(self):
1089        self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
1090
1091
1092class TestPurePythonRoughEquivalents(unittest.TestCase):
1093
1094    @staticmethod
1095    def islice(iterable, *args):
1096        s = slice(*args)
1097        start, stop, step = s.start or 0, s.stop or sys.maxint, s.step or 1
1098        it = iter(xrange(start, stop, step))
1099        try:
1100            nexti = next(it)
1101        except StopIteration:
1102            # Consume *iterable* up to the *start* position.
1103            for i, element in izip(xrange(start), iterable):
1104                pass
1105            return
1106        try:
1107            for i, element in enumerate(iterable):
1108                if i == nexti:
1109                    yield element
1110                    nexti = next(it)
1111        except StopIteration:
1112            # Consume to *stop*.
1113            for i, element in izip(xrange(i + 1, stop), iterable):
1114                pass
1115
1116    def test_islice_recipe(self):
1117        self.assertEqual(list(self.islice('ABCDEFG', 2)), list('AB'))
1118        self.assertEqual(list(self.islice('ABCDEFG', 2, 4)), list('CD'))
1119        self.assertEqual(list(self.islice('ABCDEFG', 2, None)), list('CDEFG'))
1120        self.assertEqual(list(self.islice('ABCDEFG', 0, None, 2)), list('ACEG'))
1121        # Test items consumed.
1122        it = iter(xrange(10))
1123        self.assertEqual(list(self.islice(it, 3)), range(3))
1124        self.assertEqual(list(it), range(3, 10))
1125        it = iter(xrange(10))
1126        self.assertEqual(list(self.islice(it, 3, 3)), [])
1127        self.assertEqual(list(it), range(3, 10))
1128        # Test that slice finishes in predictable state.
1129        c = count()
1130        self.assertEqual(list(self.islice(c, 1, 3, 50)), [1])
1131        self.assertEqual(next(c), 3)
1132
1133
1134class TestGC(unittest.TestCase):
1135
1136    def makecycle(self, iterator, container):
1137        container.append(iterator)
1138        iterator.next()
1139        del container, iterator
1140
1141    def test_chain(self):
1142        a = []
1143        self.makecycle(chain(a), a)
1144
1145    def test_chain_from_iterable(self):
1146        a = []
1147        self.makecycle(chain.from_iterable([a]), a)
1148
1149    def test_combinations(self):
1150        a = []
1151        self.makecycle(combinations([1,2,a,3], 3), a)
1152
1153    def test_combinations_with_replacement(self):
1154        a = []
1155        self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1156
1157    def test_compress(self):
1158        a = []
1159        self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1160
1161    def test_count(self):
1162        a = []
1163        Int = type('Int', (int,), dict(x=a))
1164        self.makecycle(count(Int(0), Int(1)), a)
1165
1166    def test_cycle(self):
1167        a = []
1168        self.makecycle(cycle([a]*2), a)
1169
1170    def test_dropwhile(self):
1171        a = []
1172        self.makecycle(dropwhile(bool, [0, a, a]), a)
1173
1174    def test_groupby(self):
1175        a = []
1176        self.makecycle(groupby([a]*2, lambda x:x), a)
1177
1178    def test_issue2246(self):
1179        # Issue 2246 -- the _grouper iterator was not included in GC
1180        n = 10
1181        keyfunc = lambda x: x
1182        for i, j in groupby(xrange(n), key=keyfunc):
1183            keyfunc.__dict__.setdefault('x',[]).append(j)
1184
1185    def test_ifilter(self):
1186        a = []
1187        self.makecycle(ifilter(lambda x:True, [a]*2), a)
1188
1189    def test_ifilterfalse(self):
1190        a = []
1191        self.makecycle(ifilterfalse(lambda x:False, a), a)
1192
1193    def test_izip(self):
1194        a = []
1195        self.makecycle(izip([a]*2, [a]*3), a)
1196
1197    def test_izip_longest(self):
1198        a = []
1199        self.makecycle(izip_longest([a]*2, [a]*3), a)
1200        b = [a, None]
1201        self.makecycle(izip_longest([a]*2, [a]*3, fillvalue=b), a)
1202
1203    def test_imap(self):
1204        a = []
1205        self.makecycle(imap(lambda x:x, [a]*2), a)
1206
1207    def test_islice(self):
1208        a = []
1209        self.makecycle(islice([a]*2, None), a)
1210
1211    def test_permutations(self):
1212        a = []
1213        self.makecycle(permutations([1,2,a,3], 3), a)
1214
1215    def test_product(self):
1216        a = []
1217        self.makecycle(product([1,2,a,3], repeat=3), a)
1218
1219    def test_repeat(self):
1220        a = []
1221        self.makecycle(repeat(a), a)
1222
1223    def test_starmap(self):
1224        a = []
1225        self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1226
1227    def test_takewhile(self):
1228        a = []
1229        self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1230
1231def R(seqn):
1232    'Regular generator'
1233    for i in seqn:
1234        yield i
1235
1236class G:
1237    'Sequence using __getitem__'
1238    def __init__(self, seqn):
1239        self.seqn = seqn
1240    def __getitem__(self, i):
1241        return self.seqn[i]
1242
1243class I:
1244    'Sequence using iterator protocol'
1245    def __init__(self, seqn):
1246        self.seqn = seqn
1247        self.i = 0
1248    def __iter__(self):
1249        return self
1250    def next(self):
1251        if self.i >= len(self.seqn): raise StopIteration
1252        v = self.seqn[self.i]
1253        self.i += 1
1254        return v
1255
1256class Ig:
1257    'Sequence using iterator protocol defined with a generator'
1258    def __init__(self, seqn):
1259        self.seqn = seqn
1260        self.i = 0
1261    def __iter__(self):
1262        for val in self.seqn:
1263            yield val
1264
1265class X:
1266    'Missing __getitem__ and __iter__'
1267    def __init__(self, seqn):
1268        self.seqn = seqn
1269        self.i = 0
1270    def next(self):
1271        if self.i >= len(self.seqn): raise StopIteration
1272        v = self.seqn[self.i]
1273        self.i += 1
1274        return v
1275
1276class N:
1277    'Iterator missing next()'
1278    def __init__(self, seqn):
1279        self.seqn = seqn
1280        self.i = 0
1281    def __iter__(self):
1282        return self
1283
1284class E:
1285    'Test propagation of exceptions'
1286    def __init__(self, seqn):
1287        self.seqn = seqn
1288        self.i = 0
1289    def __iter__(self):
1290        return self
1291    def next(self):
1292        3 // 0
1293
1294class S:
1295    'Test immediate stop'
1296    def __init__(self, seqn):
1297        pass
1298    def __iter__(self):
1299        return self
1300    def next(self):
1301        raise StopIteration
1302
1303def L(seqn):
1304    'Test multiple tiers of iterators'
1305    return chain(imap(lambda x:x, R(Ig(G(seqn)))))
1306
1307
1308class TestVariousIteratorArgs(unittest.TestCase):
1309
1310    def test_chain(self):
1311        for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1312            for g in (G, I, Ig, S, L, R):
1313                self.assertEqual(list(chain(g(s))), list(g(s)))
1314                self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
1315            self.assertRaises(TypeError, list, chain(X(s)))
1316            self.assertRaises(TypeError, list, chain(N(s)))
1317            self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1318
1319    def test_compress(self):
1320        for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1321            n = len(s)
1322            for g in (G, I, Ig, S, L, R):
1323                self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1324            self.assertRaises(TypeError, compress, X(s), repeat(1))
1325            self.assertRaises(TypeError, list, compress(N(s), repeat(1)))
1326            self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1327
1328    def test_product(self):
1329        for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1330            self.assertRaises(TypeError, product, X(s))
1331            self.assertRaises(TypeError, product, N(s))
1332            self.assertRaises(ZeroDivisionError, product, E(s))
1333
1334    def test_cycle(self):
1335        for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1336            for g in (G, I, Ig, S, L, R):
1337                tgtlen = len(s) * 3
1338                expected = list(g(s))*3
1339                actual = list(islice(cycle(g(s)), tgtlen))
1340                self.assertEqual(actual, expected)
1341            self.assertRaises(TypeError, cycle, X(s))
1342            self.assertRaises(TypeError, list, cycle(N(s)))
1343            self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1344
1345    def test_groupby(self):
1346        for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1347            for g in (G, I, Ig, S, L, R):
1348                self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1349            self.assertRaises(TypeError, groupby, X(s))
1350            self.assertRaises(TypeError, list, groupby(N(s)))
1351            self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1352
1353    def test_ifilter(self):
1354        for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1355            for g in (G, I, Ig, S, L, R):
1356                self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
1357            self.assertRaises(TypeError, ifilter, isEven, X(s))
1358            self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
1359            self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
1360
1361    def test_ifilterfalse(self):
1362        for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1363            for g in (G, I, Ig, S, L, R):
1364                self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
1365            self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
1366            self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
1367            self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
1368
1369    def test_izip(self):
1370        for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1371            for g in (G, I, Ig, S, L, R):
1372                self.assertEqual(list(izip(g(s))), zip(g(s)))
1373                self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
1374            self.assertRaises(TypeError, izip, X(s))
1375            self.assertRaises(TypeError, list, izip(N(s)))
1376            self.assertRaises(ZeroDivisionError, list, izip(E(s)))
1377
1378    def test_iziplongest(self):
1379        for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1380            for g in (G, I, Ig, S, L, R):
1381                self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
1382                self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
1383            self.assertRaises(TypeError, izip_longest, X(s))
1384            self.assertRaises(TypeError, list, izip_longest(N(s)))
1385            self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
1386
1387    def test_imap(self):
1388        for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1389            for g in (G, I, Ig, S, L, R):
1390                self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
1391                self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
1392            self.assertRaises(TypeError, imap, onearg, X(s))
1393            self.assertRaises(TypeError, list, imap(onearg, N(s)))
1394            self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
1395
1396    def test_islice(self):
1397        for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1398            for g in (G, I, Ig, S, L, R):
1399                self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1400            self.assertRaises(TypeError, islice, X(s), 10)
1401            self.assertRaises(TypeError, list, islice(N(s), 10))
1402            self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1403
1404    def test_starmap(self):
1405        for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1406            for g in (G, I, Ig, S, L, R):
1407                ss = zip(s, s)
1408                self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
1409            self.assertRaises(TypeError, starmap, operator.pow, X(ss))
1410            self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
1411            self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1412
1413    def test_takewhile(self):
1414        for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1415            for g in (G, I, Ig, S, L, R):
1416                tgt = []
1417                for elem in g(s):
1418                    if not isEven(elem): break
1419                    tgt.append(elem)
1420                self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1421            self.assertRaises(TypeError, takewhile, isEven, X(s))
1422            self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
1423            self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1424
1425    def test_dropwhile(self):
1426        for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1427            for g in (G, I, Ig, S, L, R):
1428                tgt = []
1429                for elem in g(s):
1430                    if not tgt and isOdd(elem): continue
1431                    tgt.append(elem)
1432                self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1433            self.assertRaises(TypeError, dropwhile, isOdd, X(s))
1434            self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
1435            self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1436
1437    def test_tee(self):
1438        for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1439            for g in (G, I, Ig, S, L, R):
1440                it1, it2 = tee(g(s))
1441                self.assertEqual(list(it1), list(g(s)))
1442                self.assertEqual(list(it2), list(g(s)))
1443            self.assertRaises(TypeError, tee, X(s))
1444            self.assertRaises(TypeError, list, tee(N(s))[0])
1445            self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1446
1447class LengthTransparency(unittest.TestCase):
1448
1449    def test_repeat(self):
1450        from test.test_iterlen import len
1451        self.assertEqual(len(repeat(None, 50)), 50)
1452        self.assertRaises(TypeError, len, repeat(None))
1453
1454class RegressionTests(unittest.TestCase):
1455
1456    def test_sf_793826(self):
1457        # Fix Armin Rigo's successful efforts to wreak havoc
1458
1459        def mutatingtuple(tuple1, f, tuple2):
1460            # this builds a tuple t which is a copy of tuple1,
1461            # then calls f(t), then mutates t to be equal to tuple2
1462            # (needs len(tuple1) == len(tuple2)).
1463            def g(value, first=[1]):
1464                if first:
1465                    del first[:]
1466                    f(z.next())
1467                return value
1468            items = list(tuple2)
1469            items[1:1] = list(tuple1)
1470            gen = imap(g, items)
1471            z = izip(*[gen]*len(tuple1))
1472            z.next()
1473
1474        def f(t):
1475            global T
1476            T = t
1477            first[:] = list(T)
1478
1479        first = []
1480        mutatingtuple((1,2,3), f, (4,5,6))
1481        second = list(T)
1482        self.assertEqual(first, second)
1483
1484
1485    def test_sf_950057(self):
1486        # Make sure that chain() and cycle() catch exceptions immediately
1487        # rather than when shifting between input sources
1488
1489        def gen1():
1490            hist.append(0)
1491            yield 1
1492            hist.append(1)
1493            raise AssertionError
1494            hist.append(2)
1495
1496        def gen2(x):
1497            hist.append(3)
1498            yield 2
1499            hist.append(4)
1500            if x:
1501                raise StopIteration
1502
1503        hist = []
1504        self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1505        self.assertEqual(hist, [0,1])
1506
1507        hist = []
1508        self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1509        self.assertEqual(hist, [0,1])
1510
1511        hist = []
1512        self.assertRaises(AssertionError, list, cycle(gen1()))
1513        self.assertEqual(hist, [0,1])
1514
1515    def test_long_chain_of_empty_iterables(self):
1516        # Make sure itertools.chain doesn't run into recursion limits when
1517        # dealing with long chains of empty iterables. Even with a high
1518        # number this would probably only fail in Py_DEBUG mode.
1519        it = chain.from_iterable(() for unused in xrange(10000000))
1520        with self.assertRaises(StopIteration):
1521            next(it)
1522
1523    def test_issue30347_1(self):
1524        def f(n):
1525            if n == 5:
1526                list(b)
1527            return n != 6
1528        for (k, b) in groupby(range(10), f):
1529            list(b)  # shouldn't crash
1530
1531    def test_issue30347_2(self):
1532        class K(object):
1533            i = 0
1534            def __init__(self, v):
1535                pass
1536            def __eq__(self, other):
1537                K.i += 1
1538                if K.i == 1:
1539                    next(g, None)
1540                return True
1541            def __hash__(self):
1542                return 1
1543        g = next(groupby(range(10), K))[1]
1544        for j in range(2):
1545            next(g, None)  # shouldn't crash
1546
1547
1548class SubclassWithKwargsTest(unittest.TestCase):
1549    def test_keywords_in_subclass(self):
1550        # count is not subclassable...
1551        for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
1552                    starmap, islice, takewhile, dropwhile, cycle, compress):
1553            class Subclass(cls):
1554                def __init__(self, newarg=None, *args):
1555                    cls.__init__(self, *args)
1556            try:
1557                Subclass(newarg=1)
1558            except TypeError, err:
1559                # we expect type errors because of wrong argument count
1560                self.assertNotIn("does not take keyword arguments", err.args[0])
1561
1562
1563libreftest = """ Doctest for examples in the library reference: libitertools.tex
1564
1565
1566>>> amounts = [120.15, 764.05, 823.14]
1567>>> for checknum, amount in izip(count(1200), amounts):
1568...     print 'Check %d is for $%.2f' % (checknum, amount)
1569...
1570Check 1200 is for $120.15
1571Check 1201 is for $764.05
1572Check 1202 is for $823.14
1573
1574>>> import operator
1575>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1576...    print cube
1577...
15781
15798
158027
1581
1582>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
1583>>> for name in islice(reportlines, 3, None, 2):
1584...    print name.title()
1585...
1586Alex
1587Laura
1588Martin
1589Walter
1590Samuele
1591
1592>>> from operator import itemgetter
1593>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
1594>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
1595>>> for k, g in groupby(di, itemgetter(1)):
1596...     print k, map(itemgetter(0), g)
1597...
15981 ['a', 'c', 'e']
15992 ['b', 'd', 'f']
16003 ['g']
1601
1602# Find runs of consecutive numbers using groupby.  The key to the solution
1603# is differencing with a range so that consecutive numbers all appear in
1604# same group.
1605>>> data = [ 1,  4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
1606>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
1607...     print map(operator.itemgetter(1), g)
1608...
1609[1]
1610[4, 5, 6]
1611[10]
1612[15, 16, 17, 18]
1613[22]
1614[25, 26, 27, 28]
1615
1616>>> def take(n, iterable):
1617...     "Return first n items of the iterable as a list"
1618...     return list(islice(iterable, n))
1619
1620>>> def enumerate(iterable, start=0):
1621...     return izip(count(start), iterable)
1622
1623>>> def tabulate(function, start=0):
1624...     "Return function(0), function(1), ..."
1625...     return imap(function, count(start))
1626
1627>>> import collections
1628>>> def consume(iterator, n=None):
1629...     "Advance the iterator n-steps ahead. If n is None, consume entirely."
1630...     # Use functions that consume iterators at C speed.
1631...     if n is None:
1632...         # feed the entire iterator into a zero-length deque
1633...         collections.deque(iterator, maxlen=0)
1634...     else:
1635...         # advance to the empty slice starting at position n
1636...         next(islice(iterator, n, n), None)
1637
1638>>> def nth(iterable, n, default=None):
1639...     "Returns the nth item or a default value"
1640...     return next(islice(iterable, n, None), default)
1641
1642>>> def all_equal(iterable):
1643...     "Returns True if all the elements are equal to each other"
1644...     g = groupby(iterable)
1645...     return next(g, True) and not next(g, False)
1646
1647>>> def quantify(iterable, pred=bool):
1648...     "Count how many times the predicate is true"
1649...     return sum(imap(pred, iterable))
1650
1651>>> def padnone(iterable):
1652...     "Returns the sequence elements and then returns None indefinitely"
1653...     return chain(iterable, repeat(None))
1654
1655>>> def ncycles(iterable, n):
1656...     "Returns the sequence elements n times"
1657...     return chain(*repeat(iterable, n))
1658
1659>>> def dotproduct(vec1, vec2):
1660...     return sum(imap(operator.mul, vec1, vec2))
1661
1662>>> def flatten(listOfLists):
1663...     return list(chain.from_iterable(listOfLists))
1664
1665>>> def repeatfunc(func, times=None, *args):
1666...     "Repeat calls to func with specified arguments."
1667...     "   Example:  repeatfunc(random.random)"
1668...     if times is None:
1669...         return starmap(func, repeat(args))
1670...     else:
1671...         return starmap(func, repeat(args, times))
1672
1673>>> def pairwise(iterable):
1674...     "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1675...     a, b = tee(iterable)
1676...     for elem in b:
1677...         break
1678...     return izip(a, b)
1679
1680>>> def grouper(n, iterable, fillvalue=None):
1681...     "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
1682...     args = [iter(iterable)] * n
1683...     return izip_longest(fillvalue=fillvalue, *args)
1684
1685>>> def roundrobin(*iterables):
1686...     "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
1687...     # Recipe credited to George Sakkis
1688...     pending = len(iterables)
1689...     nexts = cycle(iter(it).next for it in iterables)
1690...     while pending:
1691...         try:
1692...             for next in nexts:
1693...                 yield next()
1694...         except StopIteration:
1695...             pending -= 1
1696...             nexts = cycle(islice(nexts, pending))
1697
1698>>> def powerset(iterable):
1699...     "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1700...     s = list(iterable)
1701...     return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
1702
1703>>> def unique_everseen(iterable, key=None):
1704...     "List unique elements, preserving order. Remember all elements ever seen."
1705...     # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1706...     # unique_everseen('ABBCcAD', str.lower) --> A B C D
1707...     seen = set()
1708...     seen_add = seen.add
1709...     if key is None:
1710...         for element in iterable:
1711...             if element not in seen:
1712...                 seen_add(element)
1713...                 yield element
1714...     else:
1715...         for element in iterable:
1716...             k = key(element)
1717...             if k not in seen:
1718...                 seen_add(k)
1719...                 yield element
1720
1721>>> def unique_justseen(iterable, key=None):
1722...     "List unique elements, preserving order. Remember only the element just seen."
1723...     # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1724...     # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1725...     return imap(next, imap(itemgetter(1), groupby(iterable, key)))
1726
1727This is not part of the examples but it tests to make sure the definitions
1728perform as purported.
1729
1730>>> take(10, count())
1731[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1732
1733>>> list(enumerate('abc'))
1734[(0, 'a'), (1, 'b'), (2, 'c')]
1735
1736>>> list(islice(tabulate(lambda x: 2*x), 4))
1737[0, 2, 4, 6]
1738
1739>>> it = iter(xrange(10))
1740>>> consume(it, 3)
1741>>> next(it)
17423
1743>>> consume(it)
1744>>> next(it, 'Done')
1745'Done'
1746
1747>>> nth('abcde', 3)
1748'd'
1749
1750>>> nth('abcde', 9) is None
1751True
1752
1753>>> [all_equal(s) for s in ('', 'A', 'AAAA', 'AAAB', 'AAABA')]
1754[True, True, True, False, False]
1755
1756>>> quantify(xrange(99), lambda x: x%2==0)
175750
1758
1759>>> a = [[1, 2, 3], [4, 5, 6]]
1760>>> flatten(a)
1761[1, 2, 3, 4, 5, 6]
1762
1763>>> list(repeatfunc(pow, 5, 2, 3))
1764[8, 8, 8, 8, 8]
1765
1766>>> import random
1767>>> take(5, imap(int, repeatfunc(random.random)))
1768[0, 0, 0, 0, 0]
1769
1770>>> list(pairwise('abcd'))
1771[('a', 'b'), ('b', 'c'), ('c', 'd')]
1772
1773>>> list(pairwise([]))
1774[]
1775
1776>>> list(pairwise('a'))
1777[]
1778
1779>>> list(islice(padnone('abc'), 0, 6))
1780['a', 'b', 'c', None, None, None]
1781
1782>>> list(ncycles('abc', 3))
1783['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1784
1785>>> dotproduct([1,2,3], [4,5,6])
178632
1787
1788>>> list(grouper(3, 'abcdefg', 'x'))
1789[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1790
1791>>> list(roundrobin('abc', 'd', 'ef'))
1792['a', 'd', 'e', 'b', 'f', 'c']
1793
1794>>> list(powerset([1,2,3]))
1795[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
1796
1797>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1798True
1799
1800>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1801True
1802
1803>>> list(unique_everseen('AAAABBBCCDAABBB'))
1804['A', 'B', 'C', 'D']
1805
1806>>> list(unique_everseen('ABBCcAD', str.lower))
1807['A', 'B', 'C', 'D']
1808
1809>>> list(unique_justseen('AAAABBBCCDAABBB'))
1810['A', 'B', 'C', 'D', 'A', 'B']
1811
1812>>> list(unique_justseen('ABBCcAD', str.lower))
1813['A', 'B', 'C', 'A', 'D']
1814
1815"""
1816
1817__test__ = {'libreftest' : libreftest}
1818
1819def test_main(verbose=None):
1820    test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
1821                    RegressionTests, LengthTransparency,
1822                    SubclassWithKwargsTest, TestExamples,
1823                    TestPurePythonRoughEquivalents)
1824    test_support.run_unittest(*test_classes)
1825
1826    # verify reference counting
1827    if verbose and hasattr(sys, "gettotalrefcount"):
1828        import gc
1829        counts = [None] * 5
1830        for i in xrange(len(counts)):
1831            test_support.run_unittest(*test_classes)
1832            gc.collect()
1833            counts[i] = sys.gettotalrefcount()
1834        print counts
1835
1836    # doctest the examples in the library reference
1837    test_support.run_doctest(sys.modules[__name__], verbose)
1838
1839if __name__ == "__main__":
1840    test_main(verbose=True)
1841