blob: ee5af4b601d797f5fc47a1e738f35d99285070c1 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001import unittest
2from test import test_support
3from itertools import *
Raymond Hettingera9f60922004-10-17 16:40:14 +00004from weakref import proxy
Raymond Hettinger2012f172003-02-07 05:32:58 +00005import sys
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00006import operator
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00007import random
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +00008maxsize = test_support.MAX_Py_ssize_t
9minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000010
11def onearg(x):
12 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000013 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000014
15def errfunc(*args):
16 'Test function that raises an error'
17 raise ValueError
18
19def gen3():
20 'Non-restartable source sequence'
21 for i in (0, 1, 2):
22 yield i
23
24def isEven(x):
25 'Test predicate'
26 return x%2==0
27
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000028def isOdd(x):
29 'Test predicate'
30 return x%2==1
31
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000032class StopNow:
33 'Class emulating an empty iterable.'
34 def __iter__(self):
35 return self
36 def next(self):
37 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000038
Raymond Hettinger02420702003-06-29 20:36:23 +000039def take(n, seq):
40 'Convenience function for partially consuming a long of infinite iterable'
41 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000042
Raymond Hettingerd553d852008-03-04 04:17:08 +000043def prod(iterable):
44 return reduce(operator.mul, iterable, 1)
45
Raymond Hettinger93e804d2008-02-26 23:40:50 +000046def fact(n):
47 'Factorial'
Raymond Hettingerd553d852008-03-04 04:17:08 +000048 return prod(range(1, n+1))
49
Raymond Hettinger96ef8112003-02-01 00:10:11 +000050class TestBasicOps(unittest.TestCase):
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000051 def test_chain(self):
Raymond Hettingerad47fa12008-03-06 20:52:01 +000052
53 def chain2(*iterables):
54 'Pure python version in the docs'
55 for it in iterables:
56 for element in it:
57 yield element
58
59 for c in (chain, chain2):
60 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
61 self.assertEqual(list(c('abc')), list('abc'))
62 self.assertEqual(list(c('')), [])
63 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
64 self.assertRaises(TypeError, list,c(2, 3))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000065
Raymond Hettingerb4cbc982008-02-28 22:46:41 +000066 def test_chain_from_iterable(self):
67 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
68 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
69 self.assertEqual(list(chain.from_iterable([''])), [])
70 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
71 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
72
Raymond Hettinger93e804d2008-02-26 23:40:50 +000073 def test_combinations(self):
Raymond Hettinger825758c2009-01-08 05:20:19 +000074 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Raymond Hettinger93e804d2008-02-26 23:40:50 +000075 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Raymond Hettingerd553d852008-03-04 04:17:08 +000076 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Raymond Hettinger93e804d2008-02-26 23:40:50 +000077 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Raymond Hettinger825758c2009-01-08 05:20:19 +000078 self.assertEqual(list(combinations('abc', 32)), []) # r > n
Raymond Hettinger93e804d2008-02-26 23:40:50 +000079 self.assertEqual(list(combinations(range(4), 3)),
80 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
Raymond Hettingerd553d852008-03-04 04:17:08 +000081
82 def combinations1(iterable, r):
83 'Pure python version shown in the docs'
84 pool = tuple(iterable)
85 n = len(pool)
Raymond Hettinger825758c2009-01-08 05:20:19 +000086 if r > n:
87 return
Raymond Hettingerd553d852008-03-04 04:17:08 +000088 indices = range(r)
89 yield tuple(pool[i] for i in indices)
90 while 1:
91 for i in reversed(range(r)):
92 if indices[i] != i + n - r:
93 break
94 else:
95 return
96 indices[i] += 1
97 for j in range(i+1, r):
98 indices[j] = indices[j-1] + 1
99 yield tuple(pool[i] for i in indices)
100
101 def combinations2(iterable, r):
102 'Pure python version shown in the docs'
103 pool = tuple(iterable)
104 n = len(pool)
105 for indices in permutations(range(n), r):
106 if sorted(indices) == list(indices):
107 yield tuple(pool[i] for i in indices)
108
109 for n in range(7):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000110 values = [5*x-12 for x in range(n)]
Raymond Hettinger825758c2009-01-08 05:20:19 +0000111 for r in range(n+2):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000112 result = list(combinations(values, r))
Raymond Hettinger825758c2009-01-08 05:20:19 +0000113 self.assertEqual(len(result), 0 if r>n else fact(n) / fact(r) / fact(n-r)) # right number of combs
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000114 self.assertEqual(len(result), len(set(result))) # no repeats
115 self.assertEqual(result, sorted(result)) # lexicographic order
116 for c in result:
117 self.assertEqual(len(c), r) # r-length combinations
118 self.assertEqual(len(set(c)), r) # no duplicate elements
119 self.assertEqual(list(c), sorted(c)) # keep original ordering
120 self.assert_(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000121 self.assertEqual(list(c),
122 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Raymond Hettingerd553d852008-03-04 04:17:08 +0000123 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger825758c2009-01-08 05:20:19 +0000124 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000125
126 # Test implementation detail: tuple re-use
127 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
128 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
129
130 def test_permutations(self):
131 self.assertRaises(TypeError, permutations) # too few arguments
132 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000133 self.assertRaises(TypeError, permutations, None) # pool is not iterable
134 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger825758c2009-01-08 05:20:19 +0000135 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000136 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Raymond Hettingerd553d852008-03-04 04:17:08 +0000137 self.assertEqual(list(permutations(range(3), 2)),
138 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
139
140 def permutations1(iterable, r=None):
141 'Pure python version shown in the docs'
142 pool = tuple(iterable)
143 n = len(pool)
144 r = n if r is None else r
Raymond Hettinger825758c2009-01-08 05:20:19 +0000145 if r > n:
146 return
Raymond Hettingerd553d852008-03-04 04:17:08 +0000147 indices = range(n)
Raymond Hettingere70bb8d2008-03-23 00:55:46 +0000148 cycles = range(n, n-r, -1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000149 yield tuple(pool[i] for i in indices[:r])
150 while n:
151 for i in reversed(range(r)):
152 cycles[i] -= 1
153 if cycles[i] == 0:
154 indices[i:] = indices[i+1:] + indices[i:i+1]
155 cycles[i] = n - i
156 else:
157 j = cycles[i]
158 indices[i], indices[-j] = indices[-j], indices[i]
159 yield tuple(pool[i] for i in indices[:r])
160 break
161 else:
162 return
163
164 def permutations2(iterable, r=None):
165 'Pure python version shown in the docs'
166 pool = tuple(iterable)
167 n = len(pool)
168 r = n if r is None else r
169 for indices in product(range(n), repeat=r):
170 if len(set(indices)) == r:
171 yield tuple(pool[i] for i in indices)
172
173 for n in range(7):
174 values = [5*x-12 for x in range(n)]
Raymond Hettinger825758c2009-01-08 05:20:19 +0000175 for r in range(n+2):
Raymond Hettingerd553d852008-03-04 04:17:08 +0000176 result = list(permutations(values, r))
Raymond Hettinger825758c2009-01-08 05:20:19 +0000177 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 +0000178 self.assertEqual(len(result), len(set(result))) # no repeats
179 self.assertEqual(result, sorted(result)) # lexicographic order
180 for p in result:
181 self.assertEqual(len(p), r) # r-length permutations
182 self.assertEqual(len(set(p)), r) # no duplicate elements
183 self.assert_(all(e in values for e in p)) # elements taken from input iterable
184 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger825758c2009-01-08 05:20:19 +0000185 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000186 if r == n:
187 self.assertEqual(result, list(permutations(values, None))) # test r as None
188 self.assertEqual(result, list(permutations(values))) # test default r
189
190 # Test implementation detail: tuple re-use
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000191 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000192 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000193
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000194 def test_count(self):
195 self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
196 self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000197 self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000198 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
199 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000200 self.assertRaises(TypeError, count, 2, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000201 self.assertRaises(TypeError, count, 'a')
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000202 self.assertEqual(list(islice(count(maxsize-5), 10)), range(maxsize-5, maxsize+5))
203 self.assertEqual(list(islice(count(-maxsize-5), 10)), range(-maxsize-5, -maxsize+5))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000204 c = count(3)
205 self.assertEqual(repr(c), 'count(3)')
206 c.next()
207 self.assertEqual(repr(c), 'count(4)')
Jack Diederich36234e82006-09-21 17:50:26 +0000208 c = count(-9)
209 self.assertEqual(repr(c), 'count(-9)')
210 c.next()
211 self.assertEqual(c.next(), -8)
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000212 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 +0000213 # Test repr (ignoring the L in longs)
214 r1 = repr(count(i)).replace('L', '')
215 r2 = 'count(%r)'.__mod__(i).replace('L', '')
216 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000217
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000218 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000219 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000220 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000221 self.assertRaises(TypeError, cycle)
222 self.assertRaises(TypeError, cycle, 5)
223 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000224
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000225 def test_groupby(self):
226 # Check whether it accepts arguments correctly
227 self.assertEqual([], list(groupby([])))
228 self.assertEqual([], list(groupby([], key=id)))
229 self.assertRaises(TypeError, list, groupby('abc', []))
230 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000231 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000232
233 # Check normal input
234 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
235 (2,15,22), (3,16,23), (3,17,23)]
236 dup = []
237 for k, g in groupby(s, lambda r:r[0]):
238 for elem in g:
239 self.assertEqual(k, elem[0])
240 dup.append(elem)
241 self.assertEqual(s, dup)
242
243 # Check nested case
244 dup = []
245 for k, g in groupby(s, lambda r:r[0]):
246 for ik, ig in groupby(g, lambda r:r[2]):
247 for elem in ig:
248 self.assertEqual(k, elem[0])
249 self.assertEqual(ik, elem[2])
250 dup.append(elem)
251 self.assertEqual(s, dup)
252
253 # Check case where inner iterator is not used
254 keys = [k for k, g in groupby(s, lambda r:r[0])]
255 expectedkeys = set([r[0] for r in s])
256 self.assertEqual(set(keys), expectedkeys)
257 self.assertEqual(len(keys), len(expectedkeys))
258
259 # Exercise pipes and filters style
260 s = 'abracadabra'
261 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000262 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000263 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
264 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000265 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000266 self.assertEqual(r, ['a', 'b', 'r'])
267 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000268 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000269 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
270 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000271 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000272 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
273
274 # iter.next failure
275 class ExpectedError(Exception):
276 pass
277 def delayed_raise(n=0):
278 for i in range(n):
279 yield 'yo'
280 raise ExpectedError
281 def gulp(iterable, keyp=None, func=list):
282 return [func(g) for k, g in groupby(iterable, keyp)]
283
284 # iter.next failure on outer object
285 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
286 # iter.next failure on inner object
287 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
288
289 # __cmp__ failure
290 class DummyCmp:
291 def __cmp__(self, dst):
292 raise ExpectedError
293 s = [DummyCmp(), DummyCmp(), None]
294
295 # __cmp__ failure on outer object
296 self.assertRaises(ExpectedError, gulp, s, func=id)
297 # __cmp__ failure on inner object
298 self.assertRaises(ExpectedError, gulp, s)
299
300 # keyfunc failure
301 def keyfunc(obj):
302 if keyfunc.skip > 0:
303 keyfunc.skip -= 1
304 return obj
305 else:
306 raise ExpectedError
307
308 # keyfunc failure on outer object
309 keyfunc.skip = 0
310 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
311 keyfunc.skip = 1
312 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
313
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000314 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000315 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000316 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000317 self.assertEqual(list(ifilter(bool, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000318 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000319 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000320 self.assertRaises(TypeError, ifilter, lambda x:x)
321 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000322 self.assertRaises(TypeError, ifilter, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000323 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000324
325 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000326 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
327 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000328 self.assertEqual(list(ifilterfalse(bool, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000329 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000330 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000331 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
332 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000333 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000334 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000335
336 def test_izip(self):
337 ans = [(x,y) for x, y in izip('abc',count())]
338 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000339 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
340 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
Raymond Hettinger02420702003-06-29 20:36:23 +0000341 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000342 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000343 self.assertEqual(list(izip()), zip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000344 self.assertRaises(TypeError, izip, 3)
345 self.assertRaises(TypeError, izip, range(3), 3)
346 # Check tuple re-use (implementation detail)
347 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
348 zip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000349 self.assertEqual([pair for pair in izip('abc', 'def')],
350 zip('abc', 'def'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000351 ids = map(id, izip('abc', 'def'))
352 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000353 ids = map(id, list(izip('abc', 'def')))
354 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000355
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000356 def test_iziplongest(self):
357 for args in [
358 ['abc', range(6)],
359 [range(6), 'abc'],
360 [range(1000), range(2000,2100), range(3000,3050)],
361 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
362 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
363 ]:
364 target = map(None, *args)
365 self.assertEqual(list(izip_longest(*args)), target)
366 self.assertEqual(list(izip_longest(*args, **{})), target)
367 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
368 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Tim Petersea5962f2007-03-12 18:07:52 +0000369
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000370 self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
371
372 self.assertEqual(list(izip_longest()), zip())
373 self.assertEqual(list(izip_longest([])), zip([]))
374 self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
Tim Petersea5962f2007-03-12 18:07:52 +0000375
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000376 self.assertEqual(list(izip_longest('abc', 'defg', **{})), map(None, 'abc', 'defg')) # empty keyword dict
377 self.assertRaises(TypeError, izip_longest, 3)
378 self.assertRaises(TypeError, izip_longest, range(3), 3)
379
380 for stmt in [
381 "izip_longest('abc', fv=1)",
Tim Petersea5962f2007-03-12 18:07:52 +0000382 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000383 ]:
384 try:
385 eval(stmt, globals(), locals())
386 except TypeError:
387 pass
388 else:
389 self.fail('Did not raise Type in: ' + stmt)
Tim Petersea5962f2007-03-12 18:07:52 +0000390
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000391 # Check tuple re-use (implementation detail)
392 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
393 zip('abc', 'def'))
394 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
395 zip('abc', 'def'))
396 ids = map(id, izip_longest('abc', 'def'))
397 self.assertEqual(min(ids), max(ids))
398 ids = map(id, list(izip_longest('abc', 'def')))
399 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
400
Raymond Hettinger80d49b32009-11-01 08:53:21 +0000401 def test_bug_7244(self):
402
403 class Repeater(object):
404 # this class is similar to itertools.repeat
405 def __init__(self, o, t, e):
406 self.o = o
407 self.t = int(t)
408 self.e = e
409 def __iter__(self): # its iterator is itself
410 return self
411 def next(self):
412 if self.t > 0:
413 self.t -= 1
414 return self.o
415 else:
416 raise self.e
417
418 # Formerly this code in would fail in debug mode
419 # with Undetected Error and Stop Iteration
420 r1 = Repeater(1, 3, StopIteration)
421 r2 = Repeater(2, 4, StopIteration)
422 def run(r1, r2):
423 result = []
424 for i, j in izip_longest(r1, r2, fillvalue=0):
425 print(i, j)
426 result.append((i, j))
427 return result
428 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
429
430 # Formerly, the RuntimeError would be lost
431 # and StopIteration would stop as expected
432 r1 = Repeater(1, 3, RuntimeError)
433 r2 = Repeater(2, 4, StopIteration)
434 mylist = lambda it: [v for v in it]
435 self.assertRaises(RuntimeError, mylist, izip_longest(r1, r2, fillvalue=0))
436
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000437 def test_product(self):
438 for args, result in [
Raymond Hettingerd553d852008-03-04 04:17:08 +0000439 ([], [()]), # zero iterables
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000440 (['ab'], [('a',), ('b',)]), # one iterable
441 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
442 ([range(0), range(2), range(3)], []), # first iterable with zero length
443 ([range(2), range(0), range(3)], []), # middle iterable with zero length
444 ([range(2), range(3), range(0)], []), # last iterable with zero length
445 ]:
446 self.assertEqual(list(product(*args)), result)
Raymond Hettinger08ff6822008-02-29 02:21:48 +0000447 for r in range(4):
448 self.assertEqual(list(product(*(args*r))),
449 list(product(*args, **dict(repeat=r))))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000450 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
451 self.assertRaises(TypeError, product, range(6), None)
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000452
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000453 def product1(*args, **kwds):
454 pools = map(tuple, args) * kwds.get('repeat', 1)
455 n = len(pools)
456 if n == 0:
457 yield ()
458 return
459 if any(len(pool) == 0 for pool in pools):
460 return
461 indices = [0] * n
462 yield tuple(pool[i] for pool, i in zip(pools, indices))
463 while 1:
464 for i in reversed(range(n)): # right to left
465 if indices[i] == len(pools[i]) - 1:
466 continue
467 indices[i] += 1
468 for j in range(i+1, n):
469 indices[j] = 0
470 yield tuple(pool[i] for pool, i in zip(pools, indices))
471 break
472 else:
473 return
474
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000475 def product2(*args, **kwds):
476 'Pure python version used in docs'
477 pools = map(tuple, args) * kwds.get('repeat', 1)
478 result = [[]]
479 for pool in pools:
480 result = [x+[y] for x in result for y in pool]
481 for prod in result:
482 yield tuple(prod)
483
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000484 argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
485 set('abcdefg'), range(11), tuple(range(13))]
486 for i in range(100):
487 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Raymond Hettingerd553d852008-03-04 04:17:08 +0000488 expected_len = prod(map(len, args))
489 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000490 self.assertEqual(list(product(*args)), list(product1(*args)))
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000491 self.assertEqual(list(product(*args)), list(product2(*args)))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000492 args = map(iter, args)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000493 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000494
Raymond Hettinger73d79632008-02-23 02:20:41 +0000495 # Test implementation detail: tuple re-use
496 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
497 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000498
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000499 def test_repeat(self):
500 self.assertEqual(zip(xrange(3),repeat('a')),
501 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000502 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000503 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000504 self.assertEqual(list(repeat('a', 0)), [])
505 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000506 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000507 self.assertRaises(TypeError, repeat, None, 3, 4)
508 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000509 r = repeat(1+0j)
510 self.assertEqual(repr(r), 'repeat((1+0j))')
511 r = repeat(1+0j, 5)
512 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
513 list(r)
514 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000515
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000516 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000517 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
518 [0**1, 1**2, 2**3])
519 self.assertEqual(list(imap(None, 'abc', range(5))),
520 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000521 self.assertEqual(list(imap(None, 'abc', count())),
522 [('a',0),('b',1),('c',2)])
523 self.assertEqual(take(2,imap(None, 'abc', count())),
524 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000525 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000526 self.assertRaises(TypeError, imap)
527 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000528 self.assertRaises(TypeError, imap(10, range(5)).next)
529 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
530 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000531
532 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000533 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
534 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000535 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
536 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000537 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger47317092008-01-17 03:02:14 +0000538 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
539 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000540 self.assertRaises(TypeError, starmap)
541 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
542 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
543 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
544 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000545
546 def test_islice(self):
547 for args in [ # islice(args) should agree with range(args)
548 (10, 20, 3),
549 (10, 3, 20),
550 (10, 20),
551 (10, 3),
552 (20,)
553 ]:
554 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
555
556 for args, tgtargs in [ # Stop when seqn is exhausted
557 ((10, 110, 3), ((10, 100, 3))),
558 ((10, 110), ((10, 100))),
559 ((110,), (100,))
560 ]:
561 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
562
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000563 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000564 self.assertEqual(list(islice(xrange(10), None)), range(10))
Raymond Hettingerb2594052004-12-05 09:25:51 +0000565 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
566 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000567 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
568 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
569
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000570 # Test number of items consumed SF #1171417
571 it = iter(range(10))
572 self.assertEqual(list(islice(it, 3)), range(3))
573 self.assertEqual(list(it), range(3, 10))
574
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000575 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000576 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000577 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
578 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
579 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
580 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
581 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000582 self.assertRaises(ValueError, islice, xrange(10), 'a')
583 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
584 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
585 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
586 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +0000587 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000588
589 def test_takewhile(self):
590 data = [1, 3, 5, 20, 2, 4, 6, 8]
591 underten = lambda x: x<10
592 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000593 self.assertEqual(list(takewhile(underten, [])), [])
594 self.assertRaises(TypeError, takewhile)
595 self.assertRaises(TypeError, takewhile, operator.pow)
596 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
597 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
598 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000599 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
600 self.assertEqual(list(t), [1, 1, 1])
601 self.assertRaises(StopIteration, t.next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000602
603 def test_dropwhile(self):
604 data = [1, 3, 5, 20, 2, 4, 6, 8]
605 underten = lambda x: x<10
606 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000607 self.assertEqual(list(dropwhile(underten, [])), [])
608 self.assertRaises(TypeError, dropwhile)
609 self.assertRaises(TypeError, dropwhile, operator.pow)
610 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
611 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
612 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000613
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000614 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000615 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000616 def irange(n):
617 for i in xrange(n):
618 yield i
619
620 a, b = tee([]) # test empty iterator
621 self.assertEqual(list(a), [])
622 self.assertEqual(list(b), [])
623
624 a, b = tee(irange(n)) # test 100% interleaved
625 self.assertEqual(zip(a,b), zip(range(n),range(n)))
626
627 a, b = tee(irange(n)) # test 0% interleaved
628 self.assertEqual(list(a), range(n))
629 self.assertEqual(list(b), range(n))
630
631 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000632 for i in xrange(100):
633 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000634 del a
635 self.assertEqual(list(b), range(n))
636
637 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000638 for i in xrange(100):
639 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000640 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000641 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000642
643 for j in xrange(5): # test randomly interleaved
644 order = [0]*n + [1]*n
645 random.shuffle(order)
646 lists = ([], [])
647 its = tee(irange(n))
648 for i in order:
649 value = its[i].next()
650 lists[i].append(value)
651 self.assertEqual(lists[0], range(n))
652 self.assertEqual(lists[1], range(n))
653
Raymond Hettingerad983e72003-11-12 14:32:26 +0000654 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000655 self.assertRaises(TypeError, tee)
656 self.assertRaises(TypeError, tee, 3)
657 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000658 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000659
Raymond Hettingerad983e72003-11-12 14:32:26 +0000660 # tee object should be instantiable
661 a, b = tee('abc')
662 c = type(a)('def')
663 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000664
Raymond Hettingerad983e72003-11-12 14:32:26 +0000665 # test long-lagged and multi-way split
666 a, b, c = tee(xrange(2000), 3)
667 for i in xrange(100):
668 self.assertEqual(a.next(), i)
669 self.assertEqual(list(b), range(2000))
670 self.assertEqual([c.next(), c.next()], range(2))
671 self.assertEqual(list(a), range(100,2000))
672 self.assertEqual(list(c), range(2,2000))
673
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000674 # test values of n
675 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Neal Norwitz69e88972006-09-02 02:58:13 +0000676 self.assertRaises(ValueError, tee, [], -1)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000677 for n in xrange(5):
678 result = tee('abc', n)
679 self.assertEqual(type(result), tuple)
680 self.assertEqual(len(result), n)
681 self.assertEqual(map(list, result), [list('abc')]*n)
682
Raymond Hettingerad983e72003-11-12 14:32:26 +0000683 # tee pass-through to copyable iterator
684 a, b = tee('abc')
685 c, d = tee(a)
686 self.assert_(a is c)
687
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000688 # test tee_new
689 t1, t2 = tee('abc')
690 tnew = type(t1)
691 self.assertRaises(TypeError, tnew)
692 self.assertRaises(TypeError, tnew, 10)
693 t3 = tnew(t1)
694 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000695
Raymond Hettingera9f60922004-10-17 16:40:14 +0000696 # test that tee objects are weak referencable
697 a, b = tee(xrange(10))
698 p = proxy(a)
699 self.assertEqual(getattr(p, '__class__'), type(b))
700 del a
701 self.assertRaises(ReferenceError, getattr, p, '__class__')
702
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000703 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000704 self.assertRaises(StopIteration, izip().next)
705
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000706 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000707 self.assertRaises(StopIteration, f([]).next)
708 self.assertRaises(StopIteration, f(StopNow()).next)
709
710 self.assertRaises(StopIteration, islice([], None).next)
711 self.assertRaises(StopIteration, islice(StopNow(), None).next)
712
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000713 p, q = tee([])
714 self.assertRaises(StopIteration, p.next)
715 self.assertRaises(StopIteration, q.next)
716 p, q = tee(StopNow())
717 self.assertRaises(StopIteration, p.next)
718 self.assertRaises(StopIteration, q.next)
719
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000720 self.assertRaises(StopIteration, repeat(None, 0).next)
721
722 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
723 self.assertRaises(StopIteration, f(lambda x:x, []).next)
724 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
725
Raymond Hettinger80d49b32009-11-01 08:53:21 +0000726
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000727class TestExamples(unittest.TestCase):
728
729 def test_chain(self):
730 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
731
732 def test_chain_from_iterable(self):
733 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
734
735 def test_combinations(self):
736 self.assertEqual(list(combinations('ABCD', 2)),
737 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
738 self.assertEqual(list(combinations(range(4), 3)),
739 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
740
741 def test_count(self):
742 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
743
744 def test_cycle(self):
745 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
746
747 def test_dropwhile(self):
748 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
749
750 def test_groupby(self):
751 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
752 list('ABCDAB'))
753 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
754 [list('AAAA'), list('BBB'), list('CC'), list('D')])
755
756 def test_ifilter(self):
757 self.assertEqual(list(ifilter(lambda x: x%2, range(10))), [1,3,5,7,9])
758
759 def test_ifilterfalse(self):
760 self.assertEqual(list(ifilterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
761
762 def test_imap(self):
763 self.assertEqual(list(imap(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
764
765 def test_islice(self):
766 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
767 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
768 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
769 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
770
771 def test_izip(self):
772 self.assertEqual(list(izip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
773
774 def test_izip_longest(self):
775 self.assertEqual(list(izip_longest('ABCD', 'xy', fillvalue='-')),
776 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
777
778 def test_permutations(self):
779 self.assertEqual(list(permutations('ABCD', 2)),
780 map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split()))
781 self.assertEqual(list(permutations(range(3))),
782 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
783
784 def test_product(self):
785 self.assertEqual(list(product('ABCD', 'xy')),
786 map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split()))
787 self.assertEqual(list(product(range(2), repeat=3)),
788 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
789 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
790
791 def test_repeat(self):
792 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
793
794 def test_stapmap(self):
795 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
796 [32, 9, 1000])
797
798 def test_takewhile(self):
799 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
800
801
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000802class TestGC(unittest.TestCase):
803
804 def makecycle(self, iterator, container):
805 container.append(iterator)
806 iterator.next()
807 del container, iterator
808
809 def test_chain(self):
810 a = []
811 self.makecycle(chain(a), a)
812
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000813 def test_chain_from_iterable(self):
814 a = []
815 self.makecycle(chain.from_iterable([a]), a)
816
817 def test_combinations(self):
818 a = []
819 self.makecycle(combinations([1,2,a,3], 3), a)
820
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000821 def test_cycle(self):
822 a = []
823 self.makecycle(cycle([a]*2), a)
824
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000825 def test_dropwhile(self):
826 a = []
827 self.makecycle(dropwhile(bool, [0, a, a]), a)
828
829 def test_groupby(self):
830 a = []
831 self.makecycle(groupby([a]*2, lambda x:x), a)
832
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000833 def test_issue2246(self):
834 # Issue 2246 -- the _grouper iterator was not included in GC
835 n = 10
836 keyfunc = lambda x: x
837 for i, j in groupby(xrange(n), key=keyfunc):
838 keyfunc.__dict__.setdefault('x',[]).append(j)
839
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000840 def test_ifilter(self):
841 a = []
842 self.makecycle(ifilter(lambda x:True, [a]*2), a)
843
844 def test_ifilterfalse(self):
845 a = []
846 self.makecycle(ifilterfalse(lambda x:False, a), a)
847
848 def test_izip(self):
849 a = []
850 self.makecycle(izip([a]*2, [a]*3), a)
851
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000852 def test_izip_longest(self):
853 a = []
854 self.makecycle(izip_longest([a]*2, [a]*3), a)
855 b = [a, None]
856 self.makecycle(izip_longest([a]*2, [a]*3, fillvalue=b), a)
857
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000858 def test_imap(self):
859 a = []
860 self.makecycle(imap(lambda x:x, [a]*2), a)
861
862 def test_islice(self):
863 a = []
864 self.makecycle(islice([a]*2, None), a)
865
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000866 def test_permutations(self):
867 a = []
868 self.makecycle(permutations([1,2,a,3], 3), a)
869
870 def test_product(self):
871 a = []
872 self.makecycle(product([1,2,a,3], repeat=3), a)
873
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000874 def test_repeat(self):
875 a = []
876 self.makecycle(repeat(a), a)
877
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000878 def test_starmap(self):
879 a = []
880 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
881
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000882 def test_takewhile(self):
883 a = []
884 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
885
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000886def R(seqn):
887 'Regular generator'
888 for i in seqn:
889 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000890
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000891class G:
892 'Sequence using __getitem__'
893 def __init__(self, seqn):
894 self.seqn = seqn
895 def __getitem__(self, i):
896 return self.seqn[i]
897
898class I:
899 'Sequence using iterator protocol'
900 def __init__(self, seqn):
901 self.seqn = seqn
902 self.i = 0
903 def __iter__(self):
904 return self
905 def next(self):
906 if self.i >= len(self.seqn): raise StopIteration
907 v = self.seqn[self.i]
908 self.i += 1
909 return v
910
911class Ig:
912 'Sequence using iterator protocol defined with a generator'
913 def __init__(self, seqn):
914 self.seqn = seqn
915 self.i = 0
916 def __iter__(self):
917 for val in self.seqn:
918 yield val
919
920class X:
921 'Missing __getitem__ and __iter__'
922 def __init__(self, seqn):
923 self.seqn = seqn
924 self.i = 0
925 def next(self):
926 if self.i >= len(self.seqn): raise StopIteration
927 v = self.seqn[self.i]
928 self.i += 1
929 return v
930
931class N:
932 'Iterator missing next()'
933 def __init__(self, seqn):
934 self.seqn = seqn
935 self.i = 0
936 def __iter__(self):
937 return self
938
939class E:
940 'Test propagation of exceptions'
941 def __init__(self, seqn):
942 self.seqn = seqn
943 self.i = 0
944 def __iter__(self):
945 return self
946 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +0000947 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000948
949class S:
950 'Test immediate stop'
951 def __init__(self, seqn):
952 pass
953 def __iter__(self):
954 return self
955 def next(self):
956 raise StopIteration
957
958def L(seqn):
959 'Test multiple tiers of iterators'
960 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
961
962
963class TestVariousIteratorArgs(unittest.TestCase):
964
965 def test_chain(self):
966 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
967 for g in (G, I, Ig, S, L, R):
968 self.assertEqual(list(chain(g(s))), list(g(s)))
969 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Raymond Hettinger05bf6332008-02-28 22:30:42 +0000970 self.assertRaises(TypeError, list, chain(X(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000971 self.assertRaises(TypeError, list, chain(N(s)))
972 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
973
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000974 def test_product(self):
975 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
976 self.assertRaises(TypeError, product, X(s))
977 self.assertRaises(TypeError, product, N(s))
978 self.assertRaises(ZeroDivisionError, product, E(s))
979
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000980 def test_cycle(self):
981 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
982 for g in (G, I, Ig, S, L, R):
983 tgtlen = len(s) * 3
984 expected = list(g(s))*3
985 actual = list(islice(cycle(g(s)), tgtlen))
986 self.assertEqual(actual, expected)
987 self.assertRaises(TypeError, cycle, X(s))
988 self.assertRaises(TypeError, list, cycle(N(s)))
989 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
990
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000991 def test_groupby(self):
992 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
993 for g in (G, I, Ig, S, L, R):
994 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
995 self.assertRaises(TypeError, groupby, X(s))
996 self.assertRaises(TypeError, list, groupby(N(s)))
997 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
998
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000999 def test_ifilter(self):
1000 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1001 for g in (G, I, Ig, S, L, R):
1002 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
1003 self.assertRaises(TypeError, ifilter, isEven, X(s))
1004 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
1005 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
1006
1007 def test_ifilterfalse(self):
1008 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1009 for g in (G, I, Ig, S, L, R):
1010 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
1011 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
1012 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
1013 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
1014
1015 def test_izip(self):
1016 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1017 for g in (G, I, Ig, S, L, R):
1018 self.assertEqual(list(izip(g(s))), zip(g(s)))
1019 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
1020 self.assertRaises(TypeError, izip, X(s))
1021 self.assertRaises(TypeError, list, izip(N(s)))
1022 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
1023
Raymond Hettingerd36862c2007-02-21 05:20:38 +00001024 def test_iziplongest(self):
1025 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1026 for g in (G, I, Ig, S, L, R):
1027 self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
1028 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
1029 self.assertRaises(TypeError, izip_longest, X(s))
1030 self.assertRaises(TypeError, list, izip_longest(N(s)))
1031 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
1032
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001033 def test_imap(self):
1034 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1035 for g in (G, I, Ig, S, L, R):
1036 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
1037 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
1038 self.assertRaises(TypeError, imap, onearg, X(s))
1039 self.assertRaises(TypeError, list, imap(onearg, N(s)))
1040 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
1041
1042 def test_islice(self):
1043 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1044 for g in (G, I, Ig, S, L, R):
1045 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1046 self.assertRaises(TypeError, islice, X(s), 10)
1047 self.assertRaises(TypeError, list, islice(N(s), 10))
1048 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1049
1050 def test_starmap(self):
1051 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1052 for g in (G, I, Ig, S, L, R):
1053 ss = zip(s, s)
1054 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
1055 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
1056 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
1057 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1058
1059 def test_takewhile(self):
1060 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1061 for g in (G, I, Ig, S, L, R):
1062 tgt = []
1063 for elem in g(s):
1064 if not isEven(elem): break
1065 tgt.append(elem)
1066 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1067 self.assertRaises(TypeError, takewhile, isEven, X(s))
1068 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
1069 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1070
1071 def test_dropwhile(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 tgt and isOdd(elem): continue
1077 tgt.append(elem)
1078 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1079 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
1080 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
1081 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1082
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001083 def test_tee(self):
1084 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1085 for g in (G, I, Ig, S, L, R):
1086 it1, it2 = tee(g(s))
1087 self.assertEqual(list(it1), list(g(s)))
1088 self.assertEqual(list(it2), list(g(s)))
1089 self.assertRaises(TypeError, tee, X(s))
1090 self.assertRaises(TypeError, list, tee(N(s))[0])
1091 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1092
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001093class LengthTransparency(unittest.TestCase):
1094
1095 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001096 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001097 self.assertEqual(len(repeat(None, 50)), 50)
1098 self.assertRaises(TypeError, len, repeat(None))
1099
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001100class RegressionTests(unittest.TestCase):
1101
1102 def test_sf_793826(self):
1103 # Fix Armin Rigo's successful efforts to wreak havoc
1104
1105 def mutatingtuple(tuple1, f, tuple2):
1106 # this builds a tuple t which is a copy of tuple1,
1107 # then calls f(t), then mutates t to be equal to tuple2
1108 # (needs len(tuple1) == len(tuple2)).
1109 def g(value, first=[1]):
1110 if first:
1111 del first[:]
1112 f(z.next())
1113 return value
1114 items = list(tuple2)
1115 items[1:1] = list(tuple1)
1116 gen = imap(g, items)
1117 z = izip(*[gen]*len(tuple1))
1118 z.next()
1119
1120 def f(t):
1121 global T
1122 T = t
1123 first[:] = list(T)
1124
1125 first = []
1126 mutatingtuple((1,2,3), f, (4,5,6))
1127 second = list(T)
1128 self.assertEqual(first, second)
1129
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001130
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001131 def test_sf_950057(self):
1132 # Make sure that chain() and cycle() catch exceptions immediately
1133 # rather than when shifting between input sources
1134
1135 def gen1():
1136 hist.append(0)
1137 yield 1
1138 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001139 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001140 hist.append(2)
1141
1142 def gen2(x):
1143 hist.append(3)
1144 yield 2
1145 hist.append(4)
1146 if x:
1147 raise StopIteration
1148
1149 hist = []
1150 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1151 self.assertEqual(hist, [0,1])
1152
1153 hist = []
1154 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1155 self.assertEqual(hist, [0,1])
1156
1157 hist = []
1158 self.assertRaises(AssertionError, list, cycle(gen1()))
1159 self.assertEqual(hist, [0,1])
1160
Georg Brandlb84c1372007-01-21 10:28:43 +00001161class SubclassWithKwargsTest(unittest.TestCase):
1162 def test_keywords_in_subclass(self):
1163 # count is not subclassable...
1164 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
1165 starmap, islice, takewhile, dropwhile, cycle):
1166 class Subclass(cls):
1167 def __init__(self, newarg=None, *args):
1168 cls.__init__(self, *args)
1169 try:
1170 Subclass(newarg=1)
1171 except TypeError, err:
1172 # we expect type errors because of wrong argument count
1173 self.failIf("does not take keyword arguments" in err.args[0])
1174
1175
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001176libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001177
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001178
1179>>> amounts = [120.15, 764.05, 823.14]
1180>>> for checknum, amount in izip(count(1200), amounts):
1181... print 'Check %d is for $%.2f' % (checknum, amount)
1182...
1183Check 1200 is for $120.15
1184Check 1201 is for $764.05
1185Check 1202 is for $823.14
1186
1187>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001188>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1189... print cube
1190...
11911
11928
119327
1194
1195>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001196>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001197... print name.title()
1198...
1199Alex
1200Laura
1201Martin
1202Walter
1203Samuele
1204
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001205>>> from operator import itemgetter
1206>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Armin Rigoa3f09272006-05-28 19:13:17 +00001207>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001208>>> for k, g in groupby(di, itemgetter(1)):
1209... print k, map(itemgetter(0), g)
1210...
12111 ['a', 'c', 'e']
12122 ['b', 'd', 'f']
12133 ['g']
1214
Raymond Hettinger734fb572004-01-20 20:04:40 +00001215# Find runs of consecutive numbers using groupby. The key to the solution
1216# is differencing with a range so that consecutive numbers all appear in
1217# same group.
1218>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
1219>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
1220... print map(operator.itemgetter(1), g)
1221...
1222[1]
1223[4, 5, 6]
1224[10]
1225[15, 16, 17, 18]
1226[22]
1227[25, 26, 27, 28]
1228
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001229>>> def take(n, iterable):
1230... "Return first n items of the iterable as a list"
1231... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001232
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001233>>> def enumerate(iterable, start=0):
1234... return izip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001235
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001236>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001237... "Return function(0), function(1), ..."
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001238... return imap(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001239
Raymond Hettinger5894c2b2009-02-19 05:38:53 +00001240>>> def nth(iterable, n, default=None):
1241... "Returns the nth item or a default value"
1242... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001243
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001244>>> def quantify(iterable, pred=bool):
1245... "Count how many times the predicate is true"
1246... return sum(imap(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001247
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001248>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001249... "Returns the sequence elements and then returns None indefinitely"
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001250... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001251
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001252>>> def ncycles(iterable, n):
1253... "Returns the seqeuence elements n times"
1254... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001255
1256>>> def dotproduct(vec1, vec2):
1257... return sum(imap(operator.mul, vec1, vec2))
1258
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001259>>> def flatten(listOfLists):
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001260... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001261
1262>>> def repeatfunc(func, times=None, *args):
1263... "Repeat calls to func with specified arguments."
1264... " Example: repeatfunc(random.random)"
1265... if times is None:
1266... return starmap(func, repeat(args))
1267... else:
1268... return starmap(func, repeat(args, times))
1269
Raymond Hettingerd591f662003-10-26 15:34:50 +00001270>>> def pairwise(iterable):
1271... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1272... a, b = tee(iterable)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001273... for elem in b:
1274... break
Raymond Hettingerad983e72003-11-12 14:32:26 +00001275... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001276
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001277>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001278... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001279... args = [iter(iterable)] * n
Raymond Hettingerf080e6d2008-07-31 01:19:50 +00001280... return izip_longest(fillvalue=fillvalue, *args)
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001281
1282>>> def roundrobin(*iterables):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001283... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001284... # Recipe credited to George Sakkis
1285... pending = len(iterables)
1286... nexts = cycle(iter(it).next for it in iterables)
1287... while pending:
1288... try:
1289... for next in nexts:
1290... yield next()
1291... except StopIteration:
1292... pending -= 1
1293... nexts = cycle(islice(nexts, pending))
1294
1295>>> def powerset(iterable):
Raymond Hettinger5cd012b2009-01-26 02:17:52 +00001296... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1297... s = list(iterable)
1298... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001299
Raymond Hettingere8b4b602008-03-11 00:19:07 +00001300>>> def compress(data, selectors):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001301... "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
1302... return (d for d, s in izip(data, selectors) if s)
Raymond Hettingere8b4b602008-03-11 00:19:07 +00001303
Raymond Hettinger33691672008-07-19 00:43:00 +00001304>>> def combinations_with_replacement(iterable, r):
1305... "combinations_with_replacement('ABC', 3) --> AA AB AC BB BC CC"
1306... pool = tuple(iterable)
1307... n = len(pool)
Raymond Hettingercdc9f2c2009-01-27 06:38:00 +00001308... if not n and r:
1309... return
Raymond Hettinger33691672008-07-19 00:43:00 +00001310... indices = [0] * r
1311... yield tuple(pool[i] for i in indices)
1312... while 1:
1313... for i in reversed(range(r)):
1314... if indices[i] != n - 1:
1315... break
1316... else:
1317... return
1318... indices[i:] = [indices[i] + 1] * (r - i)
1319... yield tuple(pool[i] for i in indices)
1320
Raymond Hettingerd8461652009-01-02 21:20:38 +00001321>>> def unique_everseen(iterable, key=None):
1322... "List unique elements, preserving order. Remember all elements ever seen."
1323... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1324... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1325... seen = set()
1326... seen_add = seen.add
1327... if key is None:
1328... for element in iterable:
1329... if element not in seen:
1330... seen_add(element)
1331... yield element
1332... else:
1333... for element in iterable:
1334... k = key(element)
1335... if k not in seen:
1336... seen_add(k)
1337... yield element
1338
1339>>> def unique_justseen(iterable, key=None):
1340... "List unique elements, preserving order. Remember only the element just seen."
1341... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1342... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1343... return imap(next, imap(itemgetter(1), groupby(iterable, key)))
1344
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001345This is not part of the examples but it tests to make sure the definitions
1346perform as purported.
1347
Raymond Hettingera098b332003-09-08 23:58:40 +00001348>>> take(10, count())
1349[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1350
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001351>>> list(enumerate('abc'))
1352[(0, 'a'), (1, 'b'), (2, 'c')]
1353
1354>>> list(islice(tabulate(lambda x: 2*x), 4))
1355[0, 2, 4, 6]
1356
1357>>> nth('abcde', 3)
Raymond Hettinger11485b42009-02-04 19:34:31 +00001358'd'
1359
1360>>> nth('abcde', 9) is None
1361True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001362
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001363>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000136450
1365
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001366>>> a = [[1, 2, 3], [4, 5, 6]]
1367>>> flatten(a)
1368[1, 2, 3, 4, 5, 6]
1369
1370>>> list(repeatfunc(pow, 5, 2, 3))
1371[8, 8, 8, 8, 8]
1372
1373>>> import random
1374>>> take(5, imap(int, repeatfunc(random.random)))
1375[0, 0, 0, 0, 0]
1376
Raymond Hettingerd591f662003-10-26 15:34:50 +00001377>>> list(pairwise('abcd'))
1378[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001379
Raymond Hettingerd591f662003-10-26 15:34:50 +00001380>>> list(pairwise([]))
1381[]
1382
1383>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001384[]
1385
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001386>>> list(islice(padnone('abc'), 0, 6))
1387['a', 'b', 'c', None, None, None]
1388
1389>>> list(ncycles('abc', 3))
1390['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1391
1392>>> dotproduct([1,2,3], [4,5,6])
139332
1394
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001395>>> list(grouper(3, 'abcdefg', 'x'))
1396[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1397
1398>>> list(roundrobin('abc', 'd', 'ef'))
1399['a', 'd', 'e', 'b', 'f', 'c']
1400
Raymond Hettinger5cd012b2009-01-26 02:17:52 +00001401>>> list(powerset([1,2,3]))
1402[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001403
Raymond Hettingere8b4b602008-03-11 00:19:07 +00001404>>> list(compress('abcdef', [1,0,1,0,1,1]))
1405['a', 'c', 'e', 'f']
1406
Raymond Hettinger33691672008-07-19 00:43:00 +00001407>>> list(combinations_with_replacement('abc', 2))
1408[('a', 'a'), ('a', 'b'), ('a', 'c'), ('b', 'b'), ('b', 'c'), ('c', 'c')]
1409
Raymond Hettinger825758c2009-01-08 05:20:19 +00001410>>> list(combinations_with_replacement('01', 3))
1411[('0', '0', '0'), ('0', '0', '1'), ('0', '1', '1'), ('1', '1', '1')]
1412
1413>>> def combinations_with_replacement2(iterable, r):
1414... 'Alternate version that filters from product()'
1415... pool = tuple(iterable)
1416... n = len(pool)
1417... for indices in product(range(n), repeat=r):
1418... if sorted(indices) == list(indices):
1419... yield tuple(pool[i] for i in indices)
1420
1421>>> list(combinations_with_replacement('abc', 2)) == list(combinations_with_replacement2('abc', 2))
1422True
1423
1424>>> list(combinations_with_replacement('01', 3)) == list(combinations_with_replacement2('01', 3))
1425True
1426
1427>>> list(combinations_with_replacement('2310', 6)) == list(combinations_with_replacement2('2310', 6))
1428True
1429
Raymond Hettingerd8461652009-01-02 21:20:38 +00001430>>> list(unique_everseen('AAAABBBCCDAABBB'))
1431['A', 'B', 'C', 'D']
1432
1433>>> list(unique_everseen('ABBCcAD', str.lower))
1434['A', 'B', 'C', 'D']
1435
1436>>> list(unique_justseen('AAAABBBCCDAABBB'))
1437['A', 'B', 'C', 'D', 'A', 'B']
1438
1439>>> list(unique_justseen('ABBCcAD', str.lower))
1440['A', 'B', 'C', 'A', 'D']
1441
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001442"""
1443
1444__test__ = {'libreftest' : libreftest}
1445
1446def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001447 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Georg Brandlb84c1372007-01-21 10:28:43 +00001448 RegressionTests, LengthTransparency,
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001449 SubclassWithKwargsTest, TestExamples)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001450 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001451
1452 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001453 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001454 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001455 counts = [None] * 5
1456 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001457 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001458 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001459 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001460 print counts
1461
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001462 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001463 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001464
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001465if __name__ == "__main__":
1466 test_main(verbose=True)