blob: 9b399c0f665bd83a4fe507c76578bb375a09af6f [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 Hettinger5b913e32009-01-08 06:39:04 +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 Hettinger5b913e32009-01-08 06:39:04 +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 Hettinger5b913e32009-01-08 06:39:04 +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 Hettinger5b913e32009-01-08 06:39:04 +0000111 for r in range(n+2):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000112 result = list(combinations(values, r))
Raymond Hettinger5b913e32009-01-08 06:39:04 +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 Hettinger5b913e32009-01-08 06:39:04 +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 Hettinger5b913e32009-01-08 06:39:04 +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 Hettinger5b913e32009-01-08 06:39:04 +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 Hettinger5b913e32009-01-08 06:39:04 +0000175 for r in range(n+2):
Raymond Hettingerd553d852008-03-04 04:17:08 +0000176 result = list(permutations(values, r))
Raymond Hettinger5b913e32009-01-08 06:39:04 +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 Hettinger5b913e32009-01-08 06:39:04 +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 Hettinger2bcb8e92009-01-25 21:04:14 +0000194 def test_compress(self):
195 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
196 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
197 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
198 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
199 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
200 n = 10000
201 data = chain.from_iterable(repeat(range(6), n))
202 selectors = chain.from_iterable(repeat((0, 1)))
203 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
204 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
205 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
206 self.assertRaises(TypeError, compress, range(6)) # too few args
207 self.assertRaises(TypeError, compress, range(6), None) # too many args
208
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000209 def test_count(self):
210 self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
211 self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000212 self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000213 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
214 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000215 self.assertRaises(TypeError, count, 2, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000216 self.assertRaises(TypeError, count, 'a')
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000217 self.assertEqual(list(islice(count(maxsize-5), 10)), range(maxsize-5, maxsize+5))
218 self.assertEqual(list(islice(count(-maxsize-5), 10)), range(-maxsize-5, -maxsize+5))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000219 c = count(3)
220 self.assertEqual(repr(c), 'count(3)')
221 c.next()
222 self.assertEqual(repr(c), 'count(4)')
Jack Diederich36234e82006-09-21 17:50:26 +0000223 c = count(-9)
224 self.assertEqual(repr(c), 'count(-9)')
225 c.next()
226 self.assertEqual(c.next(), -8)
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000227 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 +0000228 # Test repr (ignoring the L in longs)
229 r1 = repr(count(i)).replace('L', '')
230 r2 = 'count(%r)'.__mod__(i).replace('L', '')
231 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000232
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000233 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000234 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000235 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000236 self.assertRaises(TypeError, cycle)
237 self.assertRaises(TypeError, cycle, 5)
238 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000239
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000240 def test_groupby(self):
241 # Check whether it accepts arguments correctly
242 self.assertEqual([], list(groupby([])))
243 self.assertEqual([], list(groupby([], key=id)))
244 self.assertRaises(TypeError, list, groupby('abc', []))
245 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000246 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000247
248 # Check normal input
249 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
250 (2,15,22), (3,16,23), (3,17,23)]
251 dup = []
252 for k, g in groupby(s, lambda r:r[0]):
253 for elem in g:
254 self.assertEqual(k, elem[0])
255 dup.append(elem)
256 self.assertEqual(s, dup)
257
258 # Check nested case
259 dup = []
260 for k, g in groupby(s, lambda r:r[0]):
261 for ik, ig in groupby(g, lambda r:r[2]):
262 for elem in ig:
263 self.assertEqual(k, elem[0])
264 self.assertEqual(ik, elem[2])
265 dup.append(elem)
266 self.assertEqual(s, dup)
267
268 # Check case where inner iterator is not used
269 keys = [k for k, g in groupby(s, lambda r:r[0])]
270 expectedkeys = set([r[0] for r in s])
271 self.assertEqual(set(keys), expectedkeys)
272 self.assertEqual(len(keys), len(expectedkeys))
273
274 # Exercise pipes and filters style
275 s = 'abracadabra'
276 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000277 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000278 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
279 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000280 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000281 self.assertEqual(r, ['a', 'b', 'r'])
282 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000283 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000284 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
285 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000286 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000287 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
288
289 # iter.next failure
290 class ExpectedError(Exception):
291 pass
292 def delayed_raise(n=0):
293 for i in range(n):
294 yield 'yo'
295 raise ExpectedError
296 def gulp(iterable, keyp=None, func=list):
297 return [func(g) for k, g in groupby(iterable, keyp)]
298
299 # iter.next failure on outer object
300 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
301 # iter.next failure on inner object
302 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
303
304 # __cmp__ failure
305 class DummyCmp:
306 def __cmp__(self, dst):
307 raise ExpectedError
308 s = [DummyCmp(), DummyCmp(), None]
309
310 # __cmp__ failure on outer object
311 self.assertRaises(ExpectedError, gulp, s, func=id)
312 # __cmp__ failure on inner object
313 self.assertRaises(ExpectedError, gulp, s)
314
315 # keyfunc failure
316 def keyfunc(obj):
317 if keyfunc.skip > 0:
318 keyfunc.skip -= 1
319 return obj
320 else:
321 raise ExpectedError
322
323 # keyfunc failure on outer object
324 keyfunc.skip = 0
325 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
326 keyfunc.skip = 1
327 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
328
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000329 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000330 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000331 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000332 self.assertEqual(list(ifilter(bool, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000333 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000334 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000335 self.assertRaises(TypeError, ifilter, lambda x:x)
336 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000337 self.assertRaises(TypeError, ifilter, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000338 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000339
340 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000341 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
342 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000343 self.assertEqual(list(ifilterfalse(bool, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000344 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000345 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000346 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
347 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000348 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000349 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000350
351 def test_izip(self):
352 ans = [(x,y) for x, y in izip('abc',count())]
353 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000354 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
355 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
Raymond Hettinger02420702003-06-29 20:36:23 +0000356 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000357 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000358 self.assertEqual(list(izip()), zip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000359 self.assertRaises(TypeError, izip, 3)
360 self.assertRaises(TypeError, izip, range(3), 3)
361 # Check tuple re-use (implementation detail)
362 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
363 zip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000364 self.assertEqual([pair for pair in izip('abc', 'def')],
365 zip('abc', 'def'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000366 ids = map(id, izip('abc', 'def'))
367 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000368 ids = map(id, list(izip('abc', 'def')))
369 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000370
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000371 def test_iziplongest(self):
372 for args in [
373 ['abc', range(6)],
374 [range(6), 'abc'],
375 [range(1000), range(2000,2100), range(3000,3050)],
376 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
377 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
378 ]:
379 target = map(None, *args)
380 self.assertEqual(list(izip_longest(*args)), target)
381 self.assertEqual(list(izip_longest(*args, **{})), target)
382 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
383 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Tim Petersea5962f2007-03-12 18:07:52 +0000384
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000385 self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
386
387 self.assertEqual(list(izip_longest()), zip())
388 self.assertEqual(list(izip_longest([])), zip([]))
389 self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
Tim Petersea5962f2007-03-12 18:07:52 +0000390
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000391 self.assertEqual(list(izip_longest('abc', 'defg', **{})), map(None, 'abc', 'defg')) # empty keyword dict
392 self.assertRaises(TypeError, izip_longest, 3)
393 self.assertRaises(TypeError, izip_longest, range(3), 3)
394
395 for stmt in [
396 "izip_longest('abc', fv=1)",
Tim Petersea5962f2007-03-12 18:07:52 +0000397 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000398 ]:
399 try:
400 eval(stmt, globals(), locals())
401 except TypeError:
402 pass
403 else:
404 self.fail('Did not raise Type in: ' + stmt)
Tim Petersea5962f2007-03-12 18:07:52 +0000405
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000406 # Check tuple re-use (implementation detail)
407 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
408 zip('abc', 'def'))
409 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
410 zip('abc', 'def'))
411 ids = map(id, izip_longest('abc', 'def'))
412 self.assertEqual(min(ids), max(ids))
413 ids = map(id, list(izip_longest('abc', 'def')))
414 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
415
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000416 def test_product(self):
417 for args, result in [
Raymond Hettingerd553d852008-03-04 04:17:08 +0000418 ([], [()]), # zero iterables
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000419 (['ab'], [('a',), ('b',)]), # one iterable
420 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
421 ([range(0), range(2), range(3)], []), # first iterable with zero length
422 ([range(2), range(0), range(3)], []), # middle iterable with zero length
423 ([range(2), range(3), range(0)], []), # last iterable with zero length
424 ]:
425 self.assertEqual(list(product(*args)), result)
Raymond Hettinger08ff6822008-02-29 02:21:48 +0000426 for r in range(4):
427 self.assertEqual(list(product(*(args*r))),
428 list(product(*args, **dict(repeat=r))))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000429 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
430 self.assertRaises(TypeError, product, range(6), None)
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000431
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000432 def product1(*args, **kwds):
433 pools = map(tuple, args) * kwds.get('repeat', 1)
434 n = len(pools)
435 if n == 0:
436 yield ()
437 return
438 if any(len(pool) == 0 for pool in pools):
439 return
440 indices = [0] * n
441 yield tuple(pool[i] for pool, i in zip(pools, indices))
442 while 1:
443 for i in reversed(range(n)): # right to left
444 if indices[i] == len(pools[i]) - 1:
445 continue
446 indices[i] += 1
447 for j in range(i+1, n):
448 indices[j] = 0
449 yield tuple(pool[i] for pool, i in zip(pools, indices))
450 break
451 else:
452 return
453
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000454 def product2(*args, **kwds):
455 'Pure python version used in docs'
456 pools = map(tuple, args) * kwds.get('repeat', 1)
457 result = [[]]
458 for pool in pools:
459 result = [x+[y] for x in result for y in pool]
460 for prod in result:
461 yield tuple(prod)
462
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000463 argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
464 set('abcdefg'), range(11), tuple(range(13))]
465 for i in range(100):
466 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Raymond Hettingerd553d852008-03-04 04:17:08 +0000467 expected_len = prod(map(len, args))
468 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000469 self.assertEqual(list(product(*args)), list(product1(*args)))
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000470 self.assertEqual(list(product(*args)), list(product2(*args)))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000471 args = map(iter, args)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000472 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000473
Raymond Hettinger73d79632008-02-23 02:20:41 +0000474 # Test implementation detail: tuple re-use
475 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
476 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000477
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000478 def test_repeat(self):
479 self.assertEqual(zip(xrange(3),repeat('a')),
480 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000481 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000482 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000483 self.assertEqual(list(repeat('a', 0)), [])
484 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000485 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000486 self.assertRaises(TypeError, repeat, None, 3, 4)
487 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000488 r = repeat(1+0j)
489 self.assertEqual(repr(r), 'repeat((1+0j))')
490 r = repeat(1+0j, 5)
491 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
492 list(r)
493 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000494
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000495 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000496 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
497 [0**1, 1**2, 2**3])
498 self.assertEqual(list(imap(None, 'abc', range(5))),
499 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000500 self.assertEqual(list(imap(None, 'abc', count())),
501 [('a',0),('b',1),('c',2)])
502 self.assertEqual(take(2,imap(None, 'abc', count())),
503 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000504 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000505 self.assertRaises(TypeError, imap)
506 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000507 self.assertRaises(TypeError, imap(10, range(5)).next)
508 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
509 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000510
511 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000512 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
513 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000514 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
515 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000516 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger47317092008-01-17 03:02:14 +0000517 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
518 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000519 self.assertRaises(TypeError, starmap)
520 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
521 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
522 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
523 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000524
525 def test_islice(self):
526 for args in [ # islice(args) should agree with range(args)
527 (10, 20, 3),
528 (10, 3, 20),
529 (10, 20),
530 (10, 3),
531 (20,)
532 ]:
533 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
534
535 for args, tgtargs in [ # Stop when seqn is exhausted
536 ((10, 110, 3), ((10, 100, 3))),
537 ((10, 110), ((10, 100))),
538 ((110,), (100,))
539 ]:
540 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
541
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000542 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000543 self.assertEqual(list(islice(xrange(10), None)), range(10))
Raymond Hettingerb2594052004-12-05 09:25:51 +0000544 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
545 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000546 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
547 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
548
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000549 # Test number of items consumed SF #1171417
550 it = iter(range(10))
551 self.assertEqual(list(islice(it, 3)), range(3))
552 self.assertEqual(list(it), range(3, 10))
553
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000554 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000555 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000556 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
557 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
558 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
559 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
560 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000561 self.assertRaises(ValueError, islice, xrange(10), 'a')
562 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
563 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
564 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
565 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +0000566 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000567
568 def test_takewhile(self):
569 data = [1, 3, 5, 20, 2, 4, 6, 8]
570 underten = lambda x: x<10
571 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000572 self.assertEqual(list(takewhile(underten, [])), [])
573 self.assertRaises(TypeError, takewhile)
574 self.assertRaises(TypeError, takewhile, operator.pow)
575 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
576 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
577 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000578 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
579 self.assertEqual(list(t), [1, 1, 1])
580 self.assertRaises(StopIteration, t.next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000581
582 def test_dropwhile(self):
583 data = [1, 3, 5, 20, 2, 4, 6, 8]
584 underten = lambda x: x<10
585 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000586 self.assertEqual(list(dropwhile(underten, [])), [])
587 self.assertRaises(TypeError, dropwhile)
588 self.assertRaises(TypeError, dropwhile, operator.pow)
589 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
590 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
591 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000592
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000593 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000594 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000595 def irange(n):
596 for i in xrange(n):
597 yield i
598
599 a, b = tee([]) # test empty iterator
600 self.assertEqual(list(a), [])
601 self.assertEqual(list(b), [])
602
603 a, b = tee(irange(n)) # test 100% interleaved
604 self.assertEqual(zip(a,b), zip(range(n),range(n)))
605
606 a, b = tee(irange(n)) # test 0% interleaved
607 self.assertEqual(list(a), range(n))
608 self.assertEqual(list(b), range(n))
609
610 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000611 for i in xrange(100):
612 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000613 del a
614 self.assertEqual(list(b), range(n))
615
616 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000617 for i in xrange(100):
618 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000619 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000620 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000621
622 for j in xrange(5): # test randomly interleaved
623 order = [0]*n + [1]*n
624 random.shuffle(order)
625 lists = ([], [])
626 its = tee(irange(n))
627 for i in order:
628 value = its[i].next()
629 lists[i].append(value)
630 self.assertEqual(lists[0], range(n))
631 self.assertEqual(lists[1], range(n))
632
Raymond Hettingerad983e72003-11-12 14:32:26 +0000633 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000634 self.assertRaises(TypeError, tee)
635 self.assertRaises(TypeError, tee, 3)
636 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000637 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000638
Raymond Hettingerad983e72003-11-12 14:32:26 +0000639 # tee object should be instantiable
640 a, b = tee('abc')
641 c = type(a)('def')
642 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000643
Raymond Hettingerad983e72003-11-12 14:32:26 +0000644 # test long-lagged and multi-way split
645 a, b, c = tee(xrange(2000), 3)
646 for i in xrange(100):
647 self.assertEqual(a.next(), i)
648 self.assertEqual(list(b), range(2000))
649 self.assertEqual([c.next(), c.next()], range(2))
650 self.assertEqual(list(a), range(100,2000))
651 self.assertEqual(list(c), range(2,2000))
652
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000653 # test values of n
654 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Neal Norwitz69e88972006-09-02 02:58:13 +0000655 self.assertRaises(ValueError, tee, [], -1)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000656 for n in xrange(5):
657 result = tee('abc', n)
658 self.assertEqual(type(result), tuple)
659 self.assertEqual(len(result), n)
660 self.assertEqual(map(list, result), [list('abc')]*n)
661
Raymond Hettingerad983e72003-11-12 14:32:26 +0000662 # tee pass-through to copyable iterator
663 a, b = tee('abc')
664 c, d = tee(a)
665 self.assert_(a is c)
666
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000667 # test tee_new
668 t1, t2 = tee('abc')
669 tnew = type(t1)
670 self.assertRaises(TypeError, tnew)
671 self.assertRaises(TypeError, tnew, 10)
672 t3 = tnew(t1)
673 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000674
Raymond Hettingera9f60922004-10-17 16:40:14 +0000675 # test that tee objects are weak referencable
676 a, b = tee(xrange(10))
677 p = proxy(a)
678 self.assertEqual(getattr(p, '__class__'), type(b))
679 del a
680 self.assertRaises(ReferenceError, getattr, p, '__class__')
681
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000682 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000683 self.assertRaises(StopIteration, izip().next)
684
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000685 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000686 self.assertRaises(StopIteration, f([]).next)
687 self.assertRaises(StopIteration, f(StopNow()).next)
688
689 self.assertRaises(StopIteration, islice([], None).next)
690 self.assertRaises(StopIteration, islice(StopNow(), None).next)
691
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000692 p, q = tee([])
693 self.assertRaises(StopIteration, p.next)
694 self.assertRaises(StopIteration, q.next)
695 p, q = tee(StopNow())
696 self.assertRaises(StopIteration, p.next)
697 self.assertRaises(StopIteration, q.next)
698
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000699 self.assertRaises(StopIteration, repeat(None, 0).next)
700
701 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
702 self.assertRaises(StopIteration, f(lambda x:x, []).next)
703 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
704
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000705class TestExamples(unittest.TestCase):
706
707 def test_chain(self):
708 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
709
710 def test_chain_from_iterable(self):
711 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
712
713 def test_combinations(self):
714 self.assertEqual(list(combinations('ABCD', 2)),
715 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
716 self.assertEqual(list(combinations(range(4), 3)),
717 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
718
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000719 def test_compress(self):
720 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
721
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000722 def test_count(self):
723 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
724
725 def test_cycle(self):
726 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
727
728 def test_dropwhile(self):
729 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
730
731 def test_groupby(self):
732 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
733 list('ABCDAB'))
734 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
735 [list('AAAA'), list('BBB'), list('CC'), list('D')])
736
737 def test_ifilter(self):
738 self.assertEqual(list(ifilter(lambda x: x%2, range(10))), [1,3,5,7,9])
739
740 def test_ifilterfalse(self):
741 self.assertEqual(list(ifilterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
742
743 def test_imap(self):
744 self.assertEqual(list(imap(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
745
746 def test_islice(self):
747 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
748 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
749 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
750 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
751
752 def test_izip(self):
753 self.assertEqual(list(izip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
754
755 def test_izip_longest(self):
756 self.assertEqual(list(izip_longest('ABCD', 'xy', fillvalue='-')),
757 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
758
759 def test_permutations(self):
760 self.assertEqual(list(permutations('ABCD', 2)),
761 map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split()))
762 self.assertEqual(list(permutations(range(3))),
763 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
764
765 def test_product(self):
766 self.assertEqual(list(product('ABCD', 'xy')),
767 map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split()))
768 self.assertEqual(list(product(range(2), repeat=3)),
769 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
770 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
771
772 def test_repeat(self):
773 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
774
775 def test_stapmap(self):
776 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
777 [32, 9, 1000])
778
779 def test_takewhile(self):
780 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
781
782
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000783class TestGC(unittest.TestCase):
784
785 def makecycle(self, iterator, container):
786 container.append(iterator)
787 iterator.next()
788 del container, iterator
789
790 def test_chain(self):
791 a = []
792 self.makecycle(chain(a), a)
793
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000794 def test_chain_from_iterable(self):
795 a = []
796 self.makecycle(chain.from_iterable([a]), a)
797
798 def test_combinations(self):
799 a = []
800 self.makecycle(combinations([1,2,a,3], 3), a)
801
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000802 def test_compress(self):
803 a = []
804 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
805
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000806 def test_cycle(self):
807 a = []
808 self.makecycle(cycle([a]*2), a)
809
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000810 def test_dropwhile(self):
811 a = []
812 self.makecycle(dropwhile(bool, [0, a, a]), a)
813
814 def test_groupby(self):
815 a = []
816 self.makecycle(groupby([a]*2, lambda x:x), a)
817
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000818 def test_issue2246(self):
819 # Issue 2246 -- the _grouper iterator was not included in GC
820 n = 10
821 keyfunc = lambda x: x
822 for i, j in groupby(xrange(n), key=keyfunc):
823 keyfunc.__dict__.setdefault('x',[]).append(j)
824
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000825 def test_ifilter(self):
826 a = []
827 self.makecycle(ifilter(lambda x:True, [a]*2), a)
828
829 def test_ifilterfalse(self):
830 a = []
831 self.makecycle(ifilterfalse(lambda x:False, a), a)
832
833 def test_izip(self):
834 a = []
835 self.makecycle(izip([a]*2, [a]*3), a)
836
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000837 def test_izip_longest(self):
838 a = []
839 self.makecycle(izip_longest([a]*2, [a]*3), a)
840 b = [a, None]
841 self.makecycle(izip_longest([a]*2, [a]*3, fillvalue=b), a)
842
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000843 def test_imap(self):
844 a = []
845 self.makecycle(imap(lambda x:x, [a]*2), a)
846
847 def test_islice(self):
848 a = []
849 self.makecycle(islice([a]*2, None), a)
850
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000851 def test_permutations(self):
852 a = []
853 self.makecycle(permutations([1,2,a,3], 3), a)
854
855 def test_product(self):
856 a = []
857 self.makecycle(product([1,2,a,3], repeat=3), a)
858
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000859 def test_repeat(self):
860 a = []
861 self.makecycle(repeat(a), a)
862
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000863 def test_starmap(self):
864 a = []
865 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
866
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000867 def test_takewhile(self):
868 a = []
869 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
870
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000871def R(seqn):
872 'Regular generator'
873 for i in seqn:
874 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000875
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000876class G:
877 'Sequence using __getitem__'
878 def __init__(self, seqn):
879 self.seqn = seqn
880 def __getitem__(self, i):
881 return self.seqn[i]
882
883class I:
884 'Sequence using iterator protocol'
885 def __init__(self, seqn):
886 self.seqn = seqn
887 self.i = 0
888 def __iter__(self):
889 return self
890 def next(self):
891 if self.i >= len(self.seqn): raise StopIteration
892 v = self.seqn[self.i]
893 self.i += 1
894 return v
895
896class Ig:
897 'Sequence using iterator protocol defined with a generator'
898 def __init__(self, seqn):
899 self.seqn = seqn
900 self.i = 0
901 def __iter__(self):
902 for val in self.seqn:
903 yield val
904
905class X:
906 'Missing __getitem__ and __iter__'
907 def __init__(self, seqn):
908 self.seqn = seqn
909 self.i = 0
910 def next(self):
911 if self.i >= len(self.seqn): raise StopIteration
912 v = self.seqn[self.i]
913 self.i += 1
914 return v
915
916class N:
917 'Iterator missing next()'
918 def __init__(self, seqn):
919 self.seqn = seqn
920 self.i = 0
921 def __iter__(self):
922 return self
923
924class E:
925 'Test propagation of exceptions'
926 def __init__(self, seqn):
927 self.seqn = seqn
928 self.i = 0
929 def __iter__(self):
930 return self
931 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +0000932 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000933
934class S:
935 'Test immediate stop'
936 def __init__(self, seqn):
937 pass
938 def __iter__(self):
939 return self
940 def next(self):
941 raise StopIteration
942
943def L(seqn):
944 'Test multiple tiers of iterators'
945 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
946
947
948class TestVariousIteratorArgs(unittest.TestCase):
949
950 def test_chain(self):
951 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
952 for g in (G, I, Ig, S, L, R):
953 self.assertEqual(list(chain(g(s))), list(g(s)))
954 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Raymond Hettinger05bf6332008-02-28 22:30:42 +0000955 self.assertRaises(TypeError, list, chain(X(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000956 self.assertRaises(TypeError, list, chain(N(s)))
957 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
958
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000959 def test_compress(self):
960 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
961 n = len(s)
962 for g in (G, I, Ig, S, L, R):
963 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
964 self.assertRaises(TypeError, compress, X(s), repeat(1))
965 self.assertRaises(TypeError, list, compress(N(s), repeat(1)))
966 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
967
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000968 def test_product(self):
969 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
970 self.assertRaises(TypeError, product, X(s))
971 self.assertRaises(TypeError, product, N(s))
972 self.assertRaises(ZeroDivisionError, product, E(s))
973
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000974 def test_cycle(self):
975 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
976 for g in (G, I, Ig, S, L, R):
977 tgtlen = len(s) * 3
978 expected = list(g(s))*3
979 actual = list(islice(cycle(g(s)), tgtlen))
980 self.assertEqual(actual, expected)
981 self.assertRaises(TypeError, cycle, X(s))
982 self.assertRaises(TypeError, list, cycle(N(s)))
983 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
984
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000985 def test_groupby(self):
986 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
987 for g in (G, I, Ig, S, L, R):
988 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
989 self.assertRaises(TypeError, groupby, X(s))
990 self.assertRaises(TypeError, list, groupby(N(s)))
991 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
992
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000993 def test_ifilter(self):
994 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
995 for g in (G, I, Ig, S, L, R):
996 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
997 self.assertRaises(TypeError, ifilter, isEven, X(s))
998 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
999 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
1000
1001 def test_ifilterfalse(self):
1002 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1003 for g in (G, I, Ig, S, L, R):
1004 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
1005 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
1006 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
1007 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
1008
1009 def test_izip(self):
1010 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1011 for g in (G, I, Ig, S, L, R):
1012 self.assertEqual(list(izip(g(s))), zip(g(s)))
1013 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
1014 self.assertRaises(TypeError, izip, X(s))
1015 self.assertRaises(TypeError, list, izip(N(s)))
1016 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
1017
Raymond Hettingerd36862c2007-02-21 05:20:38 +00001018 def test_iziplongest(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_longest(g(s))), zip(g(s)))
1022 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
1023 self.assertRaises(TypeError, izip_longest, X(s))
1024 self.assertRaises(TypeError, list, izip_longest(N(s)))
1025 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
1026
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001027 def test_imap(self):
1028 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1029 for g in (G, I, Ig, S, L, R):
1030 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
1031 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
1032 self.assertRaises(TypeError, imap, onearg, X(s))
1033 self.assertRaises(TypeError, list, imap(onearg, N(s)))
1034 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
1035
1036 def test_islice(self):
1037 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1038 for g in (G, I, Ig, S, L, R):
1039 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1040 self.assertRaises(TypeError, islice, X(s), 10)
1041 self.assertRaises(TypeError, list, islice(N(s), 10))
1042 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1043
1044 def test_starmap(self):
1045 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1046 for g in (G, I, Ig, S, L, R):
1047 ss = zip(s, s)
1048 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
1049 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
1050 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
1051 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1052
1053 def test_takewhile(self):
1054 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1055 for g in (G, I, Ig, S, L, R):
1056 tgt = []
1057 for elem in g(s):
1058 if not isEven(elem): break
1059 tgt.append(elem)
1060 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1061 self.assertRaises(TypeError, takewhile, isEven, X(s))
1062 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
1063 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1064
1065 def test_dropwhile(self):
1066 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1067 for g in (G, I, Ig, S, L, R):
1068 tgt = []
1069 for elem in g(s):
1070 if not tgt and isOdd(elem): continue
1071 tgt.append(elem)
1072 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1073 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
1074 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
1075 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1076
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001077 def test_tee(self):
1078 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1079 for g in (G, I, Ig, S, L, R):
1080 it1, it2 = tee(g(s))
1081 self.assertEqual(list(it1), list(g(s)))
1082 self.assertEqual(list(it2), list(g(s)))
1083 self.assertRaises(TypeError, tee, X(s))
1084 self.assertRaises(TypeError, list, tee(N(s))[0])
1085 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1086
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001087class LengthTransparency(unittest.TestCase):
1088
1089 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001090 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001091 self.assertEqual(len(repeat(None, 50)), 50)
1092 self.assertRaises(TypeError, len, repeat(None))
1093
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001094class RegressionTests(unittest.TestCase):
1095
1096 def test_sf_793826(self):
1097 # Fix Armin Rigo's successful efforts to wreak havoc
1098
1099 def mutatingtuple(tuple1, f, tuple2):
1100 # this builds a tuple t which is a copy of tuple1,
1101 # then calls f(t), then mutates t to be equal to tuple2
1102 # (needs len(tuple1) == len(tuple2)).
1103 def g(value, first=[1]):
1104 if first:
1105 del first[:]
1106 f(z.next())
1107 return value
1108 items = list(tuple2)
1109 items[1:1] = list(tuple1)
1110 gen = imap(g, items)
1111 z = izip(*[gen]*len(tuple1))
1112 z.next()
1113
1114 def f(t):
1115 global T
1116 T = t
1117 first[:] = list(T)
1118
1119 first = []
1120 mutatingtuple((1,2,3), f, (4,5,6))
1121 second = list(T)
1122 self.assertEqual(first, second)
1123
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001124
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001125 def test_sf_950057(self):
1126 # Make sure that chain() and cycle() catch exceptions immediately
1127 # rather than when shifting between input sources
1128
1129 def gen1():
1130 hist.append(0)
1131 yield 1
1132 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001133 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001134 hist.append(2)
1135
1136 def gen2(x):
1137 hist.append(3)
1138 yield 2
1139 hist.append(4)
1140 if x:
1141 raise StopIteration
1142
1143 hist = []
1144 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1145 self.assertEqual(hist, [0,1])
1146
1147 hist = []
1148 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1149 self.assertEqual(hist, [0,1])
1150
1151 hist = []
1152 self.assertRaises(AssertionError, list, cycle(gen1()))
1153 self.assertEqual(hist, [0,1])
1154
Georg Brandlb84c1372007-01-21 10:28:43 +00001155class SubclassWithKwargsTest(unittest.TestCase):
1156 def test_keywords_in_subclass(self):
1157 # count is not subclassable...
1158 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001159 starmap, islice, takewhile, dropwhile, cycle, compress):
Georg Brandlb84c1372007-01-21 10:28:43 +00001160 class Subclass(cls):
1161 def __init__(self, newarg=None, *args):
1162 cls.__init__(self, *args)
1163 try:
1164 Subclass(newarg=1)
1165 except TypeError, err:
1166 # we expect type errors because of wrong argument count
1167 self.failIf("does not take keyword arguments" in err.args[0])
1168
1169
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001170libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001171
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001172
1173>>> amounts = [120.15, 764.05, 823.14]
1174>>> for checknum, amount in izip(count(1200), amounts):
1175... print 'Check %d is for $%.2f' % (checknum, amount)
1176...
1177Check 1200 is for $120.15
1178Check 1201 is for $764.05
1179Check 1202 is for $823.14
1180
1181>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001182>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1183... print cube
1184...
11851
11868
118727
1188
1189>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001190>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001191... print name.title()
1192...
1193Alex
1194Laura
1195Martin
1196Walter
1197Samuele
1198
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001199>>> from operator import itemgetter
1200>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Armin Rigoa3f09272006-05-28 19:13:17 +00001201>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001202>>> for k, g in groupby(di, itemgetter(1)):
1203... print k, map(itemgetter(0), g)
1204...
12051 ['a', 'c', 'e']
12062 ['b', 'd', 'f']
12073 ['g']
1208
Raymond Hettinger734fb572004-01-20 20:04:40 +00001209# Find runs of consecutive numbers using groupby. The key to the solution
1210# is differencing with a range so that consecutive numbers all appear in
1211# same group.
1212>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
1213>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
1214... print map(operator.itemgetter(1), g)
1215...
1216[1]
1217[4, 5, 6]
1218[10]
1219[15, 16, 17, 18]
1220[22]
1221[25, 26, 27, 28]
1222
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001223>>> def take(n, iterable):
1224... "Return first n items of the iterable as a list"
1225... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001226
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001227>>> def enumerate(iterable, start=0):
1228... return izip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001229
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001230>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001231... "Return function(0), function(1), ..."
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001232... return imap(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001233
1234>>> def nth(iterable, n):
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001235... "Returns the nth item or empty list"
Raymond Hettinger60eca932003-02-09 06:40:58 +00001236... return list(islice(iterable, n, n+1))
1237
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001238>>> def quantify(iterable, pred=bool):
1239... "Count how many times the predicate is true"
1240... return sum(imap(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001241
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001242>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001243... "Returns the sequence elements and then returns None indefinitely"
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001244... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001245
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001246>>> def ncycles(iterable, n):
1247... "Returns the seqeuence elements n times"
1248... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001249
1250>>> def dotproduct(vec1, vec2):
1251... return sum(imap(operator.mul, vec1, vec2))
1252
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001253>>> def flatten(listOfLists):
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001254... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001255
1256>>> def repeatfunc(func, times=None, *args):
1257... "Repeat calls to func with specified arguments."
1258... " Example: repeatfunc(random.random)"
1259... if times is None:
1260... return starmap(func, repeat(args))
1261... else:
1262... return starmap(func, repeat(args, times))
1263
Raymond Hettingerd591f662003-10-26 15:34:50 +00001264>>> def pairwise(iterable):
1265... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1266... a, b = tee(iterable)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001267... for elem in b:
1268... break
Raymond Hettingerad983e72003-11-12 14:32:26 +00001269... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001270
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001271>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001272... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001273... args = [iter(iterable)] * n
Raymond Hettingerf080e6d2008-07-31 01:19:50 +00001274... return izip_longest(fillvalue=fillvalue, *args)
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001275
1276>>> def roundrobin(*iterables):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001277... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001278... # Recipe credited to George Sakkis
1279... pending = len(iterables)
1280... nexts = cycle(iter(it).next for it in iterables)
1281... while pending:
1282... try:
1283... for next in nexts:
1284... yield next()
1285... except StopIteration:
1286... pending -= 1
1287... nexts = cycle(islice(nexts, pending))
1288
1289>>> def powerset(iterable):
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001290... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1291... s = list(iterable)
1292... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001293
Raymond Hettinger33691672008-07-19 00:43:00 +00001294>>> def combinations_with_replacement(iterable, r):
1295... "combinations_with_replacement('ABC', 3) --> AA AB AC BB BC CC"
1296... pool = tuple(iterable)
1297... n = len(pool)
1298... indices = [0] * r
1299... yield tuple(pool[i] for i in indices)
1300... while 1:
1301... for i in reversed(range(r)):
1302... if indices[i] != n - 1:
1303... break
1304... else:
1305... return
1306... indices[i:] = [indices[i] + 1] * (r - i)
1307... yield tuple(pool[i] for i in indices)
1308
Raymond Hettinger44e15812009-01-02 21:26:45 +00001309>>> def unique_everseen(iterable, key=None):
1310... "List unique elements, preserving order. Remember all elements ever seen."
1311... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1312... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1313... seen = set()
1314... seen_add = seen.add
1315... if key is None:
1316... for element in iterable:
1317... if element not in seen:
1318... seen_add(element)
1319... yield element
1320... else:
1321... for element in iterable:
1322... k = key(element)
1323... if k not in seen:
1324... seen_add(k)
1325... yield element
1326
1327>>> def unique_justseen(iterable, key=None):
1328... "List unique elements, preserving order. Remember only the element just seen."
1329... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1330... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1331... return imap(next, imap(itemgetter(1), groupby(iterable, key)))
1332
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001333This is not part of the examples but it tests to make sure the definitions
1334perform as purported.
1335
Raymond Hettingera098b332003-09-08 23:58:40 +00001336>>> take(10, count())
1337[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1338
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001339>>> list(enumerate('abc'))
1340[(0, 'a'), (1, 'b'), (2, 'c')]
1341
1342>>> list(islice(tabulate(lambda x: 2*x), 4))
1343[0, 2, 4, 6]
1344
1345>>> nth('abcde', 3)
1346['d']
1347
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001348>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000134950
1350
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001351>>> a = [[1, 2, 3], [4, 5, 6]]
1352>>> flatten(a)
1353[1, 2, 3, 4, 5, 6]
1354
1355>>> list(repeatfunc(pow, 5, 2, 3))
1356[8, 8, 8, 8, 8]
1357
1358>>> import random
1359>>> take(5, imap(int, repeatfunc(random.random)))
1360[0, 0, 0, 0, 0]
1361
Raymond Hettingerd591f662003-10-26 15:34:50 +00001362>>> list(pairwise('abcd'))
1363[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001364
Raymond Hettingerd591f662003-10-26 15:34:50 +00001365>>> list(pairwise([]))
1366[]
1367
1368>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001369[]
1370
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001371>>> list(islice(padnone('abc'), 0, 6))
1372['a', 'b', 'c', None, None, None]
1373
1374>>> list(ncycles('abc', 3))
1375['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1376
1377>>> dotproduct([1,2,3], [4,5,6])
137832
1379
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001380>>> list(grouper(3, 'abcdefg', 'x'))
1381[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1382
1383>>> list(roundrobin('abc', 'd', 'ef'))
1384['a', 'd', 'e', 'b', 'f', 'c']
1385
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001386>>> list(powerset([1,2,3]))
1387[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001388
Raymond Hettinger33691672008-07-19 00:43:00 +00001389>>> list(combinations_with_replacement('abc', 2))
1390[('a', 'a'), ('a', 'b'), ('a', 'c'), ('b', 'b'), ('b', 'c'), ('c', 'c')]
1391
Raymond Hettinger5b913e32009-01-08 06:39:04 +00001392>>> list(combinations_with_replacement('01', 3))
1393[('0', '0', '0'), ('0', '0', '1'), ('0', '1', '1'), ('1', '1', '1')]
1394
1395>>> def combinations_with_replacement2(iterable, r):
1396... 'Alternate version that filters from product()'
1397... pool = tuple(iterable)
1398... n = len(pool)
1399... for indices in product(range(n), repeat=r):
1400... if sorted(indices) == list(indices):
1401... yield tuple(pool[i] for i in indices)
1402
1403>>> list(combinations_with_replacement('abc', 2)) == list(combinations_with_replacement2('abc', 2))
1404True
1405
1406>>> list(combinations_with_replacement('01', 3)) == list(combinations_with_replacement2('01', 3))
1407True
1408
1409>>> list(combinations_with_replacement('2310', 6)) == list(combinations_with_replacement2('2310', 6))
1410True
1411
Raymond Hettinger44e15812009-01-02 21:26:45 +00001412>>> list(unique_everseen('AAAABBBCCDAABBB'))
1413['A', 'B', 'C', 'D']
1414
1415>>> list(unique_everseen('ABBCcAD', str.lower))
1416['A', 'B', 'C', 'D']
1417
1418>>> list(unique_justseen('AAAABBBCCDAABBB'))
1419['A', 'B', 'C', 'D', 'A', 'B']
1420
1421>>> list(unique_justseen('ABBCcAD', str.lower))
1422['A', 'B', 'C', 'A', 'D']
1423
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001424"""
1425
1426__test__ = {'libreftest' : libreftest}
1427
1428def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001429 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Georg Brandlb84c1372007-01-21 10:28:43 +00001430 RegressionTests, LengthTransparency,
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001431 SubclassWithKwargsTest, TestExamples)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001432 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001433
1434 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001435 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001436 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001437 counts = [None] * 5
1438 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001439 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001440 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001441 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001442 print counts
1443
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001444 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001445 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001446
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001447if __name__ == "__main__":
1448 test_main(verbose=True)