blob: de71fda5de24654bae75a46457100b28f741fda3 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001import unittest
2from test import test_support
3from itertools import *
Raymond Hettingera9f60922004-10-17 16:40:14 +00004from weakref import proxy
Raymond Hettinger2012f172003-02-07 05:32:58 +00005import sys
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00006import operator
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00007import random
Raymond Hettinger673827c2009-11-30 11:15:28 +00008import copy
9import pickle
Ezio Melotti3efafd72010-08-02 18:40:55 +000010from functools import reduce
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +000011maxsize = test_support.MAX_Py_ssize_t
12minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000013
14def onearg(x):
15 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000016 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000017
18def errfunc(*args):
19 'Test function that raises an error'
20 raise ValueError
21
22def gen3():
23 'Non-restartable source sequence'
24 for i in (0, 1, 2):
25 yield i
26
27def isEven(x):
28 'Test predicate'
29 return x%2==0
30
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000031def isOdd(x):
32 'Test predicate'
33 return x%2==1
34
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000035class StopNow:
36 'Class emulating an empty iterable.'
37 def __iter__(self):
38 return self
39 def next(self):
40 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000041
Raymond Hettinger02420702003-06-29 20:36:23 +000042def take(n, seq):
43 'Convenience function for partially consuming a long of infinite iterable'
44 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000045
Raymond Hettingerd553d852008-03-04 04:17:08 +000046def prod(iterable):
47 return reduce(operator.mul, iterable, 1)
48
Raymond Hettinger93e804d2008-02-26 23:40:50 +000049def fact(n):
50 'Factorial'
Raymond Hettingerd553d852008-03-04 04:17:08 +000051 return prod(range(1, n+1))
52
Raymond Hettinger96ef8112003-02-01 00:10:11 +000053class TestBasicOps(unittest.TestCase):
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000054 def test_chain(self):
Raymond Hettingerad47fa12008-03-06 20:52:01 +000055
56 def chain2(*iterables):
57 'Pure python version in the docs'
58 for it in iterables:
59 for element in it:
60 yield element
61
62 for c in (chain, chain2):
63 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
64 self.assertEqual(list(c('abc')), list('abc'))
65 self.assertEqual(list(c('')), [])
66 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
67 self.assertRaises(TypeError, list,c(2, 3))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000068
Raymond Hettingerb4cbc982008-02-28 22:46:41 +000069 def test_chain_from_iterable(self):
70 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
71 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
72 self.assertEqual(list(chain.from_iterable([''])), [])
73 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
74 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
75
Raymond Hettinger93e804d2008-02-26 23:40:50 +000076 def test_combinations(self):
Raymond Hettinger825758c2009-01-08 05:20:19 +000077 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Raymond Hettinger93e804d2008-02-26 23:40:50 +000078 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Raymond Hettingerd553d852008-03-04 04:17:08 +000079 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Raymond Hettinger93e804d2008-02-26 23:40:50 +000080 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Raymond Hettinger825758c2009-01-08 05:20:19 +000081 self.assertEqual(list(combinations('abc', 32)), []) # r > n
Raymond Hettinger93e804d2008-02-26 23:40:50 +000082 self.assertEqual(list(combinations(range(4), 3)),
83 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
Raymond Hettingerd553d852008-03-04 04:17:08 +000084
85 def combinations1(iterable, r):
86 'Pure python version shown in the docs'
87 pool = tuple(iterable)
88 n = len(pool)
Raymond Hettinger825758c2009-01-08 05:20:19 +000089 if r > n:
90 return
Raymond Hettingerd553d852008-03-04 04:17:08 +000091 indices = range(r)
92 yield tuple(pool[i] for i in indices)
93 while 1:
94 for i in reversed(range(r)):
95 if indices[i] != i + n - r:
96 break
97 else:
98 return
99 indices[i] += 1
100 for j in range(i+1, r):
101 indices[j] = indices[j-1] + 1
102 yield tuple(pool[i] for i in indices)
103
104 def combinations2(iterable, r):
105 'Pure python version shown in the docs'
106 pool = tuple(iterable)
107 n = len(pool)
108 for indices in permutations(range(n), r):
109 if sorted(indices) == list(indices):
110 yield tuple(pool[i] for i in indices)
111
112 for n in range(7):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000113 values = [5*x-12 for x in range(n)]
Raymond Hettinger825758c2009-01-08 05:20:19 +0000114 for r in range(n+2):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000115 result = list(combinations(values, r))
Ezio Melotti3efafd72010-08-02 18:40:55 +0000116 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 +0000117 self.assertEqual(len(result), len(set(result))) # no repeats
118 self.assertEqual(result, sorted(result)) # lexicographic order
119 for c in result:
120 self.assertEqual(len(c), r) # r-length combinations
121 self.assertEqual(len(set(c)), r) # no duplicate elements
122 self.assertEqual(list(c), sorted(c)) # keep original ordering
123 self.assert_(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000124 self.assertEqual(list(c),
125 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Raymond Hettingerd553d852008-03-04 04:17:08 +0000126 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger825758c2009-01-08 05:20:19 +0000127 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000128
129 # Test implementation detail: tuple re-use
130 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
131 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
132
133 def test_permutations(self):
134 self.assertRaises(TypeError, permutations) # too few arguments
135 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000136 self.assertRaises(TypeError, permutations, None) # pool is not iterable
137 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger825758c2009-01-08 05:20:19 +0000138 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000139 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Raymond Hettingerd553d852008-03-04 04:17:08 +0000140 self.assertEqual(list(permutations(range(3), 2)),
141 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
142
143 def permutations1(iterable, r=None):
144 'Pure python version shown in the docs'
145 pool = tuple(iterable)
146 n = len(pool)
147 r = n if r is None else r
Raymond Hettinger825758c2009-01-08 05:20:19 +0000148 if r > n:
149 return
Raymond Hettingerd553d852008-03-04 04:17:08 +0000150 indices = range(n)
Raymond Hettingere70bb8d2008-03-23 00:55:46 +0000151 cycles = range(n, n-r, -1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000152 yield tuple(pool[i] for i in indices[:r])
153 while n:
154 for i in reversed(range(r)):
155 cycles[i] -= 1
156 if cycles[i] == 0:
157 indices[i:] = indices[i+1:] + indices[i:i+1]
158 cycles[i] = n - i
159 else:
160 j = cycles[i]
161 indices[i], indices[-j] = indices[-j], indices[i]
162 yield tuple(pool[i] for i in indices[:r])
163 break
164 else:
165 return
166
167 def permutations2(iterable, r=None):
168 'Pure python version shown in the docs'
169 pool = tuple(iterable)
170 n = len(pool)
171 r = n if r is None else r
172 for indices in product(range(n), repeat=r):
173 if len(set(indices)) == r:
174 yield tuple(pool[i] for i in indices)
175
176 for n in range(7):
177 values = [5*x-12 for x in range(n)]
Raymond Hettinger825758c2009-01-08 05:20:19 +0000178 for r in range(n+2):
Raymond Hettingerd553d852008-03-04 04:17:08 +0000179 result = list(permutations(values, r))
Ezio Melotti3efafd72010-08-02 18:40:55 +0000180 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 +0000181 self.assertEqual(len(result), len(set(result))) # no repeats
182 self.assertEqual(result, sorted(result)) # lexicographic order
183 for p in result:
184 self.assertEqual(len(p), r) # r-length permutations
185 self.assertEqual(len(set(p)), r) # no duplicate elements
186 self.assert_(all(e in values for e in p)) # elements taken from input iterable
187 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger825758c2009-01-08 05:20:19 +0000188 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000189 if r == n:
190 self.assertEqual(result, list(permutations(values, None))) # test r as None
191 self.assertEqual(result, list(permutations(values))) # test default r
192
193 # Test implementation detail: tuple re-use
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000194 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000195 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000196
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000197 def test_count(self):
198 self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
199 self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000200 self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000201 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
202 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000203 self.assertRaises(TypeError, count, 2, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000204 self.assertRaises(TypeError, count, 'a')
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000205 self.assertEqual(list(islice(count(maxsize-5), 10)), range(maxsize-5, maxsize+5))
206 self.assertEqual(list(islice(count(-maxsize-5), 10)), range(-maxsize-5, -maxsize+5))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000207 c = count(3)
208 self.assertEqual(repr(c), 'count(3)')
209 c.next()
210 self.assertEqual(repr(c), 'count(4)')
Jack Diederich36234e82006-09-21 17:50:26 +0000211 c = count(-9)
212 self.assertEqual(repr(c), 'count(-9)')
213 c.next()
214 self.assertEqual(c.next(), -8)
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000215 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 +0000216 # Test repr (ignoring the L in longs)
217 r1 = repr(count(i)).replace('L', '')
218 r2 = 'count(%r)'.__mod__(i).replace('L', '')
219 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000220
Raymond Hettinger673827c2009-11-30 11:15:28 +0000221 # check copy, deepcopy, pickle
222 for value in -3, 3, sys.maxint-5, sys.maxint+5:
223 c = count(value)
224 self.assertEqual(next(copy.copy(c)), value)
225 self.assertEqual(next(copy.deepcopy(c)), value)
226 self.assertEqual(next(pickle.loads(pickle.dumps(c))), value)
227
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000228 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000229 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000230 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000231 self.assertRaises(TypeError, cycle)
232 self.assertRaises(TypeError, cycle, 5)
233 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000234
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000235 def test_groupby(self):
236 # Check whether it accepts arguments correctly
237 self.assertEqual([], list(groupby([])))
238 self.assertEqual([], list(groupby([], key=id)))
239 self.assertRaises(TypeError, list, groupby('abc', []))
240 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000241 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000242
243 # Check normal input
244 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
245 (2,15,22), (3,16,23), (3,17,23)]
246 dup = []
247 for k, g in groupby(s, lambda r:r[0]):
248 for elem in g:
249 self.assertEqual(k, elem[0])
250 dup.append(elem)
251 self.assertEqual(s, dup)
252
253 # Check nested case
254 dup = []
255 for k, g in groupby(s, lambda r:r[0]):
256 for ik, ig in groupby(g, lambda r:r[2]):
257 for elem in ig:
258 self.assertEqual(k, elem[0])
259 self.assertEqual(ik, elem[2])
260 dup.append(elem)
261 self.assertEqual(s, dup)
262
263 # Check case where inner iterator is not used
264 keys = [k for k, g in groupby(s, lambda r:r[0])]
265 expectedkeys = set([r[0] for r in s])
266 self.assertEqual(set(keys), expectedkeys)
267 self.assertEqual(len(keys), len(expectedkeys))
268
269 # Exercise pipes and filters style
270 s = 'abracadabra'
271 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000272 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000273 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
274 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000275 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000276 self.assertEqual(r, ['a', 'b', 'r'])
277 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000278 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000279 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
280 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000281 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000282 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
283
284 # iter.next failure
285 class ExpectedError(Exception):
286 pass
287 def delayed_raise(n=0):
288 for i in range(n):
289 yield 'yo'
290 raise ExpectedError
291 def gulp(iterable, keyp=None, func=list):
292 return [func(g) for k, g in groupby(iterable, keyp)]
293
294 # iter.next failure on outer object
295 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
296 # iter.next failure on inner object
297 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
298
299 # __cmp__ failure
300 class DummyCmp:
301 def __cmp__(self, dst):
302 raise ExpectedError
303 s = [DummyCmp(), DummyCmp(), None]
304
305 # __cmp__ failure on outer object
306 self.assertRaises(ExpectedError, gulp, s, func=id)
307 # __cmp__ failure on inner object
308 self.assertRaises(ExpectedError, gulp, s)
309
310 # keyfunc failure
311 def keyfunc(obj):
312 if keyfunc.skip > 0:
313 keyfunc.skip -= 1
314 return obj
315 else:
316 raise ExpectedError
317
318 # keyfunc failure on outer object
319 keyfunc.skip = 0
320 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
321 keyfunc.skip = 1
322 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
323
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000324 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000325 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000326 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000327 self.assertEqual(list(ifilter(bool, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000328 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000329 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000330 self.assertRaises(TypeError, ifilter, lambda x:x)
331 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000332 self.assertRaises(TypeError, ifilter, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000333 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000334
335 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000336 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
337 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000338 self.assertEqual(list(ifilterfalse(bool, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000339 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000340 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000341 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
342 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000343 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000344 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000345
346 def test_izip(self):
347 ans = [(x,y) for x, y in izip('abc',count())]
348 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000349 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
350 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
Raymond Hettinger02420702003-06-29 20:36:23 +0000351 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000352 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000353 self.assertEqual(list(izip()), zip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000354 self.assertRaises(TypeError, izip, 3)
355 self.assertRaises(TypeError, izip, range(3), 3)
356 # Check tuple re-use (implementation detail)
357 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
358 zip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000359 self.assertEqual([pair for pair in izip('abc', 'def')],
360 zip('abc', 'def'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000361 ids = map(id, izip('abc', 'def'))
362 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000363 ids = map(id, list(izip('abc', 'def')))
364 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000365
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000366 def test_iziplongest(self):
367 for args in [
368 ['abc', range(6)],
369 [range(6), 'abc'],
370 [range(1000), range(2000,2100), range(3000,3050)],
371 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
372 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
373 ]:
Ezio Melotti3efafd72010-08-02 18:40:55 +0000374 # target = map(None, *args) <- this raises a py3k warning
375 # this is the replacement:
376 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
377 for i in range(max(map(len, args)))]
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000378 self.assertEqual(list(izip_longest(*args)), target)
379 self.assertEqual(list(izip_longest(*args, **{})), target)
380 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
381 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Tim Petersea5962f2007-03-12 18:07:52 +0000382
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000383 self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
384
385 self.assertEqual(list(izip_longest()), zip())
386 self.assertEqual(list(izip_longest([])), zip([]))
387 self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
Tim Petersea5962f2007-03-12 18:07:52 +0000388
Ezio Melotti3efafd72010-08-02 18:40:55 +0000389 self.assertEqual(list(izip_longest('abc', 'defg', **{})),
390 zip(list('abc') + [None], 'defg')) # empty keyword dict
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000391 self.assertRaises(TypeError, izip_longest, 3)
392 self.assertRaises(TypeError, izip_longest, range(3), 3)
393
394 for stmt in [
395 "izip_longest('abc', fv=1)",
Tim Petersea5962f2007-03-12 18:07:52 +0000396 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000397 ]:
398 try:
399 eval(stmt, globals(), locals())
400 except TypeError:
401 pass
402 else:
403 self.fail('Did not raise Type in: ' + stmt)
Tim Petersea5962f2007-03-12 18:07:52 +0000404
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000405 # Check tuple re-use (implementation detail)
406 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
407 zip('abc', 'def'))
408 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
409 zip('abc', 'def'))
410 ids = map(id, izip_longest('abc', 'def'))
411 self.assertEqual(min(ids), max(ids))
412 ids = map(id, list(izip_longest('abc', 'def')))
413 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
414
Raymond Hettinger80d49b32009-11-01 08:53:21 +0000415 def test_bug_7244(self):
416
417 class Repeater(object):
418 # this class is similar to itertools.repeat
419 def __init__(self, o, t, e):
420 self.o = o
421 self.t = int(t)
422 self.e = e
423 def __iter__(self): # its iterator is itself
424 return self
425 def next(self):
426 if self.t > 0:
427 self.t -= 1
428 return self.o
429 else:
430 raise self.e
431
432 # Formerly this code in would fail in debug mode
433 # with Undetected Error and Stop Iteration
434 r1 = Repeater(1, 3, StopIteration)
435 r2 = Repeater(2, 4, StopIteration)
436 def run(r1, r2):
437 result = []
438 for i, j in izip_longest(r1, r2, fillvalue=0):
Raymond Hettinger4da5faa2009-11-01 18:43:31 +0000439 with test_support.captured_output('stdout'):
440 print (i, j)
Raymond Hettinger80d49b32009-11-01 08:53:21 +0000441 result.append((i, j))
442 return result
443 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
444
445 # Formerly, the RuntimeError would be lost
446 # and StopIteration would stop as expected
447 r1 = Repeater(1, 3, RuntimeError)
448 r2 = Repeater(2, 4, StopIteration)
Raymond Hettinger4da5faa2009-11-01 18:43:31 +0000449 it = izip_longest(r1, r2, fillvalue=0)
450 self.assertEqual(next(it), (1, 2))
451 self.assertEqual(next(it), (1, 2))
452 self.assertEqual(next(it), (1, 2))
453 self.assertRaises(RuntimeError, next, it)
Raymond Hettinger80d49b32009-11-01 08:53:21 +0000454
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000455 def test_product(self):
456 for args, result in [
Raymond Hettingerd553d852008-03-04 04:17:08 +0000457 ([], [()]), # zero iterables
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000458 (['ab'], [('a',), ('b',)]), # one iterable
459 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
460 ([range(0), range(2), range(3)], []), # first iterable with zero length
461 ([range(2), range(0), range(3)], []), # middle iterable with zero length
462 ([range(2), range(3), range(0)], []), # last iterable with zero length
463 ]:
464 self.assertEqual(list(product(*args)), result)
Raymond Hettinger08ff6822008-02-29 02:21:48 +0000465 for r in range(4):
466 self.assertEqual(list(product(*(args*r))),
467 list(product(*args, **dict(repeat=r))))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000468 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
469 self.assertRaises(TypeError, product, range(6), None)
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000470
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000471 def product1(*args, **kwds):
472 pools = map(tuple, args) * kwds.get('repeat', 1)
473 n = len(pools)
474 if n == 0:
475 yield ()
476 return
477 if any(len(pool) == 0 for pool in pools):
478 return
479 indices = [0] * n
480 yield tuple(pool[i] for pool, i in zip(pools, indices))
481 while 1:
482 for i in reversed(range(n)): # right to left
483 if indices[i] == len(pools[i]) - 1:
484 continue
485 indices[i] += 1
486 for j in range(i+1, n):
487 indices[j] = 0
488 yield tuple(pool[i] for pool, i in zip(pools, indices))
489 break
490 else:
491 return
492
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000493 def product2(*args, **kwds):
494 'Pure python version used in docs'
495 pools = map(tuple, args) * kwds.get('repeat', 1)
496 result = [[]]
497 for pool in pools:
498 result = [x+[y] for x in result for y in pool]
499 for prod in result:
500 yield tuple(prod)
501
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000502 argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
503 set('abcdefg'), range(11), tuple(range(13))]
504 for i in range(100):
505 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Raymond Hettingerd553d852008-03-04 04:17:08 +0000506 expected_len = prod(map(len, args))
507 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000508 self.assertEqual(list(product(*args)), list(product1(*args)))
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000509 self.assertEqual(list(product(*args)), list(product2(*args)))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000510 args = map(iter, args)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000511 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000512
Raymond Hettinger73d79632008-02-23 02:20:41 +0000513 # Test implementation detail: tuple re-use
514 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
515 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000516
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000517 def test_repeat(self):
518 self.assertEqual(zip(xrange(3),repeat('a')),
519 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000520 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000521 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000522 self.assertEqual(list(repeat('a', 0)), [])
523 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000524 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000525 self.assertRaises(TypeError, repeat, None, 3, 4)
526 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000527 r = repeat(1+0j)
528 self.assertEqual(repr(r), 'repeat((1+0j))')
529 r = repeat(1+0j, 5)
530 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
531 list(r)
532 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000533
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000534 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000535 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
536 [0**1, 1**2, 2**3])
537 self.assertEqual(list(imap(None, 'abc', range(5))),
538 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000539 self.assertEqual(list(imap(None, 'abc', count())),
540 [('a',0),('b',1),('c',2)])
541 self.assertEqual(take(2,imap(None, 'abc', count())),
542 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000543 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000544 self.assertRaises(TypeError, imap)
545 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000546 self.assertRaises(TypeError, imap(10, range(5)).next)
547 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
548 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000549
550 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000551 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
552 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000553 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
554 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000555 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger47317092008-01-17 03:02:14 +0000556 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
557 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000558 self.assertRaises(TypeError, starmap)
559 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
560 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
561 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
562 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000563
564 def test_islice(self):
565 for args in [ # islice(args) should agree with range(args)
566 (10, 20, 3),
567 (10, 3, 20),
568 (10, 20),
569 (10, 3),
570 (20,)
571 ]:
572 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
573
574 for args, tgtargs in [ # Stop when seqn is exhausted
575 ((10, 110, 3), ((10, 100, 3))),
576 ((10, 110), ((10, 100))),
577 ((110,), (100,))
578 ]:
579 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
580
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000581 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000582 self.assertEqual(list(islice(xrange(10), None)), range(10))
Raymond Hettingerb2594052004-12-05 09:25:51 +0000583 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
584 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000585 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
586 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
587
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000588 # Test number of items consumed SF #1171417
589 it = iter(range(10))
590 self.assertEqual(list(islice(it, 3)), range(3))
591 self.assertEqual(list(it), range(3, 10))
592
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000593 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000594 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000595 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
596 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
597 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
598 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
599 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000600 self.assertRaises(ValueError, islice, xrange(10), 'a')
601 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
602 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
603 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
604 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +0000605 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000606
607 def test_takewhile(self):
608 data = [1, 3, 5, 20, 2, 4, 6, 8]
609 underten = lambda x: x<10
610 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000611 self.assertEqual(list(takewhile(underten, [])), [])
612 self.assertRaises(TypeError, takewhile)
613 self.assertRaises(TypeError, takewhile, operator.pow)
614 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
615 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
616 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000617 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
618 self.assertEqual(list(t), [1, 1, 1])
619 self.assertRaises(StopIteration, t.next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000620
621 def test_dropwhile(self):
622 data = [1, 3, 5, 20, 2, 4, 6, 8]
623 underten = lambda x: x<10
624 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000625 self.assertEqual(list(dropwhile(underten, [])), [])
626 self.assertRaises(TypeError, dropwhile)
627 self.assertRaises(TypeError, dropwhile, operator.pow)
628 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
629 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
630 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000631
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000632 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000633 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000634 def irange(n):
635 for i in xrange(n):
636 yield i
637
638 a, b = tee([]) # test empty iterator
639 self.assertEqual(list(a), [])
640 self.assertEqual(list(b), [])
641
642 a, b = tee(irange(n)) # test 100% interleaved
643 self.assertEqual(zip(a,b), zip(range(n),range(n)))
644
645 a, b = tee(irange(n)) # test 0% interleaved
646 self.assertEqual(list(a), range(n))
647 self.assertEqual(list(b), range(n))
648
649 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000650 for i in xrange(100):
651 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000652 del a
653 self.assertEqual(list(b), range(n))
654
655 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000656 for i in xrange(100):
657 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000658 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000659 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000660
661 for j in xrange(5): # test randomly interleaved
662 order = [0]*n + [1]*n
663 random.shuffle(order)
664 lists = ([], [])
665 its = tee(irange(n))
666 for i in order:
667 value = its[i].next()
668 lists[i].append(value)
669 self.assertEqual(lists[0], range(n))
670 self.assertEqual(lists[1], range(n))
671
Raymond Hettingerad983e72003-11-12 14:32:26 +0000672 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000673 self.assertRaises(TypeError, tee)
674 self.assertRaises(TypeError, tee, 3)
675 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000676 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000677
Raymond Hettingerad983e72003-11-12 14:32:26 +0000678 # tee object should be instantiable
679 a, b = tee('abc')
680 c = type(a)('def')
681 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000682
Raymond Hettingerad983e72003-11-12 14:32:26 +0000683 # test long-lagged and multi-way split
684 a, b, c = tee(xrange(2000), 3)
685 for i in xrange(100):
686 self.assertEqual(a.next(), i)
687 self.assertEqual(list(b), range(2000))
688 self.assertEqual([c.next(), c.next()], range(2))
689 self.assertEqual(list(a), range(100,2000))
690 self.assertEqual(list(c), range(2,2000))
691
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000692 # test values of n
693 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Neal Norwitz69e88972006-09-02 02:58:13 +0000694 self.assertRaises(ValueError, tee, [], -1)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000695 for n in xrange(5):
696 result = tee('abc', n)
697 self.assertEqual(type(result), tuple)
698 self.assertEqual(len(result), n)
699 self.assertEqual(map(list, result), [list('abc')]*n)
700
Raymond Hettingerad983e72003-11-12 14:32:26 +0000701 # tee pass-through to copyable iterator
702 a, b = tee('abc')
703 c, d = tee(a)
704 self.assert_(a is c)
705
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000706 # test tee_new
707 t1, t2 = tee('abc')
708 tnew = type(t1)
709 self.assertRaises(TypeError, tnew)
710 self.assertRaises(TypeError, tnew, 10)
711 t3 = tnew(t1)
712 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000713
Raymond Hettingera9f60922004-10-17 16:40:14 +0000714 # test that tee objects are weak referencable
715 a, b = tee(xrange(10))
716 p = proxy(a)
717 self.assertEqual(getattr(p, '__class__'), type(b))
718 del a
719 self.assertRaises(ReferenceError, getattr, p, '__class__')
720
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000721 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000722 self.assertRaises(StopIteration, izip().next)
723
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000724 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000725 self.assertRaises(StopIteration, f([]).next)
726 self.assertRaises(StopIteration, f(StopNow()).next)
727
728 self.assertRaises(StopIteration, islice([], None).next)
729 self.assertRaises(StopIteration, islice(StopNow(), None).next)
730
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000731 p, q = tee([])
732 self.assertRaises(StopIteration, p.next)
733 self.assertRaises(StopIteration, q.next)
734 p, q = tee(StopNow())
735 self.assertRaises(StopIteration, p.next)
736 self.assertRaises(StopIteration, q.next)
737
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000738 self.assertRaises(StopIteration, repeat(None, 0).next)
739
740 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
741 self.assertRaises(StopIteration, f(lambda x:x, []).next)
742 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
743
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000744class TestExamples(unittest.TestCase):
745
746 def test_chain(self):
747 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
748
749 def test_chain_from_iterable(self):
750 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
751
752 def test_combinations(self):
753 self.assertEqual(list(combinations('ABCD', 2)),
754 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
755 self.assertEqual(list(combinations(range(4), 3)),
756 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
757
758 def test_count(self):
759 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
760
761 def test_cycle(self):
762 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
763
764 def test_dropwhile(self):
765 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
766
767 def test_groupby(self):
768 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
769 list('ABCDAB'))
770 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
771 [list('AAAA'), list('BBB'), list('CC'), list('D')])
772
773 def test_ifilter(self):
774 self.assertEqual(list(ifilter(lambda x: x%2, range(10))), [1,3,5,7,9])
775
776 def test_ifilterfalse(self):
777 self.assertEqual(list(ifilterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
778
779 def test_imap(self):
780 self.assertEqual(list(imap(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
781
782 def test_islice(self):
783 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
784 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
785 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
786 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
787
788 def test_izip(self):
789 self.assertEqual(list(izip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
790
791 def test_izip_longest(self):
792 self.assertEqual(list(izip_longest('ABCD', 'xy', fillvalue='-')),
793 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
794
795 def test_permutations(self):
796 self.assertEqual(list(permutations('ABCD', 2)),
797 map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split()))
798 self.assertEqual(list(permutations(range(3))),
799 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
800
801 def test_product(self):
802 self.assertEqual(list(product('ABCD', 'xy')),
803 map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split()))
804 self.assertEqual(list(product(range(2), repeat=3)),
805 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
806 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
807
808 def test_repeat(self):
809 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
810
811 def test_stapmap(self):
812 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
813 [32, 9, 1000])
814
815 def test_takewhile(self):
816 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
817
818
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000819class TestGC(unittest.TestCase):
820
821 def makecycle(self, iterator, container):
822 container.append(iterator)
823 iterator.next()
824 del container, iterator
825
826 def test_chain(self):
827 a = []
828 self.makecycle(chain(a), a)
829
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000830 def test_chain_from_iterable(self):
831 a = []
832 self.makecycle(chain.from_iterable([a]), a)
833
834 def test_combinations(self):
835 a = []
836 self.makecycle(combinations([1,2,a,3], 3), a)
837
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000838 def test_cycle(self):
839 a = []
840 self.makecycle(cycle([a]*2), a)
841
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000842 def test_dropwhile(self):
843 a = []
844 self.makecycle(dropwhile(bool, [0, a, a]), a)
845
846 def test_groupby(self):
847 a = []
848 self.makecycle(groupby([a]*2, lambda x:x), a)
849
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000850 def test_issue2246(self):
851 # Issue 2246 -- the _grouper iterator was not included in GC
852 n = 10
853 keyfunc = lambda x: x
854 for i, j in groupby(xrange(n), key=keyfunc):
855 keyfunc.__dict__.setdefault('x',[]).append(j)
856
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000857 def test_ifilter(self):
858 a = []
859 self.makecycle(ifilter(lambda x:True, [a]*2), a)
860
861 def test_ifilterfalse(self):
862 a = []
863 self.makecycle(ifilterfalse(lambda x:False, a), a)
864
865 def test_izip(self):
866 a = []
867 self.makecycle(izip([a]*2, [a]*3), a)
868
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000869 def test_izip_longest(self):
870 a = []
871 self.makecycle(izip_longest([a]*2, [a]*3), a)
872 b = [a, None]
873 self.makecycle(izip_longest([a]*2, [a]*3, fillvalue=b), a)
874
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000875 def test_imap(self):
876 a = []
877 self.makecycle(imap(lambda x:x, [a]*2), a)
878
879 def test_islice(self):
880 a = []
881 self.makecycle(islice([a]*2, None), a)
882
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000883 def test_permutations(self):
884 a = []
885 self.makecycle(permutations([1,2,a,3], 3), a)
886
887 def test_product(self):
888 a = []
889 self.makecycle(product([1,2,a,3], repeat=3), a)
890
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000891 def test_repeat(self):
892 a = []
893 self.makecycle(repeat(a), a)
894
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000895 def test_starmap(self):
896 a = []
897 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
898
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000899 def test_takewhile(self):
900 a = []
901 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
902
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000903def R(seqn):
904 'Regular generator'
905 for i in seqn:
906 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000907
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000908class G:
909 'Sequence using __getitem__'
910 def __init__(self, seqn):
911 self.seqn = seqn
912 def __getitem__(self, i):
913 return self.seqn[i]
914
915class I:
916 'Sequence using iterator protocol'
917 def __init__(self, seqn):
918 self.seqn = seqn
919 self.i = 0
920 def __iter__(self):
921 return self
922 def next(self):
923 if self.i >= len(self.seqn): raise StopIteration
924 v = self.seqn[self.i]
925 self.i += 1
926 return v
927
928class Ig:
929 'Sequence using iterator protocol defined with a generator'
930 def __init__(self, seqn):
931 self.seqn = seqn
932 self.i = 0
933 def __iter__(self):
934 for val in self.seqn:
935 yield val
936
937class X:
938 'Missing __getitem__ and __iter__'
939 def __init__(self, seqn):
940 self.seqn = seqn
941 self.i = 0
942 def next(self):
943 if self.i >= len(self.seqn): raise StopIteration
944 v = self.seqn[self.i]
945 self.i += 1
946 return v
947
948class N:
949 'Iterator missing next()'
950 def __init__(self, seqn):
951 self.seqn = seqn
952 self.i = 0
953 def __iter__(self):
954 return self
955
956class E:
957 'Test propagation of exceptions'
958 def __init__(self, seqn):
959 self.seqn = seqn
960 self.i = 0
961 def __iter__(self):
962 return self
963 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +0000964 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000965
966class S:
967 'Test immediate stop'
968 def __init__(self, seqn):
969 pass
970 def __iter__(self):
971 return self
972 def next(self):
973 raise StopIteration
974
975def L(seqn):
976 'Test multiple tiers of iterators'
977 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
978
979
980class TestVariousIteratorArgs(unittest.TestCase):
981
982 def test_chain(self):
983 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
984 for g in (G, I, Ig, S, L, R):
985 self.assertEqual(list(chain(g(s))), list(g(s)))
986 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Raymond Hettinger05bf6332008-02-28 22:30:42 +0000987 self.assertRaises(TypeError, list, chain(X(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000988 self.assertRaises(TypeError, list, chain(N(s)))
989 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
990
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000991 def test_product(self):
992 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
993 self.assertRaises(TypeError, product, X(s))
994 self.assertRaises(TypeError, product, N(s))
995 self.assertRaises(ZeroDivisionError, product, E(s))
996
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000997 def test_cycle(self):
998 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
999 for g in (G, I, Ig, S, L, R):
1000 tgtlen = len(s) * 3
1001 expected = list(g(s))*3
1002 actual = list(islice(cycle(g(s)), tgtlen))
1003 self.assertEqual(actual, expected)
1004 self.assertRaises(TypeError, cycle, X(s))
1005 self.assertRaises(TypeError, list, cycle(N(s)))
1006 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1007
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001008 def test_groupby(self):
1009 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1010 for g in (G, I, Ig, S, L, R):
1011 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1012 self.assertRaises(TypeError, groupby, X(s))
1013 self.assertRaises(TypeError, list, groupby(N(s)))
1014 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1015
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001016 def test_ifilter(self):
1017 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1018 for g in (G, I, Ig, S, L, R):
1019 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
1020 self.assertRaises(TypeError, ifilter, isEven, X(s))
1021 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
1022 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
1023
1024 def test_ifilterfalse(self):
1025 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1026 for g in (G, I, Ig, S, L, R):
1027 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
1028 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
1029 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
1030 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
1031
1032 def test_izip(self):
1033 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1034 for g in (G, I, Ig, S, L, R):
1035 self.assertEqual(list(izip(g(s))), zip(g(s)))
1036 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
1037 self.assertRaises(TypeError, izip, X(s))
1038 self.assertRaises(TypeError, list, izip(N(s)))
1039 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
1040
Raymond Hettingerd36862c2007-02-21 05:20:38 +00001041 def test_iziplongest(self):
1042 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1043 for g in (G, I, Ig, S, L, R):
1044 self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
1045 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
1046 self.assertRaises(TypeError, izip_longest, X(s))
1047 self.assertRaises(TypeError, list, izip_longest(N(s)))
1048 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
1049
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001050 def test_imap(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 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
1054 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
1055 self.assertRaises(TypeError, imap, onearg, X(s))
1056 self.assertRaises(TypeError, list, imap(onearg, N(s)))
1057 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
1058
1059 def test_islice(self):
1060 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1061 for g in (G, I, Ig, S, L, R):
1062 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1063 self.assertRaises(TypeError, islice, X(s), 10)
1064 self.assertRaises(TypeError, list, islice(N(s), 10))
1065 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1066
1067 def test_starmap(self):
1068 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1069 for g in (G, I, Ig, S, L, R):
1070 ss = zip(s, s)
1071 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
1072 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
1073 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
1074 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1075
1076 def test_takewhile(self):
1077 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1078 for g in (G, I, Ig, S, L, R):
1079 tgt = []
1080 for elem in g(s):
1081 if not isEven(elem): break
1082 tgt.append(elem)
1083 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1084 self.assertRaises(TypeError, takewhile, isEven, X(s))
1085 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
1086 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1087
1088 def test_dropwhile(self):
1089 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1090 for g in (G, I, Ig, S, L, R):
1091 tgt = []
1092 for elem in g(s):
1093 if not tgt and isOdd(elem): continue
1094 tgt.append(elem)
1095 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1096 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
1097 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
1098 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1099
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001100 def test_tee(self):
1101 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1102 for g in (G, I, Ig, S, L, R):
1103 it1, it2 = tee(g(s))
1104 self.assertEqual(list(it1), list(g(s)))
1105 self.assertEqual(list(it2), list(g(s)))
1106 self.assertRaises(TypeError, tee, X(s))
1107 self.assertRaises(TypeError, list, tee(N(s))[0])
1108 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1109
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001110class LengthTransparency(unittest.TestCase):
1111
1112 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001113 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001114 self.assertEqual(len(repeat(None, 50)), 50)
1115 self.assertRaises(TypeError, len, repeat(None))
1116
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001117class RegressionTests(unittest.TestCase):
1118
1119 def test_sf_793826(self):
1120 # Fix Armin Rigo's successful efforts to wreak havoc
1121
1122 def mutatingtuple(tuple1, f, tuple2):
1123 # this builds a tuple t which is a copy of tuple1,
1124 # then calls f(t), then mutates t to be equal to tuple2
1125 # (needs len(tuple1) == len(tuple2)).
1126 def g(value, first=[1]):
1127 if first:
1128 del first[:]
1129 f(z.next())
1130 return value
1131 items = list(tuple2)
1132 items[1:1] = list(tuple1)
1133 gen = imap(g, items)
1134 z = izip(*[gen]*len(tuple1))
1135 z.next()
1136
1137 def f(t):
1138 global T
1139 T = t
1140 first[:] = list(T)
1141
1142 first = []
1143 mutatingtuple((1,2,3), f, (4,5,6))
1144 second = list(T)
1145 self.assertEqual(first, second)
1146
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001147
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001148 def test_sf_950057(self):
1149 # Make sure that chain() and cycle() catch exceptions immediately
1150 # rather than when shifting between input sources
1151
1152 def gen1():
1153 hist.append(0)
1154 yield 1
1155 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001156 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001157 hist.append(2)
1158
1159 def gen2(x):
1160 hist.append(3)
1161 yield 2
1162 hist.append(4)
1163 if x:
1164 raise StopIteration
1165
1166 hist = []
1167 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1168 self.assertEqual(hist, [0,1])
1169
1170 hist = []
1171 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1172 self.assertEqual(hist, [0,1])
1173
1174 hist = []
1175 self.assertRaises(AssertionError, list, cycle(gen1()))
1176 self.assertEqual(hist, [0,1])
1177
Georg Brandlb84c1372007-01-21 10:28:43 +00001178class SubclassWithKwargsTest(unittest.TestCase):
1179 def test_keywords_in_subclass(self):
1180 # count is not subclassable...
1181 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
1182 starmap, islice, takewhile, dropwhile, cycle):
1183 class Subclass(cls):
1184 def __init__(self, newarg=None, *args):
1185 cls.__init__(self, *args)
1186 try:
1187 Subclass(newarg=1)
1188 except TypeError, err:
1189 # we expect type errors because of wrong argument count
1190 self.failIf("does not take keyword arguments" in err.args[0])
1191
1192
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001193libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001194
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001195
1196>>> amounts = [120.15, 764.05, 823.14]
1197>>> for checknum, amount in izip(count(1200), amounts):
1198... print 'Check %d is for $%.2f' % (checknum, amount)
1199...
1200Check 1200 is for $120.15
1201Check 1201 is for $764.05
1202Check 1202 is for $823.14
1203
1204>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001205>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1206... print cube
1207...
12081
12098
121027
1211
1212>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001213>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001214... print name.title()
1215...
1216Alex
1217Laura
1218Martin
1219Walter
1220Samuele
1221
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001222>>> from operator import itemgetter
1223>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Armin Rigoa3f09272006-05-28 19:13:17 +00001224>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001225>>> for k, g in groupby(di, itemgetter(1)):
1226... print k, map(itemgetter(0), g)
1227...
12281 ['a', 'c', 'e']
12292 ['b', 'd', 'f']
12303 ['g']
1231
Raymond Hettinger734fb572004-01-20 20:04:40 +00001232# Find runs of consecutive numbers using groupby. The key to the solution
1233# is differencing with a range so that consecutive numbers all appear in
1234# same group.
1235>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Ezio Melotti3efafd72010-08-02 18:40:55 +00001236>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Raymond Hettinger734fb572004-01-20 20:04:40 +00001237... print map(operator.itemgetter(1), g)
1238...
1239[1]
1240[4, 5, 6]
1241[10]
1242[15, 16, 17, 18]
1243[22]
1244[25, 26, 27, 28]
1245
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001246>>> def take(n, iterable):
1247... "Return first n items of the iterable as a list"
1248... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001249
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001250>>> def enumerate(iterable, start=0):
1251... return izip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001252
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001253>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001254... "Return function(0), function(1), ..."
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001255... return imap(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001256
Raymond Hettinger5894c2b2009-02-19 05:38:53 +00001257>>> def nth(iterable, n, default=None):
1258... "Returns the nth item or a default value"
1259... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001260
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001261>>> def quantify(iterable, pred=bool):
1262... "Count how many times the predicate is true"
1263... return sum(imap(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001264
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001265>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001266... "Returns the sequence elements and then returns None indefinitely"
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001267... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001268
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001269>>> def ncycles(iterable, n):
1270... "Returns the seqeuence elements n times"
1271... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001272
1273>>> def dotproduct(vec1, vec2):
1274... return sum(imap(operator.mul, vec1, vec2))
1275
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001276>>> def flatten(listOfLists):
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001277... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001278
1279>>> def repeatfunc(func, times=None, *args):
1280... "Repeat calls to func with specified arguments."
1281... " Example: repeatfunc(random.random)"
1282... if times is None:
1283... return starmap(func, repeat(args))
1284... else:
1285... return starmap(func, repeat(args, times))
1286
Raymond Hettingerd591f662003-10-26 15:34:50 +00001287>>> def pairwise(iterable):
1288... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1289... a, b = tee(iterable)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001290... for elem in b:
1291... break
Raymond Hettingerad983e72003-11-12 14:32:26 +00001292... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001293
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001294>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001295... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001296... args = [iter(iterable)] * n
Raymond Hettingerf080e6d2008-07-31 01:19:50 +00001297... return izip_longest(fillvalue=fillvalue, *args)
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001298
1299>>> def roundrobin(*iterables):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001300... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001301... # Recipe credited to George Sakkis
1302... pending = len(iterables)
1303... nexts = cycle(iter(it).next for it in iterables)
1304... while pending:
1305... try:
1306... for next in nexts:
1307... yield next()
1308... except StopIteration:
1309... pending -= 1
1310... nexts = cycle(islice(nexts, pending))
1311
1312>>> def powerset(iterable):
Raymond Hettinger5cd012b2009-01-26 02:17:52 +00001313... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1314... s = list(iterable)
1315... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001316
Raymond Hettingere8b4b602008-03-11 00:19:07 +00001317>>> def compress(data, selectors):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001318... "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
1319... return (d for d, s in izip(data, selectors) if s)
Raymond Hettingere8b4b602008-03-11 00:19:07 +00001320
Raymond Hettinger33691672008-07-19 00:43:00 +00001321>>> def combinations_with_replacement(iterable, r):
1322... "combinations_with_replacement('ABC', 3) --> AA AB AC BB BC CC"
1323... pool = tuple(iterable)
1324... n = len(pool)
Raymond Hettingercdc9f2c2009-01-27 06:38:00 +00001325... if not n and r:
1326... return
Raymond Hettinger33691672008-07-19 00:43:00 +00001327... indices = [0] * r
1328... yield tuple(pool[i] for i in indices)
1329... while 1:
1330... for i in reversed(range(r)):
1331... if indices[i] != n - 1:
1332... break
1333... else:
1334... return
1335... indices[i:] = [indices[i] + 1] * (r - i)
1336... yield tuple(pool[i] for i in indices)
1337
Raymond Hettingerd8461652009-01-02 21:20:38 +00001338>>> def unique_everseen(iterable, key=None):
1339... "List unique elements, preserving order. Remember all elements ever seen."
1340... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1341... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1342... seen = set()
1343... seen_add = seen.add
1344... if key is None:
1345... for element in iterable:
1346... if element not in seen:
1347... seen_add(element)
1348... yield element
1349... else:
1350... for element in iterable:
1351... k = key(element)
1352... if k not in seen:
1353... seen_add(k)
1354... yield element
1355
1356>>> def unique_justseen(iterable, key=None):
1357... "List unique elements, preserving order. Remember only the element just seen."
1358... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1359... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1360... return imap(next, imap(itemgetter(1), groupby(iterable, key)))
1361
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001362This is not part of the examples but it tests to make sure the definitions
1363perform as purported.
1364
Raymond Hettingera098b332003-09-08 23:58:40 +00001365>>> take(10, count())
1366[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1367
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001368>>> list(enumerate('abc'))
1369[(0, 'a'), (1, 'b'), (2, 'c')]
1370
1371>>> list(islice(tabulate(lambda x: 2*x), 4))
1372[0, 2, 4, 6]
1373
1374>>> nth('abcde', 3)
Raymond Hettinger11485b42009-02-04 19:34:31 +00001375'd'
1376
1377>>> nth('abcde', 9) is None
1378True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001379
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001380>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000138150
1382
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001383>>> a = [[1, 2, 3], [4, 5, 6]]
1384>>> flatten(a)
1385[1, 2, 3, 4, 5, 6]
1386
1387>>> list(repeatfunc(pow, 5, 2, 3))
1388[8, 8, 8, 8, 8]
1389
1390>>> import random
1391>>> take(5, imap(int, repeatfunc(random.random)))
1392[0, 0, 0, 0, 0]
1393
Raymond Hettingerd591f662003-10-26 15:34:50 +00001394>>> list(pairwise('abcd'))
1395[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001396
Raymond Hettingerd591f662003-10-26 15:34:50 +00001397>>> list(pairwise([]))
1398[]
1399
1400>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001401[]
1402
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001403>>> list(islice(padnone('abc'), 0, 6))
1404['a', 'b', 'c', None, None, None]
1405
1406>>> list(ncycles('abc', 3))
1407['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1408
1409>>> dotproduct([1,2,3], [4,5,6])
141032
1411
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001412>>> list(grouper(3, 'abcdefg', 'x'))
1413[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1414
1415>>> list(roundrobin('abc', 'd', 'ef'))
1416['a', 'd', 'e', 'b', 'f', 'c']
1417
Raymond Hettinger5cd012b2009-01-26 02:17:52 +00001418>>> list(powerset([1,2,3]))
1419[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001420
Raymond Hettingere8b4b602008-03-11 00:19:07 +00001421>>> list(compress('abcdef', [1,0,1,0,1,1]))
1422['a', 'c', 'e', 'f']
1423
Raymond Hettinger33691672008-07-19 00:43:00 +00001424>>> list(combinations_with_replacement('abc', 2))
1425[('a', 'a'), ('a', 'b'), ('a', 'c'), ('b', 'b'), ('b', 'c'), ('c', 'c')]
1426
Raymond Hettinger825758c2009-01-08 05:20:19 +00001427>>> list(combinations_with_replacement('01', 3))
1428[('0', '0', '0'), ('0', '0', '1'), ('0', '1', '1'), ('1', '1', '1')]
1429
1430>>> def combinations_with_replacement2(iterable, r):
1431... 'Alternate version that filters from product()'
1432... pool = tuple(iterable)
1433... n = len(pool)
1434... for indices in product(range(n), repeat=r):
1435... if sorted(indices) == list(indices):
1436... yield tuple(pool[i] for i in indices)
1437
1438>>> list(combinations_with_replacement('abc', 2)) == list(combinations_with_replacement2('abc', 2))
1439True
1440
1441>>> list(combinations_with_replacement('01', 3)) == list(combinations_with_replacement2('01', 3))
1442True
1443
1444>>> list(combinations_with_replacement('2310', 6)) == list(combinations_with_replacement2('2310', 6))
1445True
1446
Raymond Hettingerd8461652009-01-02 21:20:38 +00001447>>> list(unique_everseen('AAAABBBCCDAABBB'))
1448['A', 'B', 'C', 'D']
1449
1450>>> list(unique_everseen('ABBCcAD', str.lower))
1451['A', 'B', 'C', 'D']
1452
1453>>> list(unique_justseen('AAAABBBCCDAABBB'))
1454['A', 'B', 'C', 'D', 'A', 'B']
1455
1456>>> list(unique_justseen('ABBCcAD', str.lower))
1457['A', 'B', 'C', 'A', 'D']
1458
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001459"""
1460
1461__test__ = {'libreftest' : libreftest}
1462
1463def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001464 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Georg Brandlb84c1372007-01-21 10:28:43 +00001465 RegressionTests, LengthTransparency,
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001466 SubclassWithKwargsTest, TestExamples)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001467 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001468
1469 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001470 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001471 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001472 counts = [None] * 5
1473 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001474 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001475 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001476 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001477 print counts
1478
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001479 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001480 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001481
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001482if __name__ == "__main__":
1483 test_main(verbose=True)