blob: dbc941fd77ea709f234e3924aafb7a504aec995a [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):
Raymond Hettinger4da5faa2009-11-01 18:43:31 +0000425 with test_support.captured_output('stdout'):
426 print (i, j)
Raymond Hettinger80d49b32009-11-01 08:53:21 +0000427 result.append((i, j))
428 return result
429 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
430
431 # Formerly, the RuntimeError would be lost
432 # and StopIteration would stop as expected
433 r1 = Repeater(1, 3, RuntimeError)
434 r2 = Repeater(2, 4, StopIteration)
Raymond Hettinger4da5faa2009-11-01 18:43:31 +0000435 it = izip_longest(r1, r2, fillvalue=0)
436 self.assertEqual(next(it), (1, 2))
437 self.assertEqual(next(it), (1, 2))
438 self.assertEqual(next(it), (1, 2))
439 self.assertRaises(RuntimeError, next, it)
Raymond Hettinger80d49b32009-11-01 08:53:21 +0000440
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000441 def test_product(self):
442 for args, result in [
Raymond Hettingerd553d852008-03-04 04:17:08 +0000443 ([], [()]), # zero iterables
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000444 (['ab'], [('a',), ('b',)]), # one iterable
445 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
446 ([range(0), range(2), range(3)], []), # first iterable with zero length
447 ([range(2), range(0), range(3)], []), # middle iterable with zero length
448 ([range(2), range(3), range(0)], []), # last iterable with zero length
449 ]:
450 self.assertEqual(list(product(*args)), result)
Raymond Hettinger08ff6822008-02-29 02:21:48 +0000451 for r in range(4):
452 self.assertEqual(list(product(*(args*r))),
453 list(product(*args, **dict(repeat=r))))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000454 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
455 self.assertRaises(TypeError, product, range(6), None)
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000456
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000457 def product1(*args, **kwds):
458 pools = map(tuple, args) * kwds.get('repeat', 1)
459 n = len(pools)
460 if n == 0:
461 yield ()
462 return
463 if any(len(pool) == 0 for pool in pools):
464 return
465 indices = [0] * n
466 yield tuple(pool[i] for pool, i in zip(pools, indices))
467 while 1:
468 for i in reversed(range(n)): # right to left
469 if indices[i] == len(pools[i]) - 1:
470 continue
471 indices[i] += 1
472 for j in range(i+1, n):
473 indices[j] = 0
474 yield tuple(pool[i] for pool, i in zip(pools, indices))
475 break
476 else:
477 return
478
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000479 def product2(*args, **kwds):
480 'Pure python version used in docs'
481 pools = map(tuple, args) * kwds.get('repeat', 1)
482 result = [[]]
483 for pool in pools:
484 result = [x+[y] for x in result for y in pool]
485 for prod in result:
486 yield tuple(prod)
487
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000488 argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
489 set('abcdefg'), range(11), tuple(range(13))]
490 for i in range(100):
491 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Raymond Hettingerd553d852008-03-04 04:17:08 +0000492 expected_len = prod(map(len, args))
493 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000494 self.assertEqual(list(product(*args)), list(product1(*args)))
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000495 self.assertEqual(list(product(*args)), list(product2(*args)))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000496 args = map(iter, args)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000497 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000498
Raymond Hettinger73d79632008-02-23 02:20:41 +0000499 # Test implementation detail: tuple re-use
500 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
501 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000502
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000503 def test_repeat(self):
504 self.assertEqual(zip(xrange(3),repeat('a')),
505 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000506 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000507 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000508 self.assertEqual(list(repeat('a', 0)), [])
509 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000510 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000511 self.assertRaises(TypeError, repeat, None, 3, 4)
512 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000513 r = repeat(1+0j)
514 self.assertEqual(repr(r), 'repeat((1+0j))')
515 r = repeat(1+0j, 5)
516 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
517 list(r)
518 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000519
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000520 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000521 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
522 [0**1, 1**2, 2**3])
523 self.assertEqual(list(imap(None, 'abc', range(5))),
524 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000525 self.assertEqual(list(imap(None, 'abc', count())),
526 [('a',0),('b',1),('c',2)])
527 self.assertEqual(take(2,imap(None, 'abc', count())),
528 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000529 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000530 self.assertRaises(TypeError, imap)
531 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000532 self.assertRaises(TypeError, imap(10, range(5)).next)
533 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
534 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000535
536 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000537 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
538 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000539 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
540 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000541 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger47317092008-01-17 03:02:14 +0000542 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
543 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000544 self.assertRaises(TypeError, starmap)
545 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
546 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
547 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
548 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000549
550 def test_islice(self):
551 for args in [ # islice(args) should agree with range(args)
552 (10, 20, 3),
553 (10, 3, 20),
554 (10, 20),
555 (10, 3),
556 (20,)
557 ]:
558 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
559
560 for args, tgtargs in [ # Stop when seqn is exhausted
561 ((10, 110, 3), ((10, 100, 3))),
562 ((10, 110), ((10, 100))),
563 ((110,), (100,))
564 ]:
565 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
566
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000567 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000568 self.assertEqual(list(islice(xrange(10), None)), range(10))
Raymond Hettingerb2594052004-12-05 09:25:51 +0000569 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
570 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000571 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
572 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
573
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000574 # Test number of items consumed SF #1171417
575 it = iter(range(10))
576 self.assertEqual(list(islice(it, 3)), range(3))
577 self.assertEqual(list(it), range(3, 10))
578
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000579 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000580 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000581 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
582 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
583 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
584 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
585 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000586 self.assertRaises(ValueError, islice, xrange(10), 'a')
587 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
588 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
589 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
590 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +0000591 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000592
593 def test_takewhile(self):
594 data = [1, 3, 5, 20, 2, 4, 6, 8]
595 underten = lambda x: x<10
596 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000597 self.assertEqual(list(takewhile(underten, [])), [])
598 self.assertRaises(TypeError, takewhile)
599 self.assertRaises(TypeError, takewhile, operator.pow)
600 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
601 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
602 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000603 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
604 self.assertEqual(list(t), [1, 1, 1])
605 self.assertRaises(StopIteration, t.next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000606
607 def test_dropwhile(self):
608 data = [1, 3, 5, 20, 2, 4, 6, 8]
609 underten = lambda x: x<10
610 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000611 self.assertEqual(list(dropwhile(underten, [])), [])
612 self.assertRaises(TypeError, dropwhile)
613 self.assertRaises(TypeError, dropwhile, operator.pow)
614 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
615 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
616 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000617
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000618 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000619 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000620 def irange(n):
621 for i in xrange(n):
622 yield i
623
624 a, b = tee([]) # test empty iterator
625 self.assertEqual(list(a), [])
626 self.assertEqual(list(b), [])
627
628 a, b = tee(irange(n)) # test 100% interleaved
629 self.assertEqual(zip(a,b), zip(range(n),range(n)))
630
631 a, b = tee(irange(n)) # test 0% interleaved
632 self.assertEqual(list(a), range(n))
633 self.assertEqual(list(b), range(n))
634
635 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000636 for i in xrange(100):
637 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000638 del a
639 self.assertEqual(list(b), range(n))
640
641 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000642 for i in xrange(100):
643 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000644 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000645 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000646
647 for j in xrange(5): # test randomly interleaved
648 order = [0]*n + [1]*n
649 random.shuffle(order)
650 lists = ([], [])
651 its = tee(irange(n))
652 for i in order:
653 value = its[i].next()
654 lists[i].append(value)
655 self.assertEqual(lists[0], range(n))
656 self.assertEqual(lists[1], range(n))
657
Raymond Hettingerad983e72003-11-12 14:32:26 +0000658 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000659 self.assertRaises(TypeError, tee)
660 self.assertRaises(TypeError, tee, 3)
661 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000662 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000663
Raymond Hettingerad983e72003-11-12 14:32:26 +0000664 # tee object should be instantiable
665 a, b = tee('abc')
666 c = type(a)('def')
667 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000668
Raymond Hettingerad983e72003-11-12 14:32:26 +0000669 # test long-lagged and multi-way split
670 a, b, c = tee(xrange(2000), 3)
671 for i in xrange(100):
672 self.assertEqual(a.next(), i)
673 self.assertEqual(list(b), range(2000))
674 self.assertEqual([c.next(), c.next()], range(2))
675 self.assertEqual(list(a), range(100,2000))
676 self.assertEqual(list(c), range(2,2000))
677
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000678 # test values of n
679 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Neal Norwitz69e88972006-09-02 02:58:13 +0000680 self.assertRaises(ValueError, tee, [], -1)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000681 for n in xrange(5):
682 result = tee('abc', n)
683 self.assertEqual(type(result), tuple)
684 self.assertEqual(len(result), n)
685 self.assertEqual(map(list, result), [list('abc')]*n)
686
Raymond Hettingerad983e72003-11-12 14:32:26 +0000687 # tee pass-through to copyable iterator
688 a, b = tee('abc')
689 c, d = tee(a)
690 self.assert_(a is c)
691
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000692 # test tee_new
693 t1, t2 = tee('abc')
694 tnew = type(t1)
695 self.assertRaises(TypeError, tnew)
696 self.assertRaises(TypeError, tnew, 10)
697 t3 = tnew(t1)
698 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000699
Raymond Hettingera9f60922004-10-17 16:40:14 +0000700 # test that tee objects are weak referencable
701 a, b = tee(xrange(10))
702 p = proxy(a)
703 self.assertEqual(getattr(p, '__class__'), type(b))
704 del a
705 self.assertRaises(ReferenceError, getattr, p, '__class__')
706
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000707 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000708 self.assertRaises(StopIteration, izip().next)
709
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000710 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000711 self.assertRaises(StopIteration, f([]).next)
712 self.assertRaises(StopIteration, f(StopNow()).next)
713
714 self.assertRaises(StopIteration, islice([], None).next)
715 self.assertRaises(StopIteration, islice(StopNow(), None).next)
716
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000717 p, q = tee([])
718 self.assertRaises(StopIteration, p.next)
719 self.assertRaises(StopIteration, q.next)
720 p, q = tee(StopNow())
721 self.assertRaises(StopIteration, p.next)
722 self.assertRaises(StopIteration, q.next)
723
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000724 self.assertRaises(StopIteration, repeat(None, 0).next)
725
726 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
727 self.assertRaises(StopIteration, f(lambda x:x, []).next)
728 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
729
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000730class TestExamples(unittest.TestCase):
731
732 def test_chain(self):
733 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
734
735 def test_chain_from_iterable(self):
736 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
737
738 def test_combinations(self):
739 self.assertEqual(list(combinations('ABCD', 2)),
740 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
741 self.assertEqual(list(combinations(range(4), 3)),
742 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
743
744 def test_count(self):
745 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
746
747 def test_cycle(self):
748 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
749
750 def test_dropwhile(self):
751 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
752
753 def test_groupby(self):
754 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
755 list('ABCDAB'))
756 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
757 [list('AAAA'), list('BBB'), list('CC'), list('D')])
758
759 def test_ifilter(self):
760 self.assertEqual(list(ifilter(lambda x: x%2, range(10))), [1,3,5,7,9])
761
762 def test_ifilterfalse(self):
763 self.assertEqual(list(ifilterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
764
765 def test_imap(self):
766 self.assertEqual(list(imap(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
767
768 def test_islice(self):
769 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
770 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
771 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
772 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
773
774 def test_izip(self):
775 self.assertEqual(list(izip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
776
777 def test_izip_longest(self):
778 self.assertEqual(list(izip_longest('ABCD', 'xy', fillvalue='-')),
779 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
780
781 def test_permutations(self):
782 self.assertEqual(list(permutations('ABCD', 2)),
783 map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split()))
784 self.assertEqual(list(permutations(range(3))),
785 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
786
787 def test_product(self):
788 self.assertEqual(list(product('ABCD', 'xy')),
789 map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split()))
790 self.assertEqual(list(product(range(2), repeat=3)),
791 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
792 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
793
794 def test_repeat(self):
795 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
796
797 def test_stapmap(self):
798 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
799 [32, 9, 1000])
800
801 def test_takewhile(self):
802 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
803
804
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000805class TestGC(unittest.TestCase):
806
807 def makecycle(self, iterator, container):
808 container.append(iterator)
809 iterator.next()
810 del container, iterator
811
812 def test_chain(self):
813 a = []
814 self.makecycle(chain(a), a)
815
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000816 def test_chain_from_iterable(self):
817 a = []
818 self.makecycle(chain.from_iterable([a]), a)
819
820 def test_combinations(self):
821 a = []
822 self.makecycle(combinations([1,2,a,3], 3), a)
823
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000824 def test_cycle(self):
825 a = []
826 self.makecycle(cycle([a]*2), a)
827
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000828 def test_dropwhile(self):
829 a = []
830 self.makecycle(dropwhile(bool, [0, a, a]), a)
831
832 def test_groupby(self):
833 a = []
834 self.makecycle(groupby([a]*2, lambda x:x), a)
835
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000836 def test_issue2246(self):
837 # Issue 2246 -- the _grouper iterator was not included in GC
838 n = 10
839 keyfunc = lambda x: x
840 for i, j in groupby(xrange(n), key=keyfunc):
841 keyfunc.__dict__.setdefault('x',[]).append(j)
842
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000843 def test_ifilter(self):
844 a = []
845 self.makecycle(ifilter(lambda x:True, [a]*2), a)
846
847 def test_ifilterfalse(self):
848 a = []
849 self.makecycle(ifilterfalse(lambda x:False, a), a)
850
851 def test_izip(self):
852 a = []
853 self.makecycle(izip([a]*2, [a]*3), a)
854
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000855 def test_izip_longest(self):
856 a = []
857 self.makecycle(izip_longest([a]*2, [a]*3), a)
858 b = [a, None]
859 self.makecycle(izip_longest([a]*2, [a]*3, fillvalue=b), a)
860
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000861 def test_imap(self):
862 a = []
863 self.makecycle(imap(lambda x:x, [a]*2), a)
864
865 def test_islice(self):
866 a = []
867 self.makecycle(islice([a]*2, None), a)
868
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000869 def test_permutations(self):
870 a = []
871 self.makecycle(permutations([1,2,a,3], 3), a)
872
873 def test_product(self):
874 a = []
875 self.makecycle(product([1,2,a,3], repeat=3), a)
876
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000877 def test_repeat(self):
878 a = []
879 self.makecycle(repeat(a), a)
880
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000881 def test_starmap(self):
882 a = []
883 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
884
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000885 def test_takewhile(self):
886 a = []
887 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
888
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000889def R(seqn):
890 'Regular generator'
891 for i in seqn:
892 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000893
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000894class G:
895 'Sequence using __getitem__'
896 def __init__(self, seqn):
897 self.seqn = seqn
898 def __getitem__(self, i):
899 return self.seqn[i]
900
901class I:
902 'Sequence using iterator protocol'
903 def __init__(self, seqn):
904 self.seqn = seqn
905 self.i = 0
906 def __iter__(self):
907 return self
908 def next(self):
909 if self.i >= len(self.seqn): raise StopIteration
910 v = self.seqn[self.i]
911 self.i += 1
912 return v
913
914class Ig:
915 'Sequence using iterator protocol defined with a generator'
916 def __init__(self, seqn):
917 self.seqn = seqn
918 self.i = 0
919 def __iter__(self):
920 for val in self.seqn:
921 yield val
922
923class X:
924 'Missing __getitem__ and __iter__'
925 def __init__(self, seqn):
926 self.seqn = seqn
927 self.i = 0
928 def next(self):
929 if self.i >= len(self.seqn): raise StopIteration
930 v = self.seqn[self.i]
931 self.i += 1
932 return v
933
934class N:
935 'Iterator missing next()'
936 def __init__(self, seqn):
937 self.seqn = seqn
938 self.i = 0
939 def __iter__(self):
940 return self
941
942class E:
943 'Test propagation of exceptions'
944 def __init__(self, seqn):
945 self.seqn = seqn
946 self.i = 0
947 def __iter__(self):
948 return self
949 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +0000950 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000951
952class S:
953 'Test immediate stop'
954 def __init__(self, seqn):
955 pass
956 def __iter__(self):
957 return self
958 def next(self):
959 raise StopIteration
960
961def L(seqn):
962 'Test multiple tiers of iterators'
963 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
964
965
966class TestVariousIteratorArgs(unittest.TestCase):
967
968 def test_chain(self):
969 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
970 for g in (G, I, Ig, S, L, R):
971 self.assertEqual(list(chain(g(s))), list(g(s)))
972 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Raymond Hettinger05bf6332008-02-28 22:30:42 +0000973 self.assertRaises(TypeError, list, chain(X(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000974 self.assertRaises(TypeError, list, chain(N(s)))
975 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
976
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000977 def test_product(self):
978 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
979 self.assertRaises(TypeError, product, X(s))
980 self.assertRaises(TypeError, product, N(s))
981 self.assertRaises(ZeroDivisionError, product, E(s))
982
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000983 def test_cycle(self):
984 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
985 for g in (G, I, Ig, S, L, R):
986 tgtlen = len(s) * 3
987 expected = list(g(s))*3
988 actual = list(islice(cycle(g(s)), tgtlen))
989 self.assertEqual(actual, expected)
990 self.assertRaises(TypeError, cycle, X(s))
991 self.assertRaises(TypeError, list, cycle(N(s)))
992 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
993
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000994 def test_groupby(self):
995 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
996 for g in (G, I, Ig, S, L, R):
997 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
998 self.assertRaises(TypeError, groupby, X(s))
999 self.assertRaises(TypeError, list, groupby(N(s)))
1000 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1001
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001002 def test_ifilter(self):
1003 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1004 for g in (G, I, Ig, S, L, R):
1005 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
1006 self.assertRaises(TypeError, ifilter, isEven, X(s))
1007 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
1008 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
1009
1010 def test_ifilterfalse(self):
1011 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1012 for g in (G, I, Ig, S, L, R):
1013 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
1014 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
1015 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
1016 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
1017
1018 def test_izip(self):
1019 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1020 for g in (G, I, Ig, S, L, R):
1021 self.assertEqual(list(izip(g(s))), zip(g(s)))
1022 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
1023 self.assertRaises(TypeError, izip, X(s))
1024 self.assertRaises(TypeError, list, izip(N(s)))
1025 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
1026
Raymond Hettingerd36862c2007-02-21 05:20:38 +00001027 def test_iziplongest(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_longest(g(s))), zip(g(s)))
1031 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
1032 self.assertRaises(TypeError, izip_longest, X(s))
1033 self.assertRaises(TypeError, list, izip_longest(N(s)))
1034 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
1035
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001036 def test_imap(self):
1037 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1038 for g in (G, I, Ig, S, L, R):
1039 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
1040 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
1041 self.assertRaises(TypeError, imap, onearg, X(s))
1042 self.assertRaises(TypeError, list, imap(onearg, N(s)))
1043 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
1044
1045 def test_islice(self):
1046 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1047 for g in (G, I, Ig, S, L, R):
1048 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1049 self.assertRaises(TypeError, islice, X(s), 10)
1050 self.assertRaises(TypeError, list, islice(N(s), 10))
1051 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1052
1053 def test_starmap(self):
1054 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1055 for g in (G, I, Ig, S, L, R):
1056 ss = zip(s, s)
1057 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
1058 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
1059 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
1060 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1061
1062 def test_takewhile(self):
1063 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1064 for g in (G, I, Ig, S, L, R):
1065 tgt = []
1066 for elem in g(s):
1067 if not isEven(elem): break
1068 tgt.append(elem)
1069 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1070 self.assertRaises(TypeError, takewhile, isEven, X(s))
1071 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
1072 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1073
1074 def test_dropwhile(self):
1075 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1076 for g in (G, I, Ig, S, L, R):
1077 tgt = []
1078 for elem in g(s):
1079 if not tgt and isOdd(elem): continue
1080 tgt.append(elem)
1081 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1082 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
1083 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
1084 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1085
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001086 def test_tee(self):
1087 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1088 for g in (G, I, Ig, S, L, R):
1089 it1, it2 = tee(g(s))
1090 self.assertEqual(list(it1), list(g(s)))
1091 self.assertEqual(list(it2), list(g(s)))
1092 self.assertRaises(TypeError, tee, X(s))
1093 self.assertRaises(TypeError, list, tee(N(s))[0])
1094 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1095
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001096class LengthTransparency(unittest.TestCase):
1097
1098 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001099 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001100 self.assertEqual(len(repeat(None, 50)), 50)
1101 self.assertRaises(TypeError, len, repeat(None))
1102
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001103class RegressionTests(unittest.TestCase):
1104
1105 def test_sf_793826(self):
1106 # Fix Armin Rigo's successful efforts to wreak havoc
1107
1108 def mutatingtuple(tuple1, f, tuple2):
1109 # this builds a tuple t which is a copy of tuple1,
1110 # then calls f(t), then mutates t to be equal to tuple2
1111 # (needs len(tuple1) == len(tuple2)).
1112 def g(value, first=[1]):
1113 if first:
1114 del first[:]
1115 f(z.next())
1116 return value
1117 items = list(tuple2)
1118 items[1:1] = list(tuple1)
1119 gen = imap(g, items)
1120 z = izip(*[gen]*len(tuple1))
1121 z.next()
1122
1123 def f(t):
1124 global T
1125 T = t
1126 first[:] = list(T)
1127
1128 first = []
1129 mutatingtuple((1,2,3), f, (4,5,6))
1130 second = list(T)
1131 self.assertEqual(first, second)
1132
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001133
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001134 def test_sf_950057(self):
1135 # Make sure that chain() and cycle() catch exceptions immediately
1136 # rather than when shifting between input sources
1137
1138 def gen1():
1139 hist.append(0)
1140 yield 1
1141 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001142 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001143 hist.append(2)
1144
1145 def gen2(x):
1146 hist.append(3)
1147 yield 2
1148 hist.append(4)
1149 if x:
1150 raise StopIteration
1151
1152 hist = []
1153 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1154 self.assertEqual(hist, [0,1])
1155
1156 hist = []
1157 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1158 self.assertEqual(hist, [0,1])
1159
1160 hist = []
1161 self.assertRaises(AssertionError, list, cycle(gen1()))
1162 self.assertEqual(hist, [0,1])
1163
Georg Brandlb84c1372007-01-21 10:28:43 +00001164class SubclassWithKwargsTest(unittest.TestCase):
1165 def test_keywords_in_subclass(self):
1166 # count is not subclassable...
1167 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
1168 starmap, islice, takewhile, dropwhile, cycle):
1169 class Subclass(cls):
1170 def __init__(self, newarg=None, *args):
1171 cls.__init__(self, *args)
1172 try:
1173 Subclass(newarg=1)
1174 except TypeError, err:
1175 # we expect type errors because of wrong argument count
1176 self.failIf("does not take keyword arguments" in err.args[0])
1177
1178
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001179libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001180
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001181
1182>>> amounts = [120.15, 764.05, 823.14]
1183>>> for checknum, amount in izip(count(1200), amounts):
1184... print 'Check %d is for $%.2f' % (checknum, amount)
1185...
1186Check 1200 is for $120.15
1187Check 1201 is for $764.05
1188Check 1202 is for $823.14
1189
1190>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001191>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1192... print cube
1193...
11941
11958
119627
1197
1198>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001199>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001200... print name.title()
1201...
1202Alex
1203Laura
1204Martin
1205Walter
1206Samuele
1207
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001208>>> from operator import itemgetter
1209>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Armin Rigoa3f09272006-05-28 19:13:17 +00001210>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001211>>> for k, g in groupby(di, itemgetter(1)):
1212... print k, map(itemgetter(0), g)
1213...
12141 ['a', 'c', 'e']
12152 ['b', 'd', 'f']
12163 ['g']
1217
Raymond Hettinger734fb572004-01-20 20:04:40 +00001218# Find runs of consecutive numbers using groupby. The key to the solution
1219# is differencing with a range so that consecutive numbers all appear in
1220# same group.
1221>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
1222>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
1223... print map(operator.itemgetter(1), g)
1224...
1225[1]
1226[4, 5, 6]
1227[10]
1228[15, 16, 17, 18]
1229[22]
1230[25, 26, 27, 28]
1231
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001232>>> def take(n, iterable):
1233... "Return first n items of the iterable as a list"
1234... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001235
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001236>>> def enumerate(iterable, start=0):
1237... return izip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001238
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001239>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001240... "Return function(0), function(1), ..."
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001241... return imap(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001242
Raymond Hettinger5894c2b2009-02-19 05:38:53 +00001243>>> def nth(iterable, n, default=None):
1244... "Returns the nth item or a default value"
1245... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001246
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001247>>> def quantify(iterable, pred=bool):
1248... "Count how many times the predicate is true"
1249... return sum(imap(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001250
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001251>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001252... "Returns the sequence elements and then returns None indefinitely"
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001253... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001254
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001255>>> def ncycles(iterable, n):
1256... "Returns the seqeuence elements n times"
1257... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001258
1259>>> def dotproduct(vec1, vec2):
1260... return sum(imap(operator.mul, vec1, vec2))
1261
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001262>>> def flatten(listOfLists):
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001263... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001264
1265>>> def repeatfunc(func, times=None, *args):
1266... "Repeat calls to func with specified arguments."
1267... " Example: repeatfunc(random.random)"
1268... if times is None:
1269... return starmap(func, repeat(args))
1270... else:
1271... return starmap(func, repeat(args, times))
1272
Raymond Hettingerd591f662003-10-26 15:34:50 +00001273>>> def pairwise(iterable):
1274... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1275... a, b = tee(iterable)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001276... for elem in b:
1277... break
Raymond Hettingerad983e72003-11-12 14:32:26 +00001278... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001279
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001280>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001281... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001282... args = [iter(iterable)] * n
Raymond Hettingerf080e6d2008-07-31 01:19:50 +00001283... return izip_longest(fillvalue=fillvalue, *args)
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001284
1285>>> def roundrobin(*iterables):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001286... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001287... # Recipe credited to George Sakkis
1288... pending = len(iterables)
1289... nexts = cycle(iter(it).next for it in iterables)
1290... while pending:
1291... try:
1292... for next in nexts:
1293... yield next()
1294... except StopIteration:
1295... pending -= 1
1296... nexts = cycle(islice(nexts, pending))
1297
1298>>> def powerset(iterable):
Raymond Hettinger5cd012b2009-01-26 02:17:52 +00001299... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1300... s = list(iterable)
1301... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001302
Raymond Hettingere8b4b602008-03-11 00:19:07 +00001303>>> def compress(data, selectors):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001304... "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
1305... return (d for d, s in izip(data, selectors) if s)
Raymond Hettingere8b4b602008-03-11 00:19:07 +00001306
Raymond Hettinger33691672008-07-19 00:43:00 +00001307>>> def combinations_with_replacement(iterable, r):
1308... "combinations_with_replacement('ABC', 3) --> AA AB AC BB BC CC"
1309... pool = tuple(iterable)
1310... n = len(pool)
Raymond Hettingercdc9f2c2009-01-27 06:38:00 +00001311... if not n and r:
1312... return
Raymond Hettinger33691672008-07-19 00:43:00 +00001313... indices = [0] * r
1314... yield tuple(pool[i] for i in indices)
1315... while 1:
1316... for i in reversed(range(r)):
1317... if indices[i] != n - 1:
1318... break
1319... else:
1320... return
1321... indices[i:] = [indices[i] + 1] * (r - i)
1322... yield tuple(pool[i] for i in indices)
1323
Raymond Hettingerd8461652009-01-02 21:20:38 +00001324>>> def unique_everseen(iterable, key=None):
1325... "List unique elements, preserving order. Remember all elements ever seen."
1326... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1327... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1328... seen = set()
1329... seen_add = seen.add
1330... if key is None:
1331... for element in iterable:
1332... if element not in seen:
1333... seen_add(element)
1334... yield element
1335... else:
1336... for element in iterable:
1337... k = key(element)
1338... if k not in seen:
1339... seen_add(k)
1340... yield element
1341
1342>>> def unique_justseen(iterable, key=None):
1343... "List unique elements, preserving order. Remember only the element just seen."
1344... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1345... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1346... return imap(next, imap(itemgetter(1), groupby(iterable, key)))
1347
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001348This is not part of the examples but it tests to make sure the definitions
1349perform as purported.
1350
Raymond Hettingera098b332003-09-08 23:58:40 +00001351>>> take(10, count())
1352[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1353
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001354>>> list(enumerate('abc'))
1355[(0, 'a'), (1, 'b'), (2, 'c')]
1356
1357>>> list(islice(tabulate(lambda x: 2*x), 4))
1358[0, 2, 4, 6]
1359
1360>>> nth('abcde', 3)
Raymond Hettinger11485b42009-02-04 19:34:31 +00001361'd'
1362
1363>>> nth('abcde', 9) is None
1364True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001365
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001366>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000136750
1368
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001369>>> a = [[1, 2, 3], [4, 5, 6]]
1370>>> flatten(a)
1371[1, 2, 3, 4, 5, 6]
1372
1373>>> list(repeatfunc(pow, 5, 2, 3))
1374[8, 8, 8, 8, 8]
1375
1376>>> import random
1377>>> take(5, imap(int, repeatfunc(random.random)))
1378[0, 0, 0, 0, 0]
1379
Raymond Hettingerd591f662003-10-26 15:34:50 +00001380>>> list(pairwise('abcd'))
1381[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001382
Raymond Hettingerd591f662003-10-26 15:34:50 +00001383>>> list(pairwise([]))
1384[]
1385
1386>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001387[]
1388
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001389>>> list(islice(padnone('abc'), 0, 6))
1390['a', 'b', 'c', None, None, None]
1391
1392>>> list(ncycles('abc', 3))
1393['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1394
1395>>> dotproduct([1,2,3], [4,5,6])
139632
1397
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001398>>> list(grouper(3, 'abcdefg', 'x'))
1399[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1400
1401>>> list(roundrobin('abc', 'd', 'ef'))
1402['a', 'd', 'e', 'b', 'f', 'c']
1403
Raymond Hettinger5cd012b2009-01-26 02:17:52 +00001404>>> list(powerset([1,2,3]))
1405[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001406
Raymond Hettingere8b4b602008-03-11 00:19:07 +00001407>>> list(compress('abcdef', [1,0,1,0,1,1]))
1408['a', 'c', 'e', 'f']
1409
Raymond Hettinger33691672008-07-19 00:43:00 +00001410>>> list(combinations_with_replacement('abc', 2))
1411[('a', 'a'), ('a', 'b'), ('a', 'c'), ('b', 'b'), ('b', 'c'), ('c', 'c')]
1412
Raymond Hettinger825758c2009-01-08 05:20:19 +00001413>>> list(combinations_with_replacement('01', 3))
1414[('0', '0', '0'), ('0', '0', '1'), ('0', '1', '1'), ('1', '1', '1')]
1415
1416>>> def combinations_with_replacement2(iterable, r):
1417... 'Alternate version that filters from product()'
1418... pool = tuple(iterable)
1419... n = len(pool)
1420... for indices in product(range(n), repeat=r):
1421... if sorted(indices) == list(indices):
1422... yield tuple(pool[i] for i in indices)
1423
1424>>> list(combinations_with_replacement('abc', 2)) == list(combinations_with_replacement2('abc', 2))
1425True
1426
1427>>> list(combinations_with_replacement('01', 3)) == list(combinations_with_replacement2('01', 3))
1428True
1429
1430>>> list(combinations_with_replacement('2310', 6)) == list(combinations_with_replacement2('2310', 6))
1431True
1432
Raymond Hettingerd8461652009-01-02 21:20:38 +00001433>>> list(unique_everseen('AAAABBBCCDAABBB'))
1434['A', 'B', 'C', 'D']
1435
1436>>> list(unique_everseen('ABBCcAD', str.lower))
1437['A', 'B', 'C', 'D']
1438
1439>>> list(unique_justseen('AAAABBBCCDAABBB'))
1440['A', 'B', 'C', 'D', 'A', 'B']
1441
1442>>> list(unique_justseen('ABBCcAD', str.lower))
1443['A', 'B', 'C', 'A', 'D']
1444
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001445"""
1446
1447__test__ = {'libreftest' : libreftest}
1448
1449def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001450 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Georg Brandlb84c1372007-01-21 10:28:43 +00001451 RegressionTests, LengthTransparency,
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001452 SubclassWithKwargsTest, TestExamples)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001453 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001454
1455 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001456 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001457 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001458 counts = [None] * 5
1459 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001460 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001461 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001462 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001463 print counts
1464
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001465 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001466 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001467
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001468if __name__ == "__main__":
1469 test_main(verbose=True)