blob: 2ba1e8ee4145ee41a5b7457c1f9bc9cf523b0049 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001import unittest
2from test import test_support
3from itertools import *
Raymond Hettingera9f60922004-10-17 16:40:14 +00004from weakref import proxy
Raymond Hettinger2012f172003-02-07 05:32:58 +00005import sys
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00006import operator
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00007import random
Raymond Hettinger673827c2009-11-30 11:15:28 +00008import copy
9import pickle
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +000010maxsize = test_support.MAX_Py_ssize_t
11minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000012
13def onearg(x):
14 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000015 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000016
17def errfunc(*args):
18 'Test function that raises an error'
19 raise ValueError
20
21def gen3():
22 'Non-restartable source sequence'
23 for i in (0, 1, 2):
24 yield i
25
26def isEven(x):
27 'Test predicate'
28 return x%2==0
29
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000030def isOdd(x):
31 'Test predicate'
32 return x%2==1
33
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000034class StopNow:
35 'Class emulating an empty iterable.'
36 def __iter__(self):
37 return self
38 def next(self):
39 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000040
Raymond Hettinger02420702003-06-29 20:36:23 +000041def take(n, seq):
42 'Convenience function for partially consuming a long of infinite iterable'
43 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000044
Raymond Hettingerd553d852008-03-04 04:17:08 +000045def prod(iterable):
46 return reduce(operator.mul, iterable, 1)
47
Raymond Hettinger93e804d2008-02-26 23:40:50 +000048def fact(n):
49 'Factorial'
Raymond Hettingerd553d852008-03-04 04:17:08 +000050 return prod(range(1, n+1))
51
Raymond Hettinger96ef8112003-02-01 00:10:11 +000052class TestBasicOps(unittest.TestCase):
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000053 def test_chain(self):
Raymond Hettingerad47fa12008-03-06 20:52:01 +000054
55 def chain2(*iterables):
56 'Pure python version in the docs'
57 for it in iterables:
58 for element in it:
59 yield element
60
61 for c in (chain, chain2):
62 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
63 self.assertEqual(list(c('abc')), list('abc'))
64 self.assertEqual(list(c('')), [])
65 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
66 self.assertRaises(TypeError, list,c(2, 3))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000067
Raymond Hettingerb4cbc982008-02-28 22:46:41 +000068 def test_chain_from_iterable(self):
69 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
70 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
71 self.assertEqual(list(chain.from_iterable([''])), [])
72 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
73 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
74
Raymond Hettinger93e804d2008-02-26 23:40:50 +000075 def test_combinations(self):
Raymond Hettinger825758c2009-01-08 05:20:19 +000076 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Raymond Hettinger93e804d2008-02-26 23:40:50 +000077 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Raymond Hettingerd553d852008-03-04 04:17:08 +000078 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Raymond Hettinger93e804d2008-02-26 23:40:50 +000079 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Raymond Hettinger825758c2009-01-08 05:20:19 +000080 self.assertEqual(list(combinations('abc', 32)), []) # r > n
Raymond Hettinger93e804d2008-02-26 23:40:50 +000081 self.assertEqual(list(combinations(range(4), 3)),
82 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
Raymond Hettingerd553d852008-03-04 04:17:08 +000083
84 def combinations1(iterable, r):
85 'Pure python version shown in the docs'
86 pool = tuple(iterable)
87 n = len(pool)
Raymond Hettinger825758c2009-01-08 05:20:19 +000088 if r > n:
89 return
Raymond Hettingerd553d852008-03-04 04:17:08 +000090 indices = range(r)
91 yield tuple(pool[i] for i in indices)
92 while 1:
93 for i in reversed(range(r)):
94 if indices[i] != i + n - r:
95 break
96 else:
97 return
98 indices[i] += 1
99 for j in range(i+1, r):
100 indices[j] = indices[j-1] + 1
101 yield tuple(pool[i] for i in indices)
102
103 def combinations2(iterable, r):
104 'Pure python version shown in the docs'
105 pool = tuple(iterable)
106 n = len(pool)
107 for indices in permutations(range(n), r):
108 if sorted(indices) == list(indices):
109 yield tuple(pool[i] for i in indices)
110
111 for n in range(7):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000112 values = [5*x-12 for x in range(n)]
Raymond Hettinger825758c2009-01-08 05:20:19 +0000113 for r in range(n+2):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000114 result = list(combinations(values, r))
Raymond Hettinger825758c2009-01-08 05:20:19 +0000115 self.assertEqual(len(result), 0 if r>n else fact(n) / fact(r) / fact(n-r)) # right number of combs
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000116 self.assertEqual(len(result), len(set(result))) # no repeats
117 self.assertEqual(result, sorted(result)) # lexicographic order
118 for c in result:
119 self.assertEqual(len(c), r) # r-length combinations
120 self.assertEqual(len(set(c)), r) # no duplicate elements
121 self.assertEqual(list(c), sorted(c)) # keep original ordering
122 self.assert_(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000123 self.assertEqual(list(c),
124 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Raymond Hettingerd553d852008-03-04 04:17:08 +0000125 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger825758c2009-01-08 05:20:19 +0000126 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000127
128 # Test implementation detail: tuple re-use
129 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
130 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
131
132 def test_permutations(self):
133 self.assertRaises(TypeError, permutations) # too few arguments
134 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000135 self.assertRaises(TypeError, permutations, None) # pool is not iterable
136 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger825758c2009-01-08 05:20:19 +0000137 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000138 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Raymond Hettingerd553d852008-03-04 04:17:08 +0000139 self.assertEqual(list(permutations(range(3), 2)),
140 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
141
142 def permutations1(iterable, r=None):
143 'Pure python version shown in the docs'
144 pool = tuple(iterable)
145 n = len(pool)
146 r = n if r is None else r
Raymond Hettinger825758c2009-01-08 05:20:19 +0000147 if r > n:
148 return
Raymond Hettingerd553d852008-03-04 04:17:08 +0000149 indices = range(n)
Raymond Hettingere70bb8d2008-03-23 00:55:46 +0000150 cycles = range(n, n-r, -1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000151 yield tuple(pool[i] for i in indices[:r])
152 while n:
153 for i in reversed(range(r)):
154 cycles[i] -= 1
155 if cycles[i] == 0:
156 indices[i:] = indices[i+1:] + indices[i:i+1]
157 cycles[i] = n - i
158 else:
159 j = cycles[i]
160 indices[i], indices[-j] = indices[-j], indices[i]
161 yield tuple(pool[i] for i in indices[:r])
162 break
163 else:
164 return
165
166 def permutations2(iterable, r=None):
167 'Pure python version shown in the docs'
168 pool = tuple(iterable)
169 n = len(pool)
170 r = n if r is None else r
171 for indices in product(range(n), repeat=r):
172 if len(set(indices)) == r:
173 yield tuple(pool[i] for i in indices)
174
175 for n in range(7):
176 values = [5*x-12 for x in range(n)]
Raymond Hettinger825758c2009-01-08 05:20:19 +0000177 for r in range(n+2):
Raymond Hettingerd553d852008-03-04 04:17:08 +0000178 result = list(permutations(values, r))
Raymond Hettinger825758c2009-01-08 05:20:19 +0000179 self.assertEqual(len(result), 0 if r>n else fact(n) / fact(n-r)) # right number of perms
Raymond Hettingerd553d852008-03-04 04:17:08 +0000180 self.assertEqual(len(result), len(set(result))) # no repeats
181 self.assertEqual(result, sorted(result)) # lexicographic order
182 for p in result:
183 self.assertEqual(len(p), r) # r-length permutations
184 self.assertEqual(len(set(p)), r) # no duplicate elements
185 self.assert_(all(e in values for e in p)) # elements taken from input iterable
186 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger825758c2009-01-08 05:20:19 +0000187 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000188 if r == n:
189 self.assertEqual(result, list(permutations(values, None))) # test r as None
190 self.assertEqual(result, list(permutations(values))) # test default r
191
192 # Test implementation detail: tuple re-use
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000193 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000194 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000195
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000196 def test_count(self):
197 self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
198 self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000199 self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000200 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
201 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000202 self.assertRaises(TypeError, count, 2, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000203 self.assertRaises(TypeError, count, 'a')
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000204 self.assertEqual(list(islice(count(maxsize-5), 10)), range(maxsize-5, maxsize+5))
205 self.assertEqual(list(islice(count(-maxsize-5), 10)), range(-maxsize-5, -maxsize+5))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000206 c = count(3)
207 self.assertEqual(repr(c), 'count(3)')
208 c.next()
209 self.assertEqual(repr(c), 'count(4)')
Jack Diederich36234e82006-09-21 17:50:26 +0000210 c = count(-9)
211 self.assertEqual(repr(c), 'count(-9)')
212 c.next()
213 self.assertEqual(c.next(), -8)
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000214 for i in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 10, sys.maxint-5, sys.maxint+5):
Raymond Hettinger3a0de082007-10-12 17:53:11 +0000215 # Test repr (ignoring the L in longs)
216 r1 = repr(count(i)).replace('L', '')
217 r2 = 'count(%r)'.__mod__(i).replace('L', '')
218 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000219
Raymond Hettinger673827c2009-11-30 11:15:28 +0000220 # check copy, deepcopy, pickle
221 for value in -3, 3, sys.maxint-5, sys.maxint+5:
222 c = count(value)
223 self.assertEqual(next(copy.copy(c)), value)
224 self.assertEqual(next(copy.deepcopy(c)), value)
225 self.assertEqual(next(pickle.loads(pickle.dumps(c))), value)
226
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000227 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000228 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000229 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000230 self.assertRaises(TypeError, cycle)
231 self.assertRaises(TypeError, cycle, 5)
232 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000233
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000234 def test_groupby(self):
235 # Check whether it accepts arguments correctly
236 self.assertEqual([], list(groupby([])))
237 self.assertEqual([], list(groupby([], key=id)))
238 self.assertRaises(TypeError, list, groupby('abc', []))
239 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000240 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000241
242 # Check normal input
243 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
244 (2,15,22), (3,16,23), (3,17,23)]
245 dup = []
246 for k, g in groupby(s, lambda r:r[0]):
247 for elem in g:
248 self.assertEqual(k, elem[0])
249 dup.append(elem)
250 self.assertEqual(s, dup)
251
252 # Check nested case
253 dup = []
254 for k, g in groupby(s, lambda r:r[0]):
255 for ik, ig in groupby(g, lambda r:r[2]):
256 for elem in ig:
257 self.assertEqual(k, elem[0])
258 self.assertEqual(ik, elem[2])
259 dup.append(elem)
260 self.assertEqual(s, dup)
261
262 # Check case where inner iterator is not used
263 keys = [k for k, g in groupby(s, lambda r:r[0])]
264 expectedkeys = set([r[0] for r in s])
265 self.assertEqual(set(keys), expectedkeys)
266 self.assertEqual(len(keys), len(expectedkeys))
267
268 # Exercise pipes and filters style
269 s = 'abracadabra'
270 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000271 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000272 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
273 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000274 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000275 self.assertEqual(r, ['a', 'b', 'r'])
276 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000277 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000278 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
279 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000280 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000281 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
282
283 # iter.next failure
284 class ExpectedError(Exception):
285 pass
286 def delayed_raise(n=0):
287 for i in range(n):
288 yield 'yo'
289 raise ExpectedError
290 def gulp(iterable, keyp=None, func=list):
291 return [func(g) for k, g in groupby(iterable, keyp)]
292
293 # iter.next failure on outer object
294 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
295 # iter.next failure on inner object
296 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
297
298 # __cmp__ failure
299 class DummyCmp:
300 def __cmp__(self, dst):
301 raise ExpectedError
302 s = [DummyCmp(), DummyCmp(), None]
303
304 # __cmp__ failure on outer object
305 self.assertRaises(ExpectedError, gulp, s, func=id)
306 # __cmp__ failure on inner object
307 self.assertRaises(ExpectedError, gulp, s)
308
309 # keyfunc failure
310 def keyfunc(obj):
311 if keyfunc.skip > 0:
312 keyfunc.skip -= 1
313 return obj
314 else:
315 raise ExpectedError
316
317 # keyfunc failure on outer object
318 keyfunc.skip = 0
319 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
320 keyfunc.skip = 1
321 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
322
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000323 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000324 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000325 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000326 self.assertEqual(list(ifilter(bool, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000327 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000328 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000329 self.assertRaises(TypeError, ifilter, lambda x:x)
330 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000331 self.assertRaises(TypeError, ifilter, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000332 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000333
334 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000335 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
336 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000337 self.assertEqual(list(ifilterfalse(bool, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000338 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000339 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000340 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
341 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000342 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000343 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000344
345 def test_izip(self):
346 ans = [(x,y) for x, y in izip('abc',count())]
347 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000348 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
349 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
Raymond Hettinger02420702003-06-29 20:36:23 +0000350 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000351 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000352 self.assertEqual(list(izip()), zip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000353 self.assertRaises(TypeError, izip, 3)
354 self.assertRaises(TypeError, izip, range(3), 3)
355 # Check tuple re-use (implementation detail)
356 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
357 zip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000358 self.assertEqual([pair for pair in izip('abc', 'def')],
359 zip('abc', 'def'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000360 ids = map(id, izip('abc', 'def'))
361 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000362 ids = map(id, list(izip('abc', 'def')))
363 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000364
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000365 def test_iziplongest(self):
366 for args in [
367 ['abc', range(6)],
368 [range(6), 'abc'],
369 [range(1000), range(2000,2100), range(3000,3050)],
370 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
371 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
372 ]:
373 target = map(None, *args)
374 self.assertEqual(list(izip_longest(*args)), target)
375 self.assertEqual(list(izip_longest(*args, **{})), target)
376 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
377 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Tim Petersea5962f2007-03-12 18:07:52 +0000378
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000379 self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
380
381 self.assertEqual(list(izip_longest()), zip())
382 self.assertEqual(list(izip_longest([])), zip([]))
383 self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
Tim Petersea5962f2007-03-12 18:07:52 +0000384
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000385 self.assertEqual(list(izip_longest('abc', 'defg', **{})), map(None, 'abc', 'defg')) # empty keyword dict
386 self.assertRaises(TypeError, izip_longest, 3)
387 self.assertRaises(TypeError, izip_longest, range(3), 3)
388
389 for stmt in [
390 "izip_longest('abc', fv=1)",
Tim Petersea5962f2007-03-12 18:07:52 +0000391 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000392 ]:
393 try:
394 eval(stmt, globals(), locals())
395 except TypeError:
396 pass
397 else:
398 self.fail('Did not raise Type in: ' + stmt)
Tim Petersea5962f2007-03-12 18:07:52 +0000399
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000400 # Check tuple re-use (implementation detail)
401 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
402 zip('abc', 'def'))
403 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
404 zip('abc', 'def'))
405 ids = map(id, izip_longest('abc', 'def'))
406 self.assertEqual(min(ids), max(ids))
407 ids = map(id, list(izip_longest('abc', 'def')))
408 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
409
Raymond Hettinger80d49b32009-11-01 08:53:21 +0000410 def test_bug_7244(self):
411
412 class Repeater(object):
413 # this class is similar to itertools.repeat
414 def __init__(self, o, t, e):
415 self.o = o
416 self.t = int(t)
417 self.e = e
418 def __iter__(self): # its iterator is itself
419 return self
420 def next(self):
421 if self.t > 0:
422 self.t -= 1
423 return self.o
424 else:
425 raise self.e
426
427 # Formerly this code in would fail in debug mode
428 # with Undetected Error and Stop Iteration
429 r1 = Repeater(1, 3, StopIteration)
430 r2 = Repeater(2, 4, StopIteration)
431 def run(r1, r2):
432 result = []
433 for i, j in izip_longest(r1, r2, fillvalue=0):
Raymond Hettinger4da5faa2009-11-01 18:43:31 +0000434 with test_support.captured_output('stdout'):
435 print (i, j)
Raymond Hettinger80d49b32009-11-01 08:53:21 +0000436 result.append((i, j))
437 return result
438 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
439
440 # Formerly, the RuntimeError would be lost
441 # and StopIteration would stop as expected
442 r1 = Repeater(1, 3, RuntimeError)
443 r2 = Repeater(2, 4, StopIteration)
Raymond Hettinger4da5faa2009-11-01 18:43:31 +0000444 it = izip_longest(r1, r2, fillvalue=0)
445 self.assertEqual(next(it), (1, 2))
446 self.assertEqual(next(it), (1, 2))
447 self.assertEqual(next(it), (1, 2))
448 self.assertRaises(RuntimeError, next, it)
Raymond Hettinger80d49b32009-11-01 08:53:21 +0000449
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000450 def test_product(self):
451 for args, result in [
Raymond Hettingerd553d852008-03-04 04:17:08 +0000452 ([], [()]), # zero iterables
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000453 (['ab'], [('a',), ('b',)]), # one iterable
454 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
455 ([range(0), range(2), range(3)], []), # first iterable with zero length
456 ([range(2), range(0), range(3)], []), # middle iterable with zero length
457 ([range(2), range(3), range(0)], []), # last iterable with zero length
458 ]:
459 self.assertEqual(list(product(*args)), result)
Raymond Hettinger08ff6822008-02-29 02:21:48 +0000460 for r in range(4):
461 self.assertEqual(list(product(*(args*r))),
462 list(product(*args, **dict(repeat=r))))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000463 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
464 self.assertRaises(TypeError, product, range(6), None)
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000465
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000466 def product1(*args, **kwds):
467 pools = map(tuple, args) * kwds.get('repeat', 1)
468 n = len(pools)
469 if n == 0:
470 yield ()
471 return
472 if any(len(pool) == 0 for pool in pools):
473 return
474 indices = [0] * n
475 yield tuple(pool[i] for pool, i in zip(pools, indices))
476 while 1:
477 for i in reversed(range(n)): # right to left
478 if indices[i] == len(pools[i]) - 1:
479 continue
480 indices[i] += 1
481 for j in range(i+1, n):
482 indices[j] = 0
483 yield tuple(pool[i] for pool, i in zip(pools, indices))
484 break
485 else:
486 return
487
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000488 def product2(*args, **kwds):
489 'Pure python version used in docs'
490 pools = map(tuple, args) * kwds.get('repeat', 1)
491 result = [[]]
492 for pool in pools:
493 result = [x+[y] for x in result for y in pool]
494 for prod in result:
495 yield tuple(prod)
496
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000497 argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
498 set('abcdefg'), range(11), tuple(range(13))]
499 for i in range(100):
500 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Raymond Hettingerd553d852008-03-04 04:17:08 +0000501 expected_len = prod(map(len, args))
502 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000503 self.assertEqual(list(product(*args)), list(product1(*args)))
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000504 self.assertEqual(list(product(*args)), list(product2(*args)))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000505 args = map(iter, args)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000506 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000507
Raymond Hettinger73d79632008-02-23 02:20:41 +0000508 # Test implementation detail: tuple re-use
509 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
510 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000511
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000512 def test_repeat(self):
513 self.assertEqual(zip(xrange(3),repeat('a')),
514 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000515 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000516 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000517 self.assertEqual(list(repeat('a', 0)), [])
518 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000519 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000520 self.assertRaises(TypeError, repeat, None, 3, 4)
521 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000522 r = repeat(1+0j)
523 self.assertEqual(repr(r), 'repeat((1+0j))')
524 r = repeat(1+0j, 5)
525 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
526 list(r)
527 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000528
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000529 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000530 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
531 [0**1, 1**2, 2**3])
532 self.assertEqual(list(imap(None, 'abc', range(5))),
533 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000534 self.assertEqual(list(imap(None, 'abc', count())),
535 [('a',0),('b',1),('c',2)])
536 self.assertEqual(take(2,imap(None, 'abc', count())),
537 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000538 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000539 self.assertRaises(TypeError, imap)
540 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000541 self.assertRaises(TypeError, imap(10, range(5)).next)
542 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
543 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000544
545 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000546 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
547 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000548 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
549 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000550 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger47317092008-01-17 03:02:14 +0000551 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
552 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000553 self.assertRaises(TypeError, starmap)
554 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
555 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
556 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
557 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000558
559 def test_islice(self):
560 for args in [ # islice(args) should agree with range(args)
561 (10, 20, 3),
562 (10, 3, 20),
563 (10, 20),
564 (10, 3),
565 (20,)
566 ]:
567 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
568
569 for args, tgtargs in [ # Stop when seqn is exhausted
570 ((10, 110, 3), ((10, 100, 3))),
571 ((10, 110), ((10, 100))),
572 ((110,), (100,))
573 ]:
574 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
575
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000576 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000577 self.assertEqual(list(islice(xrange(10), None)), range(10))
Raymond Hettingerb2594052004-12-05 09:25:51 +0000578 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
579 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000580 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
581 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
582
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000583 # Test number of items consumed SF #1171417
584 it = iter(range(10))
585 self.assertEqual(list(islice(it, 3)), range(3))
586 self.assertEqual(list(it), range(3, 10))
587
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000588 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000589 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000590 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
591 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
592 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
593 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
594 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000595 self.assertRaises(ValueError, islice, xrange(10), 'a')
596 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
597 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
598 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
599 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +0000600 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000601
602 def test_takewhile(self):
603 data = [1, 3, 5, 20, 2, 4, 6, 8]
604 underten = lambda x: x<10
605 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000606 self.assertEqual(list(takewhile(underten, [])), [])
607 self.assertRaises(TypeError, takewhile)
608 self.assertRaises(TypeError, takewhile, operator.pow)
609 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
610 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
611 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000612 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
613 self.assertEqual(list(t), [1, 1, 1])
614 self.assertRaises(StopIteration, t.next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000615
616 def test_dropwhile(self):
617 data = [1, 3, 5, 20, 2, 4, 6, 8]
618 underten = lambda x: x<10
619 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000620 self.assertEqual(list(dropwhile(underten, [])), [])
621 self.assertRaises(TypeError, dropwhile)
622 self.assertRaises(TypeError, dropwhile, operator.pow)
623 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
624 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
625 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000626
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000627 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000628 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000629 def irange(n):
630 for i in xrange(n):
631 yield i
632
633 a, b = tee([]) # test empty iterator
634 self.assertEqual(list(a), [])
635 self.assertEqual(list(b), [])
636
637 a, b = tee(irange(n)) # test 100% interleaved
638 self.assertEqual(zip(a,b), zip(range(n),range(n)))
639
640 a, b = tee(irange(n)) # test 0% interleaved
641 self.assertEqual(list(a), range(n))
642 self.assertEqual(list(b), range(n))
643
644 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000645 for i in xrange(100):
646 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000647 del a
648 self.assertEqual(list(b), range(n))
649
650 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000651 for i in xrange(100):
652 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000653 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000654 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000655
656 for j in xrange(5): # test randomly interleaved
657 order = [0]*n + [1]*n
658 random.shuffle(order)
659 lists = ([], [])
660 its = tee(irange(n))
661 for i in order:
662 value = its[i].next()
663 lists[i].append(value)
664 self.assertEqual(lists[0], range(n))
665 self.assertEqual(lists[1], range(n))
666
Raymond Hettingerad983e72003-11-12 14:32:26 +0000667 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000668 self.assertRaises(TypeError, tee)
669 self.assertRaises(TypeError, tee, 3)
670 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000671 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000672
Raymond Hettingerad983e72003-11-12 14:32:26 +0000673 # tee object should be instantiable
674 a, b = tee('abc')
675 c = type(a)('def')
676 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000677
Raymond Hettingerad983e72003-11-12 14:32:26 +0000678 # test long-lagged and multi-way split
679 a, b, c = tee(xrange(2000), 3)
680 for i in xrange(100):
681 self.assertEqual(a.next(), i)
682 self.assertEqual(list(b), range(2000))
683 self.assertEqual([c.next(), c.next()], range(2))
684 self.assertEqual(list(a), range(100,2000))
685 self.assertEqual(list(c), range(2,2000))
686
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000687 # test values of n
688 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Neal Norwitz69e88972006-09-02 02:58:13 +0000689 self.assertRaises(ValueError, tee, [], -1)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000690 for n in xrange(5):
691 result = tee('abc', n)
692 self.assertEqual(type(result), tuple)
693 self.assertEqual(len(result), n)
694 self.assertEqual(map(list, result), [list('abc')]*n)
695
Raymond Hettingerad983e72003-11-12 14:32:26 +0000696 # tee pass-through to copyable iterator
697 a, b = tee('abc')
698 c, d = tee(a)
699 self.assert_(a is c)
700
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000701 # test tee_new
702 t1, t2 = tee('abc')
703 tnew = type(t1)
704 self.assertRaises(TypeError, tnew)
705 self.assertRaises(TypeError, tnew, 10)
706 t3 = tnew(t1)
707 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000708
Raymond Hettingera9f60922004-10-17 16:40:14 +0000709 # test that tee objects are weak referencable
710 a, b = tee(xrange(10))
711 p = proxy(a)
712 self.assertEqual(getattr(p, '__class__'), type(b))
713 del a
714 self.assertRaises(ReferenceError, getattr, p, '__class__')
715
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000716 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000717 self.assertRaises(StopIteration, izip().next)
718
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000719 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000720 self.assertRaises(StopIteration, f([]).next)
721 self.assertRaises(StopIteration, f(StopNow()).next)
722
723 self.assertRaises(StopIteration, islice([], None).next)
724 self.assertRaises(StopIteration, islice(StopNow(), None).next)
725
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000726 p, q = tee([])
727 self.assertRaises(StopIteration, p.next)
728 self.assertRaises(StopIteration, q.next)
729 p, q = tee(StopNow())
730 self.assertRaises(StopIteration, p.next)
731 self.assertRaises(StopIteration, q.next)
732
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000733 self.assertRaises(StopIteration, repeat(None, 0).next)
734
735 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
736 self.assertRaises(StopIteration, f(lambda x:x, []).next)
737 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
738
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000739class TestExamples(unittest.TestCase):
740
741 def test_chain(self):
742 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
743
744 def test_chain_from_iterable(self):
745 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
746
747 def test_combinations(self):
748 self.assertEqual(list(combinations('ABCD', 2)),
749 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
750 self.assertEqual(list(combinations(range(4), 3)),
751 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
752
753 def test_count(self):
754 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
755
756 def test_cycle(self):
757 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
758
759 def test_dropwhile(self):
760 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
761
762 def test_groupby(self):
763 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
764 list('ABCDAB'))
765 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
766 [list('AAAA'), list('BBB'), list('CC'), list('D')])
767
768 def test_ifilter(self):
769 self.assertEqual(list(ifilter(lambda x: x%2, range(10))), [1,3,5,7,9])
770
771 def test_ifilterfalse(self):
772 self.assertEqual(list(ifilterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
773
774 def test_imap(self):
775 self.assertEqual(list(imap(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
776
777 def test_islice(self):
778 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
779 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
780 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
781 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
782
783 def test_izip(self):
784 self.assertEqual(list(izip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
785
786 def test_izip_longest(self):
787 self.assertEqual(list(izip_longest('ABCD', 'xy', fillvalue='-')),
788 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
789
790 def test_permutations(self):
791 self.assertEqual(list(permutations('ABCD', 2)),
792 map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split()))
793 self.assertEqual(list(permutations(range(3))),
794 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
795
796 def test_product(self):
797 self.assertEqual(list(product('ABCD', 'xy')),
798 map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split()))
799 self.assertEqual(list(product(range(2), repeat=3)),
800 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
801 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
802
803 def test_repeat(self):
804 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
805
806 def test_stapmap(self):
807 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
808 [32, 9, 1000])
809
810 def test_takewhile(self):
811 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
812
813
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000814class TestGC(unittest.TestCase):
815
816 def makecycle(self, iterator, container):
817 container.append(iterator)
818 iterator.next()
819 del container, iterator
820
821 def test_chain(self):
822 a = []
823 self.makecycle(chain(a), a)
824
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000825 def test_chain_from_iterable(self):
826 a = []
827 self.makecycle(chain.from_iterable([a]), a)
828
829 def test_combinations(self):
830 a = []
831 self.makecycle(combinations([1,2,a,3], 3), a)
832
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000833 def test_cycle(self):
834 a = []
835 self.makecycle(cycle([a]*2), a)
836
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000837 def test_dropwhile(self):
838 a = []
839 self.makecycle(dropwhile(bool, [0, a, a]), a)
840
841 def test_groupby(self):
842 a = []
843 self.makecycle(groupby([a]*2, lambda x:x), a)
844
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000845 def test_issue2246(self):
846 # Issue 2246 -- the _grouper iterator was not included in GC
847 n = 10
848 keyfunc = lambda x: x
849 for i, j in groupby(xrange(n), key=keyfunc):
850 keyfunc.__dict__.setdefault('x',[]).append(j)
851
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000852 def test_ifilter(self):
853 a = []
854 self.makecycle(ifilter(lambda x:True, [a]*2), a)
855
856 def test_ifilterfalse(self):
857 a = []
858 self.makecycle(ifilterfalse(lambda x:False, a), a)
859
860 def test_izip(self):
861 a = []
862 self.makecycle(izip([a]*2, [a]*3), a)
863
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000864 def test_izip_longest(self):
865 a = []
866 self.makecycle(izip_longest([a]*2, [a]*3), a)
867 b = [a, None]
868 self.makecycle(izip_longest([a]*2, [a]*3, fillvalue=b), a)
869
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000870 def test_imap(self):
871 a = []
872 self.makecycle(imap(lambda x:x, [a]*2), a)
873
874 def test_islice(self):
875 a = []
876 self.makecycle(islice([a]*2, None), a)
877
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000878 def test_permutations(self):
879 a = []
880 self.makecycle(permutations([1,2,a,3], 3), a)
881
882 def test_product(self):
883 a = []
884 self.makecycle(product([1,2,a,3], repeat=3), a)
885
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000886 def test_repeat(self):
887 a = []
888 self.makecycle(repeat(a), a)
889
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000890 def test_starmap(self):
891 a = []
892 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
893
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000894 def test_takewhile(self):
895 a = []
896 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
897
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000898def R(seqn):
899 'Regular generator'
900 for i in seqn:
901 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000902
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000903class G:
904 'Sequence using __getitem__'
905 def __init__(self, seqn):
906 self.seqn = seqn
907 def __getitem__(self, i):
908 return self.seqn[i]
909
910class I:
911 'Sequence using iterator protocol'
912 def __init__(self, seqn):
913 self.seqn = seqn
914 self.i = 0
915 def __iter__(self):
916 return self
917 def next(self):
918 if self.i >= len(self.seqn): raise StopIteration
919 v = self.seqn[self.i]
920 self.i += 1
921 return v
922
923class Ig:
924 'Sequence using iterator protocol defined with a generator'
925 def __init__(self, seqn):
926 self.seqn = seqn
927 self.i = 0
928 def __iter__(self):
929 for val in self.seqn:
930 yield val
931
932class X:
933 'Missing __getitem__ and __iter__'
934 def __init__(self, seqn):
935 self.seqn = seqn
936 self.i = 0
937 def next(self):
938 if self.i >= len(self.seqn): raise StopIteration
939 v = self.seqn[self.i]
940 self.i += 1
941 return v
942
943class N:
944 'Iterator missing next()'
945 def __init__(self, seqn):
946 self.seqn = seqn
947 self.i = 0
948 def __iter__(self):
949 return self
950
951class E:
952 'Test propagation of exceptions'
953 def __init__(self, seqn):
954 self.seqn = seqn
955 self.i = 0
956 def __iter__(self):
957 return self
958 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +0000959 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000960
961class S:
962 'Test immediate stop'
963 def __init__(self, seqn):
964 pass
965 def __iter__(self):
966 return self
967 def next(self):
968 raise StopIteration
969
970def L(seqn):
971 'Test multiple tiers of iterators'
972 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
973
974
975class TestVariousIteratorArgs(unittest.TestCase):
976
977 def test_chain(self):
978 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
979 for g in (G, I, Ig, S, L, R):
980 self.assertEqual(list(chain(g(s))), list(g(s)))
981 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Raymond Hettinger05bf6332008-02-28 22:30:42 +0000982 self.assertRaises(TypeError, list, chain(X(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000983 self.assertRaises(TypeError, list, chain(N(s)))
984 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
985
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000986 def test_product(self):
987 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
988 self.assertRaises(TypeError, product, X(s))
989 self.assertRaises(TypeError, product, N(s))
990 self.assertRaises(ZeroDivisionError, product, E(s))
991
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000992 def test_cycle(self):
993 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
994 for g in (G, I, Ig, S, L, R):
995 tgtlen = len(s) * 3
996 expected = list(g(s))*3
997 actual = list(islice(cycle(g(s)), tgtlen))
998 self.assertEqual(actual, expected)
999 self.assertRaises(TypeError, cycle, X(s))
1000 self.assertRaises(TypeError, list, cycle(N(s)))
1001 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1002
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001003 def test_groupby(self):
1004 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1005 for g in (G, I, Ig, S, L, R):
1006 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1007 self.assertRaises(TypeError, groupby, X(s))
1008 self.assertRaises(TypeError, list, groupby(N(s)))
1009 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1010
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001011 def test_ifilter(self):
1012 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1013 for g in (G, I, Ig, S, L, R):
1014 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
1015 self.assertRaises(TypeError, ifilter, isEven, X(s))
1016 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
1017 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
1018
1019 def test_ifilterfalse(self):
1020 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1021 for g in (G, I, Ig, S, L, R):
1022 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
1023 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
1024 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
1025 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
1026
1027 def test_izip(self):
1028 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1029 for g in (G, I, Ig, S, L, R):
1030 self.assertEqual(list(izip(g(s))), zip(g(s)))
1031 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
1032 self.assertRaises(TypeError, izip, X(s))
1033 self.assertRaises(TypeError, list, izip(N(s)))
1034 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
1035
Raymond Hettingerd36862c2007-02-21 05:20:38 +00001036 def test_iziplongest(self):
1037 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1038 for g in (G, I, Ig, S, L, R):
1039 self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
1040 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
1041 self.assertRaises(TypeError, izip_longest, X(s))
1042 self.assertRaises(TypeError, list, izip_longest(N(s)))
1043 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
1044
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001045 def test_imap(self):
1046 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1047 for g in (G, I, Ig, S, L, R):
1048 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
1049 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
1050 self.assertRaises(TypeError, imap, onearg, X(s))
1051 self.assertRaises(TypeError, list, imap(onearg, N(s)))
1052 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
1053
1054 def test_islice(self):
1055 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1056 for g in (G, I, Ig, S, L, R):
1057 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1058 self.assertRaises(TypeError, islice, X(s), 10)
1059 self.assertRaises(TypeError, list, islice(N(s), 10))
1060 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1061
1062 def test_starmap(self):
1063 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1064 for g in (G, I, Ig, S, L, R):
1065 ss = zip(s, s)
1066 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
1067 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
1068 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
1069 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1070
1071 def test_takewhile(self):
1072 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1073 for g in (G, I, Ig, S, L, R):
1074 tgt = []
1075 for elem in g(s):
1076 if not isEven(elem): break
1077 tgt.append(elem)
1078 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1079 self.assertRaises(TypeError, takewhile, isEven, X(s))
1080 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
1081 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1082
1083 def test_dropwhile(self):
1084 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1085 for g in (G, I, Ig, S, L, R):
1086 tgt = []
1087 for elem in g(s):
1088 if not tgt and isOdd(elem): continue
1089 tgt.append(elem)
1090 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1091 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
1092 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
1093 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1094
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001095 def test_tee(self):
1096 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1097 for g in (G, I, Ig, S, L, R):
1098 it1, it2 = tee(g(s))
1099 self.assertEqual(list(it1), list(g(s)))
1100 self.assertEqual(list(it2), list(g(s)))
1101 self.assertRaises(TypeError, tee, X(s))
1102 self.assertRaises(TypeError, list, tee(N(s))[0])
1103 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1104
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001105class LengthTransparency(unittest.TestCase):
1106
1107 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001108 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001109 self.assertEqual(len(repeat(None, 50)), 50)
1110 self.assertRaises(TypeError, len, repeat(None))
1111
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001112class RegressionTests(unittest.TestCase):
1113
1114 def test_sf_793826(self):
1115 # Fix Armin Rigo's successful efforts to wreak havoc
1116
1117 def mutatingtuple(tuple1, f, tuple2):
1118 # this builds a tuple t which is a copy of tuple1,
1119 # then calls f(t), then mutates t to be equal to tuple2
1120 # (needs len(tuple1) == len(tuple2)).
1121 def g(value, first=[1]):
1122 if first:
1123 del first[:]
1124 f(z.next())
1125 return value
1126 items = list(tuple2)
1127 items[1:1] = list(tuple1)
1128 gen = imap(g, items)
1129 z = izip(*[gen]*len(tuple1))
1130 z.next()
1131
1132 def f(t):
1133 global T
1134 T = t
1135 first[:] = list(T)
1136
1137 first = []
1138 mutatingtuple((1,2,3), f, (4,5,6))
1139 second = list(T)
1140 self.assertEqual(first, second)
1141
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001142
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001143 def test_sf_950057(self):
1144 # Make sure that chain() and cycle() catch exceptions immediately
1145 # rather than when shifting between input sources
1146
1147 def gen1():
1148 hist.append(0)
1149 yield 1
1150 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001151 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001152 hist.append(2)
1153
1154 def gen2(x):
1155 hist.append(3)
1156 yield 2
1157 hist.append(4)
1158 if x:
1159 raise StopIteration
1160
1161 hist = []
1162 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1163 self.assertEqual(hist, [0,1])
1164
1165 hist = []
1166 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1167 self.assertEqual(hist, [0,1])
1168
1169 hist = []
1170 self.assertRaises(AssertionError, list, cycle(gen1()))
1171 self.assertEqual(hist, [0,1])
1172
Georg Brandlb84c1372007-01-21 10:28:43 +00001173class SubclassWithKwargsTest(unittest.TestCase):
1174 def test_keywords_in_subclass(self):
1175 # count is not subclassable...
1176 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
1177 starmap, islice, takewhile, dropwhile, cycle):
1178 class Subclass(cls):
1179 def __init__(self, newarg=None, *args):
1180 cls.__init__(self, *args)
1181 try:
1182 Subclass(newarg=1)
1183 except TypeError, err:
1184 # we expect type errors because of wrong argument count
1185 self.failIf("does not take keyword arguments" in err.args[0])
1186
1187
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001188libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001189
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001190
1191>>> amounts = [120.15, 764.05, 823.14]
1192>>> for checknum, amount in izip(count(1200), amounts):
1193... print 'Check %d is for $%.2f' % (checknum, amount)
1194...
1195Check 1200 is for $120.15
1196Check 1201 is for $764.05
1197Check 1202 is for $823.14
1198
1199>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001200>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1201... print cube
1202...
12031
12048
120527
1206
1207>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001208>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001209... print name.title()
1210...
1211Alex
1212Laura
1213Martin
1214Walter
1215Samuele
1216
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001217>>> from operator import itemgetter
1218>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Armin Rigoa3f09272006-05-28 19:13:17 +00001219>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001220>>> for k, g in groupby(di, itemgetter(1)):
1221... print k, map(itemgetter(0), g)
1222...
12231 ['a', 'c', 'e']
12242 ['b', 'd', 'f']
12253 ['g']
1226
Raymond Hettinger734fb572004-01-20 20:04:40 +00001227# Find runs of consecutive numbers using groupby. The key to the solution
1228# is differencing with a range so that consecutive numbers all appear in
1229# same group.
1230>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
1231>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
1232... print map(operator.itemgetter(1), g)
1233...
1234[1]
1235[4, 5, 6]
1236[10]
1237[15, 16, 17, 18]
1238[22]
1239[25, 26, 27, 28]
1240
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001241>>> def take(n, iterable):
1242... "Return first n items of the iterable as a list"
1243... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001244
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001245>>> def enumerate(iterable, start=0):
1246... return izip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001247
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001248>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001249... "Return function(0), function(1), ..."
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001250... return imap(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001251
Raymond Hettinger5894c2b2009-02-19 05:38:53 +00001252>>> def nth(iterable, n, default=None):
1253... "Returns the nth item or a default value"
1254... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001255
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001256>>> def quantify(iterable, pred=bool):
1257... "Count how many times the predicate is true"
1258... return sum(imap(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001259
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001260>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001261... "Returns the sequence elements and then returns None indefinitely"
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001262... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001263
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001264>>> def ncycles(iterable, n):
1265... "Returns the seqeuence elements n times"
1266... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001267
1268>>> def dotproduct(vec1, vec2):
1269... return sum(imap(operator.mul, vec1, vec2))
1270
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001271>>> def flatten(listOfLists):
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001272... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001273
1274>>> def repeatfunc(func, times=None, *args):
1275... "Repeat calls to func with specified arguments."
1276... " Example: repeatfunc(random.random)"
1277... if times is None:
1278... return starmap(func, repeat(args))
1279... else:
1280... return starmap(func, repeat(args, times))
1281
Raymond Hettingerd591f662003-10-26 15:34:50 +00001282>>> def pairwise(iterable):
1283... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1284... a, b = tee(iterable)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001285... for elem in b:
1286... break
Raymond Hettingerad983e72003-11-12 14:32:26 +00001287... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001288
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001289>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001290... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001291... args = [iter(iterable)] * n
Raymond Hettingerf080e6d2008-07-31 01:19:50 +00001292... return izip_longest(fillvalue=fillvalue, *args)
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001293
1294>>> def roundrobin(*iterables):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001295... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001296... # Recipe credited to George Sakkis
1297... pending = len(iterables)
1298... nexts = cycle(iter(it).next for it in iterables)
1299... while pending:
1300... try:
1301... for next in nexts:
1302... yield next()
1303... except StopIteration:
1304... pending -= 1
1305... nexts = cycle(islice(nexts, pending))
1306
1307>>> def powerset(iterable):
Raymond Hettinger5cd012b2009-01-26 02:17:52 +00001308... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1309... s = list(iterable)
1310... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001311
Raymond Hettingere8b4b602008-03-11 00:19:07 +00001312>>> def compress(data, selectors):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001313... "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
1314... return (d for d, s in izip(data, selectors) if s)
Raymond Hettingere8b4b602008-03-11 00:19:07 +00001315
Raymond Hettinger33691672008-07-19 00:43:00 +00001316>>> def combinations_with_replacement(iterable, r):
1317... "combinations_with_replacement('ABC', 3) --> AA AB AC BB BC CC"
1318... pool = tuple(iterable)
1319... n = len(pool)
Raymond Hettingercdc9f2c2009-01-27 06:38:00 +00001320... if not n and r:
1321... return
Raymond Hettinger33691672008-07-19 00:43:00 +00001322... indices = [0] * r
1323... yield tuple(pool[i] for i in indices)
1324... while 1:
1325... for i in reversed(range(r)):
1326... if indices[i] != n - 1:
1327... break
1328... else:
1329... return
1330... indices[i:] = [indices[i] + 1] * (r - i)
1331... yield tuple(pool[i] for i in indices)
1332
Raymond Hettingerd8461652009-01-02 21:20:38 +00001333>>> def unique_everseen(iterable, key=None):
1334... "List unique elements, preserving order. Remember all elements ever seen."
1335... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1336... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1337... seen = set()
1338... seen_add = seen.add
1339... if key is None:
1340... for element in iterable:
1341... if element not in seen:
1342... seen_add(element)
1343... yield element
1344... else:
1345... for element in iterable:
1346... k = key(element)
1347... if k not in seen:
1348... seen_add(k)
1349... yield element
1350
1351>>> def unique_justseen(iterable, key=None):
1352... "List unique elements, preserving order. Remember only the element just seen."
1353... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1354... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1355... return imap(next, imap(itemgetter(1), groupby(iterable, key)))
1356
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001357This is not part of the examples but it tests to make sure the definitions
1358perform as purported.
1359
Raymond Hettingera098b332003-09-08 23:58:40 +00001360>>> take(10, count())
1361[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1362
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001363>>> list(enumerate('abc'))
1364[(0, 'a'), (1, 'b'), (2, 'c')]
1365
1366>>> list(islice(tabulate(lambda x: 2*x), 4))
1367[0, 2, 4, 6]
1368
1369>>> nth('abcde', 3)
Raymond Hettinger11485b42009-02-04 19:34:31 +00001370'd'
1371
1372>>> nth('abcde', 9) is None
1373True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001374
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001375>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000137650
1377
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001378>>> a = [[1, 2, 3], [4, 5, 6]]
1379>>> flatten(a)
1380[1, 2, 3, 4, 5, 6]
1381
1382>>> list(repeatfunc(pow, 5, 2, 3))
1383[8, 8, 8, 8, 8]
1384
1385>>> import random
1386>>> take(5, imap(int, repeatfunc(random.random)))
1387[0, 0, 0, 0, 0]
1388
Raymond Hettingerd591f662003-10-26 15:34:50 +00001389>>> list(pairwise('abcd'))
1390[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001391
Raymond Hettingerd591f662003-10-26 15:34:50 +00001392>>> list(pairwise([]))
1393[]
1394
1395>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001396[]
1397
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001398>>> list(islice(padnone('abc'), 0, 6))
1399['a', 'b', 'c', None, None, None]
1400
1401>>> list(ncycles('abc', 3))
1402['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1403
1404>>> dotproduct([1,2,3], [4,5,6])
140532
1406
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001407>>> list(grouper(3, 'abcdefg', 'x'))
1408[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1409
1410>>> list(roundrobin('abc', 'd', 'ef'))
1411['a', 'd', 'e', 'b', 'f', 'c']
1412
Raymond Hettinger5cd012b2009-01-26 02:17:52 +00001413>>> list(powerset([1,2,3]))
1414[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001415
Raymond Hettingere8b4b602008-03-11 00:19:07 +00001416>>> list(compress('abcdef', [1,0,1,0,1,1]))
1417['a', 'c', 'e', 'f']
1418
Raymond Hettinger33691672008-07-19 00:43:00 +00001419>>> list(combinations_with_replacement('abc', 2))
1420[('a', 'a'), ('a', 'b'), ('a', 'c'), ('b', 'b'), ('b', 'c'), ('c', 'c')]
1421
Raymond Hettinger825758c2009-01-08 05:20:19 +00001422>>> list(combinations_with_replacement('01', 3))
1423[('0', '0', '0'), ('0', '0', '1'), ('0', '1', '1'), ('1', '1', '1')]
1424
1425>>> def combinations_with_replacement2(iterable, r):
1426... 'Alternate version that filters from product()'
1427... pool = tuple(iterable)
1428... n = len(pool)
1429... for indices in product(range(n), repeat=r):
1430... if sorted(indices) == list(indices):
1431... yield tuple(pool[i] for i in indices)
1432
1433>>> list(combinations_with_replacement('abc', 2)) == list(combinations_with_replacement2('abc', 2))
1434True
1435
1436>>> list(combinations_with_replacement('01', 3)) == list(combinations_with_replacement2('01', 3))
1437True
1438
1439>>> list(combinations_with_replacement('2310', 6)) == list(combinations_with_replacement2('2310', 6))
1440True
1441
Raymond Hettingerd8461652009-01-02 21:20:38 +00001442>>> list(unique_everseen('AAAABBBCCDAABBB'))
1443['A', 'B', 'C', 'D']
1444
1445>>> list(unique_everseen('ABBCcAD', str.lower))
1446['A', 'B', 'C', 'D']
1447
1448>>> list(unique_justseen('AAAABBBCCDAABBB'))
1449['A', 'B', 'C', 'D', 'A', 'B']
1450
1451>>> list(unique_justseen('ABBCcAD', str.lower))
1452['A', 'B', 'C', 'A', 'D']
1453
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001454"""
1455
1456__test__ = {'libreftest' : libreftest}
1457
1458def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001459 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Georg Brandlb84c1372007-01-21 10:28:43 +00001460 RegressionTests, LengthTransparency,
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001461 SubclassWithKwargsTest, TestExamples)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001462 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001463
1464 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001465 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001466 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001467 counts = [None] * 5
1468 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001469 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001470 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001471 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001472 print counts
1473
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001474 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001475 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001476
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001477if __name__ == "__main__":
1478 test_main(verbose=True)