blob: 0f1f695a2dff92e046d1d664efbe3f8b2ac4e3d9 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001import unittest
2from test import test_support
3from itertools import *
Raymond Hettingera9f60922004-10-17 16:40:14 +00004from weakref import proxy
Raymond Hettinger2012f172003-02-07 05:32:58 +00005import sys
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00006import operator
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00007import random
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +00008maxsize = test_support.MAX_Py_ssize_t
9minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000010
11def onearg(x):
12 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000013 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000014
15def errfunc(*args):
16 'Test function that raises an error'
17 raise ValueError
18
19def gen3():
20 'Non-restartable source sequence'
21 for i in (0, 1, 2):
22 yield i
23
24def isEven(x):
25 'Test predicate'
26 return x%2==0
27
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000028def isOdd(x):
29 'Test predicate'
30 return x%2==1
31
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000032class StopNow:
33 'Class emulating an empty iterable.'
34 def __iter__(self):
35 return self
36 def next(self):
37 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000038
Raymond Hettinger02420702003-06-29 20:36:23 +000039def take(n, seq):
40 'Convenience function for partially consuming a long of infinite iterable'
41 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000042
Raymond Hettingerd553d852008-03-04 04:17:08 +000043def prod(iterable):
44 return reduce(operator.mul, iterable, 1)
45
Raymond Hettinger93e804d2008-02-26 23:40:50 +000046def fact(n):
47 'Factorial'
Raymond Hettingerd553d852008-03-04 04:17:08 +000048 return prod(range(1, n+1))
49
Raymond Hettinger96ef8112003-02-01 00:10:11 +000050class TestBasicOps(unittest.TestCase):
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000051 def test_chain(self):
Raymond Hettingerad47fa12008-03-06 20:52:01 +000052
53 def chain2(*iterables):
54 'Pure python version in the docs'
55 for it in iterables:
56 for element in it:
57 yield element
58
59 for c in (chain, chain2):
60 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
61 self.assertEqual(list(c('abc')), list('abc'))
62 self.assertEqual(list(c('')), [])
63 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
64 self.assertRaises(TypeError, list,c(2, 3))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000065
Raymond Hettingerb4cbc982008-02-28 22:46:41 +000066 def test_chain_from_iterable(self):
67 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
68 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
69 self.assertEqual(list(chain.from_iterable([''])), [])
70 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
71 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
72
Raymond Hettinger93e804d2008-02-26 23:40:50 +000073 def test_combinations(self):
Raymond Hettinger825758c2009-01-08 05:20:19 +000074 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Raymond Hettinger93e804d2008-02-26 23:40:50 +000075 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Raymond Hettingerd553d852008-03-04 04:17:08 +000076 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Raymond Hettinger93e804d2008-02-26 23:40:50 +000077 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Raymond Hettinger825758c2009-01-08 05:20:19 +000078 self.assertEqual(list(combinations('abc', 32)), []) # r > n
Raymond Hettinger93e804d2008-02-26 23:40:50 +000079 self.assertEqual(list(combinations(range(4), 3)),
80 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
Raymond Hettingerd553d852008-03-04 04:17:08 +000081
82 def combinations1(iterable, r):
83 'Pure python version shown in the docs'
84 pool = tuple(iterable)
85 n = len(pool)
Raymond Hettinger825758c2009-01-08 05:20:19 +000086 if r > n:
87 return
Raymond Hettingerd553d852008-03-04 04:17:08 +000088 indices = range(r)
89 yield tuple(pool[i] for i in indices)
90 while 1:
91 for i in reversed(range(r)):
92 if indices[i] != i + n - r:
93 break
94 else:
95 return
96 indices[i] += 1
97 for j in range(i+1, r):
98 indices[j] = indices[j-1] + 1
99 yield tuple(pool[i] for i in indices)
100
101 def combinations2(iterable, r):
102 'Pure python version shown in the docs'
103 pool = tuple(iterable)
104 n = len(pool)
105 for indices in permutations(range(n), r):
106 if sorted(indices) == list(indices):
107 yield tuple(pool[i] for i in indices)
108
109 for n in range(7):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000110 values = [5*x-12 for x in range(n)]
Raymond Hettinger825758c2009-01-08 05:20:19 +0000111 for r in range(n+2):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000112 result = list(combinations(values, r))
Raymond Hettinger825758c2009-01-08 05:20:19 +0000113 self.assertEqual(len(result), 0 if r>n else fact(n) / fact(r) / fact(n-r)) # right number of combs
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000114 self.assertEqual(len(result), len(set(result))) # no repeats
115 self.assertEqual(result, sorted(result)) # lexicographic order
116 for c in result:
117 self.assertEqual(len(c), r) # r-length combinations
118 self.assertEqual(len(set(c)), r) # no duplicate elements
119 self.assertEqual(list(c), sorted(c)) # keep original ordering
120 self.assert_(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000121 self.assertEqual(list(c),
122 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Raymond Hettingerd553d852008-03-04 04:17:08 +0000123 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger825758c2009-01-08 05:20:19 +0000124 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000125
126 # Test implementation detail: tuple re-use
127 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
128 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
129
130 def test_permutations(self):
131 self.assertRaises(TypeError, permutations) # too few arguments
132 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000133 self.assertRaises(TypeError, permutations, None) # pool is not iterable
134 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger825758c2009-01-08 05:20:19 +0000135 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000136 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Raymond Hettingerd553d852008-03-04 04:17:08 +0000137 self.assertEqual(list(permutations(range(3), 2)),
138 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
139
140 def permutations1(iterable, r=None):
141 'Pure python version shown in the docs'
142 pool = tuple(iterable)
143 n = len(pool)
144 r = n if r is None else r
Raymond Hettinger825758c2009-01-08 05:20:19 +0000145 if r > n:
146 return
Raymond Hettingerd553d852008-03-04 04:17:08 +0000147 indices = range(n)
Raymond Hettingere70bb8d2008-03-23 00:55:46 +0000148 cycles = range(n, n-r, -1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000149 yield tuple(pool[i] for i in indices[:r])
150 while n:
151 for i in reversed(range(r)):
152 cycles[i] -= 1
153 if cycles[i] == 0:
154 indices[i:] = indices[i+1:] + indices[i:i+1]
155 cycles[i] = n - i
156 else:
157 j = cycles[i]
158 indices[i], indices[-j] = indices[-j], indices[i]
159 yield tuple(pool[i] for i in indices[:r])
160 break
161 else:
162 return
163
164 def permutations2(iterable, r=None):
165 'Pure python version shown in the docs'
166 pool = tuple(iterable)
167 n = len(pool)
168 r = n if r is None else r
169 for indices in product(range(n), repeat=r):
170 if len(set(indices)) == r:
171 yield tuple(pool[i] for i in indices)
172
173 for n in range(7):
174 values = [5*x-12 for x in range(n)]
Raymond Hettinger825758c2009-01-08 05:20:19 +0000175 for r in range(n+2):
Raymond Hettingerd553d852008-03-04 04:17:08 +0000176 result = list(permutations(values, r))
Raymond Hettinger825758c2009-01-08 05:20:19 +0000177 self.assertEqual(len(result), 0 if r>n else fact(n) / fact(n-r)) # right number of perms
Raymond Hettingerd553d852008-03-04 04:17:08 +0000178 self.assertEqual(len(result), len(set(result))) # no repeats
179 self.assertEqual(result, sorted(result)) # lexicographic order
180 for p in result:
181 self.assertEqual(len(p), r) # r-length permutations
182 self.assertEqual(len(set(p)), r) # no duplicate elements
183 self.assert_(all(e in values for e in p)) # elements taken from input iterable
184 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger825758c2009-01-08 05:20:19 +0000185 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000186 if r == n:
187 self.assertEqual(result, list(permutations(values, None))) # test r as None
188 self.assertEqual(result, list(permutations(values))) # test default r
189
190 # Test implementation detail: tuple re-use
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000191 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000192 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000193
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000194 def test_count(self):
195 self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
196 self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000197 self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000198 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
199 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000200 self.assertRaises(TypeError, count, 2, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000201 self.assertRaises(TypeError, count, 'a')
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000202 self.assertEqual(list(islice(count(maxsize-5), 10)), range(maxsize-5, maxsize+5))
203 self.assertEqual(list(islice(count(-maxsize-5), 10)), range(-maxsize-5, -maxsize+5))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000204 c = count(3)
205 self.assertEqual(repr(c), 'count(3)')
206 c.next()
207 self.assertEqual(repr(c), 'count(4)')
Jack Diederich36234e82006-09-21 17:50:26 +0000208 c = count(-9)
209 self.assertEqual(repr(c), 'count(-9)')
210 c.next()
211 self.assertEqual(c.next(), -8)
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000212 for i in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 10, sys.maxint-5, sys.maxint+5):
Raymond Hettinger3a0de082007-10-12 17:53:11 +0000213 # Test repr (ignoring the L in longs)
214 r1 = repr(count(i)).replace('L', '')
215 r2 = 'count(%r)'.__mod__(i).replace('L', '')
216 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000217
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000218 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000219 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000220 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000221 self.assertRaises(TypeError, cycle)
222 self.assertRaises(TypeError, cycle, 5)
223 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000224
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000225 def test_groupby(self):
226 # Check whether it accepts arguments correctly
227 self.assertEqual([], list(groupby([])))
228 self.assertEqual([], list(groupby([], key=id)))
229 self.assertRaises(TypeError, list, groupby('abc', []))
230 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000231 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000232
233 # Check normal input
234 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
235 (2,15,22), (3,16,23), (3,17,23)]
236 dup = []
237 for k, g in groupby(s, lambda r:r[0]):
238 for elem in g:
239 self.assertEqual(k, elem[0])
240 dup.append(elem)
241 self.assertEqual(s, dup)
242
243 # Check nested case
244 dup = []
245 for k, g in groupby(s, lambda r:r[0]):
246 for ik, ig in groupby(g, lambda r:r[2]):
247 for elem in ig:
248 self.assertEqual(k, elem[0])
249 self.assertEqual(ik, elem[2])
250 dup.append(elem)
251 self.assertEqual(s, dup)
252
253 # Check case where inner iterator is not used
254 keys = [k for k, g in groupby(s, lambda r:r[0])]
255 expectedkeys = set([r[0] for r in s])
256 self.assertEqual(set(keys), expectedkeys)
257 self.assertEqual(len(keys), len(expectedkeys))
258
259 # Exercise pipes and filters style
260 s = 'abracadabra'
261 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000262 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000263 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
264 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000265 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000266 self.assertEqual(r, ['a', 'b', 'r'])
267 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000268 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000269 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
270 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000271 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000272 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
273
274 # iter.next failure
275 class ExpectedError(Exception):
276 pass
277 def delayed_raise(n=0):
278 for i in range(n):
279 yield 'yo'
280 raise ExpectedError
281 def gulp(iterable, keyp=None, func=list):
282 return [func(g) for k, g in groupby(iterable, keyp)]
283
284 # iter.next failure on outer object
285 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
286 # iter.next failure on inner object
287 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
288
289 # __cmp__ failure
290 class DummyCmp:
291 def __cmp__(self, dst):
292 raise ExpectedError
293 s = [DummyCmp(), DummyCmp(), None]
294
295 # __cmp__ failure on outer object
296 self.assertRaises(ExpectedError, gulp, s, func=id)
297 # __cmp__ failure on inner object
298 self.assertRaises(ExpectedError, gulp, s)
299
300 # keyfunc failure
301 def keyfunc(obj):
302 if keyfunc.skip > 0:
303 keyfunc.skip -= 1
304 return obj
305 else:
306 raise ExpectedError
307
308 # keyfunc failure on outer object
309 keyfunc.skip = 0
310 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
311 keyfunc.skip = 1
312 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
313
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000314 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000315 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000316 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000317 self.assertEqual(list(ifilter(bool, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000318 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000319 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000320 self.assertRaises(TypeError, ifilter, lambda x:x)
321 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000322 self.assertRaises(TypeError, ifilter, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000323 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000324
325 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000326 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
327 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000328 self.assertEqual(list(ifilterfalse(bool, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000329 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000330 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000331 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
332 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000333 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000334 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000335
336 def test_izip(self):
337 ans = [(x,y) for x, y in izip('abc',count())]
338 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000339 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
340 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
Raymond Hettinger02420702003-06-29 20:36:23 +0000341 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000342 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000343 self.assertEqual(list(izip()), zip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000344 self.assertRaises(TypeError, izip, 3)
345 self.assertRaises(TypeError, izip, range(3), 3)
346 # Check tuple re-use (implementation detail)
347 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
348 zip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000349 self.assertEqual([pair for pair in izip('abc', 'def')],
350 zip('abc', 'def'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000351 ids = map(id, izip('abc', 'def'))
352 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000353 ids = map(id, list(izip('abc', 'def')))
354 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000355
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000356 def test_iziplongest(self):
357 for args in [
358 ['abc', range(6)],
359 [range(6), 'abc'],
360 [range(1000), range(2000,2100), range(3000,3050)],
361 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
362 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
363 ]:
364 target = map(None, *args)
365 self.assertEqual(list(izip_longest(*args)), target)
366 self.assertEqual(list(izip_longest(*args, **{})), target)
367 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
368 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Tim Petersea5962f2007-03-12 18:07:52 +0000369
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000370 self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
371
372 self.assertEqual(list(izip_longest()), zip())
373 self.assertEqual(list(izip_longest([])), zip([]))
374 self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
Tim Petersea5962f2007-03-12 18:07:52 +0000375
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000376 self.assertEqual(list(izip_longest('abc', 'defg', **{})), map(None, 'abc', 'defg')) # empty keyword dict
377 self.assertRaises(TypeError, izip_longest, 3)
378 self.assertRaises(TypeError, izip_longest, range(3), 3)
379
380 for stmt in [
381 "izip_longest('abc', fv=1)",
Tim Petersea5962f2007-03-12 18:07:52 +0000382 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000383 ]:
384 try:
385 eval(stmt, globals(), locals())
386 except TypeError:
387 pass
388 else:
389 self.fail('Did not raise Type in: ' + stmt)
Tim Petersea5962f2007-03-12 18:07:52 +0000390
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000391 # Check tuple re-use (implementation detail)
392 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
393 zip('abc', 'def'))
394 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
395 zip('abc', 'def'))
396 ids = map(id, izip_longest('abc', 'def'))
397 self.assertEqual(min(ids), max(ids))
398 ids = map(id, list(izip_longest('abc', 'def')))
399 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
400
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000401 def test_product(self):
402 for args, result in [
Raymond Hettingerd553d852008-03-04 04:17:08 +0000403 ([], [()]), # zero iterables
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000404 (['ab'], [('a',), ('b',)]), # one iterable
405 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
406 ([range(0), range(2), range(3)], []), # first iterable with zero length
407 ([range(2), range(0), range(3)], []), # middle iterable with zero length
408 ([range(2), range(3), range(0)], []), # last iterable with zero length
409 ]:
410 self.assertEqual(list(product(*args)), result)
Raymond Hettinger08ff6822008-02-29 02:21:48 +0000411 for r in range(4):
412 self.assertEqual(list(product(*(args*r))),
413 list(product(*args, **dict(repeat=r))))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000414 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
415 self.assertRaises(TypeError, product, range(6), None)
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000416
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000417 def product1(*args, **kwds):
418 pools = map(tuple, args) * kwds.get('repeat', 1)
419 n = len(pools)
420 if n == 0:
421 yield ()
422 return
423 if any(len(pool) == 0 for pool in pools):
424 return
425 indices = [0] * n
426 yield tuple(pool[i] for pool, i in zip(pools, indices))
427 while 1:
428 for i in reversed(range(n)): # right to left
429 if indices[i] == len(pools[i]) - 1:
430 continue
431 indices[i] += 1
432 for j in range(i+1, n):
433 indices[j] = 0
434 yield tuple(pool[i] for pool, i in zip(pools, indices))
435 break
436 else:
437 return
438
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000439 def product2(*args, **kwds):
440 'Pure python version used in docs'
441 pools = map(tuple, args) * kwds.get('repeat', 1)
442 result = [[]]
443 for pool in pools:
444 result = [x+[y] for x in result for y in pool]
445 for prod in result:
446 yield tuple(prod)
447
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000448 argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
449 set('abcdefg'), range(11), tuple(range(13))]
450 for i in range(100):
451 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Raymond Hettingerd553d852008-03-04 04:17:08 +0000452 expected_len = prod(map(len, args))
453 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000454 self.assertEqual(list(product(*args)), list(product1(*args)))
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000455 self.assertEqual(list(product(*args)), list(product2(*args)))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000456 args = map(iter, args)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000457 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000458
Raymond Hettinger73d79632008-02-23 02:20:41 +0000459 # Test implementation detail: tuple re-use
460 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
461 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000462
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000463 def test_repeat(self):
464 self.assertEqual(zip(xrange(3),repeat('a')),
465 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000466 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000467 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000468 self.assertEqual(list(repeat('a', 0)), [])
469 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000470 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000471 self.assertRaises(TypeError, repeat, None, 3, 4)
472 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000473 r = repeat(1+0j)
474 self.assertEqual(repr(r), 'repeat((1+0j))')
475 r = repeat(1+0j, 5)
476 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
477 list(r)
478 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000479
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000480 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000481 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
482 [0**1, 1**2, 2**3])
483 self.assertEqual(list(imap(None, 'abc', range(5))),
484 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000485 self.assertEqual(list(imap(None, 'abc', count())),
486 [('a',0),('b',1),('c',2)])
487 self.assertEqual(take(2,imap(None, 'abc', count())),
488 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000489 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000490 self.assertRaises(TypeError, imap)
491 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000492 self.assertRaises(TypeError, imap(10, range(5)).next)
493 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
494 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000495
496 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000497 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
498 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000499 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
500 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000501 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger47317092008-01-17 03:02:14 +0000502 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
503 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000504 self.assertRaises(TypeError, starmap)
505 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
506 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
507 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
508 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000509
510 def test_islice(self):
511 for args in [ # islice(args) should agree with range(args)
512 (10, 20, 3),
513 (10, 3, 20),
514 (10, 20),
515 (10, 3),
516 (20,)
517 ]:
518 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
519
520 for args, tgtargs in [ # Stop when seqn is exhausted
521 ((10, 110, 3), ((10, 100, 3))),
522 ((10, 110), ((10, 100))),
523 ((110,), (100,))
524 ]:
525 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
526
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000527 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000528 self.assertEqual(list(islice(xrange(10), None)), range(10))
Raymond Hettingerb2594052004-12-05 09:25:51 +0000529 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
530 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000531 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
532 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
533
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000534 # Test number of items consumed SF #1171417
535 it = iter(range(10))
536 self.assertEqual(list(islice(it, 3)), range(3))
537 self.assertEqual(list(it), range(3, 10))
538
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000539 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000540 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000541 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
542 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
543 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
544 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
545 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000546 self.assertRaises(ValueError, islice, xrange(10), 'a')
547 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
548 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
549 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
550 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +0000551 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000552
553 def test_takewhile(self):
554 data = [1, 3, 5, 20, 2, 4, 6, 8]
555 underten = lambda x: x<10
556 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000557 self.assertEqual(list(takewhile(underten, [])), [])
558 self.assertRaises(TypeError, takewhile)
559 self.assertRaises(TypeError, takewhile, operator.pow)
560 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
561 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
562 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000563 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
564 self.assertEqual(list(t), [1, 1, 1])
565 self.assertRaises(StopIteration, t.next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000566
567 def test_dropwhile(self):
568 data = [1, 3, 5, 20, 2, 4, 6, 8]
569 underten = lambda x: x<10
570 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000571 self.assertEqual(list(dropwhile(underten, [])), [])
572 self.assertRaises(TypeError, dropwhile)
573 self.assertRaises(TypeError, dropwhile, operator.pow)
574 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
575 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
576 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000577
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000578 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000579 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000580 def irange(n):
581 for i in xrange(n):
582 yield i
583
584 a, b = tee([]) # test empty iterator
585 self.assertEqual(list(a), [])
586 self.assertEqual(list(b), [])
587
588 a, b = tee(irange(n)) # test 100% interleaved
589 self.assertEqual(zip(a,b), zip(range(n),range(n)))
590
591 a, b = tee(irange(n)) # test 0% interleaved
592 self.assertEqual(list(a), range(n))
593 self.assertEqual(list(b), range(n))
594
595 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000596 for i in xrange(100):
597 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000598 del a
599 self.assertEqual(list(b), range(n))
600
601 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000602 for i in xrange(100):
603 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000604 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000605 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000606
607 for j in xrange(5): # test randomly interleaved
608 order = [0]*n + [1]*n
609 random.shuffle(order)
610 lists = ([], [])
611 its = tee(irange(n))
612 for i in order:
613 value = its[i].next()
614 lists[i].append(value)
615 self.assertEqual(lists[0], range(n))
616 self.assertEqual(lists[1], range(n))
617
Raymond Hettingerad983e72003-11-12 14:32:26 +0000618 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000619 self.assertRaises(TypeError, tee)
620 self.assertRaises(TypeError, tee, 3)
621 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000622 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000623
Raymond Hettingerad983e72003-11-12 14:32:26 +0000624 # tee object should be instantiable
625 a, b = tee('abc')
626 c = type(a)('def')
627 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000628
Raymond Hettingerad983e72003-11-12 14:32:26 +0000629 # test long-lagged and multi-way split
630 a, b, c = tee(xrange(2000), 3)
631 for i in xrange(100):
632 self.assertEqual(a.next(), i)
633 self.assertEqual(list(b), range(2000))
634 self.assertEqual([c.next(), c.next()], range(2))
635 self.assertEqual(list(a), range(100,2000))
636 self.assertEqual(list(c), range(2,2000))
637
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000638 # test values of n
639 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Neal Norwitz69e88972006-09-02 02:58:13 +0000640 self.assertRaises(ValueError, tee, [], -1)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000641 for n in xrange(5):
642 result = tee('abc', n)
643 self.assertEqual(type(result), tuple)
644 self.assertEqual(len(result), n)
645 self.assertEqual(map(list, result), [list('abc')]*n)
646
Raymond Hettingerad983e72003-11-12 14:32:26 +0000647 # tee pass-through to copyable iterator
648 a, b = tee('abc')
649 c, d = tee(a)
650 self.assert_(a is c)
651
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000652 # test tee_new
653 t1, t2 = tee('abc')
654 tnew = type(t1)
655 self.assertRaises(TypeError, tnew)
656 self.assertRaises(TypeError, tnew, 10)
657 t3 = tnew(t1)
658 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000659
Raymond Hettingera9f60922004-10-17 16:40:14 +0000660 # test that tee objects are weak referencable
661 a, b = tee(xrange(10))
662 p = proxy(a)
663 self.assertEqual(getattr(p, '__class__'), type(b))
664 del a
665 self.assertRaises(ReferenceError, getattr, p, '__class__')
666
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000667 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000668 self.assertRaises(StopIteration, izip().next)
669
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000670 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000671 self.assertRaises(StopIteration, f([]).next)
672 self.assertRaises(StopIteration, f(StopNow()).next)
673
674 self.assertRaises(StopIteration, islice([], None).next)
675 self.assertRaises(StopIteration, islice(StopNow(), None).next)
676
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000677 p, q = tee([])
678 self.assertRaises(StopIteration, p.next)
679 self.assertRaises(StopIteration, q.next)
680 p, q = tee(StopNow())
681 self.assertRaises(StopIteration, p.next)
682 self.assertRaises(StopIteration, q.next)
683
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000684 self.assertRaises(StopIteration, repeat(None, 0).next)
685
686 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
687 self.assertRaises(StopIteration, f(lambda x:x, []).next)
688 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
689
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000690class TestExamples(unittest.TestCase):
691
692 def test_chain(self):
693 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
694
695 def test_chain_from_iterable(self):
696 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
697
698 def test_combinations(self):
699 self.assertEqual(list(combinations('ABCD', 2)),
700 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
701 self.assertEqual(list(combinations(range(4), 3)),
702 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
703
704 def test_count(self):
705 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
706
707 def test_cycle(self):
708 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
709
710 def test_dropwhile(self):
711 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
712
713 def test_groupby(self):
714 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
715 list('ABCDAB'))
716 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
717 [list('AAAA'), list('BBB'), list('CC'), list('D')])
718
719 def test_ifilter(self):
720 self.assertEqual(list(ifilter(lambda x: x%2, range(10))), [1,3,5,7,9])
721
722 def test_ifilterfalse(self):
723 self.assertEqual(list(ifilterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
724
725 def test_imap(self):
726 self.assertEqual(list(imap(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
727
728 def test_islice(self):
729 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
730 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
731 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
732 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
733
734 def test_izip(self):
735 self.assertEqual(list(izip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
736
737 def test_izip_longest(self):
738 self.assertEqual(list(izip_longest('ABCD', 'xy', fillvalue='-')),
739 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
740
741 def test_permutations(self):
742 self.assertEqual(list(permutations('ABCD', 2)),
743 map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split()))
744 self.assertEqual(list(permutations(range(3))),
745 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
746
747 def test_product(self):
748 self.assertEqual(list(product('ABCD', 'xy')),
749 map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split()))
750 self.assertEqual(list(product(range(2), repeat=3)),
751 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
752 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
753
754 def test_repeat(self):
755 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
756
757 def test_stapmap(self):
758 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
759 [32, 9, 1000])
760
761 def test_takewhile(self):
762 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
763
764
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000765class TestGC(unittest.TestCase):
766
767 def makecycle(self, iterator, container):
768 container.append(iterator)
769 iterator.next()
770 del container, iterator
771
772 def test_chain(self):
773 a = []
774 self.makecycle(chain(a), a)
775
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000776 def test_chain_from_iterable(self):
777 a = []
778 self.makecycle(chain.from_iterable([a]), a)
779
780 def test_combinations(self):
781 a = []
782 self.makecycle(combinations([1,2,a,3], 3), a)
783
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000784 def test_cycle(self):
785 a = []
786 self.makecycle(cycle([a]*2), a)
787
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000788 def test_dropwhile(self):
789 a = []
790 self.makecycle(dropwhile(bool, [0, a, a]), a)
791
792 def test_groupby(self):
793 a = []
794 self.makecycle(groupby([a]*2, lambda x:x), a)
795
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000796 def test_issue2246(self):
797 # Issue 2246 -- the _grouper iterator was not included in GC
798 n = 10
799 keyfunc = lambda x: x
800 for i, j in groupby(xrange(n), key=keyfunc):
801 keyfunc.__dict__.setdefault('x',[]).append(j)
802
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000803 def test_ifilter(self):
804 a = []
805 self.makecycle(ifilter(lambda x:True, [a]*2), a)
806
807 def test_ifilterfalse(self):
808 a = []
809 self.makecycle(ifilterfalse(lambda x:False, a), a)
810
811 def test_izip(self):
812 a = []
813 self.makecycle(izip([a]*2, [a]*3), a)
814
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000815 def test_izip_longest(self):
816 a = []
817 self.makecycle(izip_longest([a]*2, [a]*3), a)
818 b = [a, None]
819 self.makecycle(izip_longest([a]*2, [a]*3, fillvalue=b), a)
820
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000821 def test_imap(self):
822 a = []
823 self.makecycle(imap(lambda x:x, [a]*2), a)
824
825 def test_islice(self):
826 a = []
827 self.makecycle(islice([a]*2, None), a)
828
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000829 def test_permutations(self):
830 a = []
831 self.makecycle(permutations([1,2,a,3], 3), a)
832
833 def test_product(self):
834 a = []
835 self.makecycle(product([1,2,a,3], repeat=3), a)
836
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000837 def test_repeat(self):
838 a = []
839 self.makecycle(repeat(a), a)
840
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000841 def test_starmap(self):
842 a = []
843 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
844
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000845 def test_takewhile(self):
846 a = []
847 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
848
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000849def R(seqn):
850 'Regular generator'
851 for i in seqn:
852 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000853
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000854class G:
855 'Sequence using __getitem__'
856 def __init__(self, seqn):
857 self.seqn = seqn
858 def __getitem__(self, i):
859 return self.seqn[i]
860
861class I:
862 'Sequence using iterator protocol'
863 def __init__(self, seqn):
864 self.seqn = seqn
865 self.i = 0
866 def __iter__(self):
867 return self
868 def next(self):
869 if self.i >= len(self.seqn): raise StopIteration
870 v = self.seqn[self.i]
871 self.i += 1
872 return v
873
874class Ig:
875 'Sequence using iterator protocol defined with a generator'
876 def __init__(self, seqn):
877 self.seqn = seqn
878 self.i = 0
879 def __iter__(self):
880 for val in self.seqn:
881 yield val
882
883class X:
884 'Missing __getitem__ and __iter__'
885 def __init__(self, seqn):
886 self.seqn = seqn
887 self.i = 0
888 def next(self):
889 if self.i >= len(self.seqn): raise StopIteration
890 v = self.seqn[self.i]
891 self.i += 1
892 return v
893
894class N:
895 'Iterator missing next()'
896 def __init__(self, seqn):
897 self.seqn = seqn
898 self.i = 0
899 def __iter__(self):
900 return self
901
902class E:
903 'Test propagation of exceptions'
904 def __init__(self, seqn):
905 self.seqn = seqn
906 self.i = 0
907 def __iter__(self):
908 return self
909 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +0000910 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000911
912class S:
913 'Test immediate stop'
914 def __init__(self, seqn):
915 pass
916 def __iter__(self):
917 return self
918 def next(self):
919 raise StopIteration
920
921def L(seqn):
922 'Test multiple tiers of iterators'
923 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
924
925
926class TestVariousIteratorArgs(unittest.TestCase):
927
928 def test_chain(self):
929 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
930 for g in (G, I, Ig, S, L, R):
931 self.assertEqual(list(chain(g(s))), list(g(s)))
932 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Raymond Hettinger05bf6332008-02-28 22:30:42 +0000933 self.assertRaises(TypeError, list, chain(X(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000934 self.assertRaises(TypeError, list, chain(N(s)))
935 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
936
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000937 def test_product(self):
938 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
939 self.assertRaises(TypeError, product, X(s))
940 self.assertRaises(TypeError, product, N(s))
941 self.assertRaises(ZeroDivisionError, product, E(s))
942
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000943 def test_cycle(self):
944 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
945 for g in (G, I, Ig, S, L, R):
946 tgtlen = len(s) * 3
947 expected = list(g(s))*3
948 actual = list(islice(cycle(g(s)), tgtlen))
949 self.assertEqual(actual, expected)
950 self.assertRaises(TypeError, cycle, X(s))
951 self.assertRaises(TypeError, list, cycle(N(s)))
952 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
953
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000954 def test_groupby(self):
955 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
956 for g in (G, I, Ig, S, L, R):
957 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
958 self.assertRaises(TypeError, groupby, X(s))
959 self.assertRaises(TypeError, list, groupby(N(s)))
960 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
961
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000962 def test_ifilter(self):
963 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
964 for g in (G, I, Ig, S, L, R):
965 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
966 self.assertRaises(TypeError, ifilter, isEven, X(s))
967 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
968 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
969
970 def test_ifilterfalse(self):
971 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
972 for g in (G, I, Ig, S, L, R):
973 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
974 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
975 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
976 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
977
978 def test_izip(self):
979 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
980 for g in (G, I, Ig, S, L, R):
981 self.assertEqual(list(izip(g(s))), zip(g(s)))
982 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
983 self.assertRaises(TypeError, izip, X(s))
984 self.assertRaises(TypeError, list, izip(N(s)))
985 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
986
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000987 def test_iziplongest(self):
988 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
989 for g in (G, I, Ig, S, L, R):
990 self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
991 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
992 self.assertRaises(TypeError, izip_longest, X(s))
993 self.assertRaises(TypeError, list, izip_longest(N(s)))
994 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
995
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000996 def test_imap(self):
997 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
998 for g in (G, I, Ig, S, L, R):
999 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
1000 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
1001 self.assertRaises(TypeError, imap, onearg, X(s))
1002 self.assertRaises(TypeError, list, imap(onearg, N(s)))
1003 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
1004
1005 def test_islice(self):
1006 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1007 for g in (G, I, Ig, S, L, R):
1008 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1009 self.assertRaises(TypeError, islice, X(s), 10)
1010 self.assertRaises(TypeError, list, islice(N(s), 10))
1011 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1012
1013 def test_starmap(self):
1014 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1015 for g in (G, I, Ig, S, L, R):
1016 ss = zip(s, s)
1017 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
1018 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
1019 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
1020 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1021
1022 def test_takewhile(self):
1023 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1024 for g in (G, I, Ig, S, L, R):
1025 tgt = []
1026 for elem in g(s):
1027 if not isEven(elem): break
1028 tgt.append(elem)
1029 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1030 self.assertRaises(TypeError, takewhile, isEven, X(s))
1031 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
1032 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1033
1034 def test_dropwhile(self):
1035 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1036 for g in (G, I, Ig, S, L, R):
1037 tgt = []
1038 for elem in g(s):
1039 if not tgt and isOdd(elem): continue
1040 tgt.append(elem)
1041 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1042 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
1043 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
1044 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1045
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001046 def test_tee(self):
1047 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1048 for g in (G, I, Ig, S, L, R):
1049 it1, it2 = tee(g(s))
1050 self.assertEqual(list(it1), list(g(s)))
1051 self.assertEqual(list(it2), list(g(s)))
1052 self.assertRaises(TypeError, tee, X(s))
1053 self.assertRaises(TypeError, list, tee(N(s))[0])
1054 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1055
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001056class LengthTransparency(unittest.TestCase):
1057
1058 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001059 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001060 self.assertEqual(len(repeat(None, 50)), 50)
1061 self.assertRaises(TypeError, len, repeat(None))
1062
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001063class RegressionTests(unittest.TestCase):
1064
1065 def test_sf_793826(self):
1066 # Fix Armin Rigo's successful efforts to wreak havoc
1067
1068 def mutatingtuple(tuple1, f, tuple2):
1069 # this builds a tuple t which is a copy of tuple1,
1070 # then calls f(t), then mutates t to be equal to tuple2
1071 # (needs len(tuple1) == len(tuple2)).
1072 def g(value, first=[1]):
1073 if first:
1074 del first[:]
1075 f(z.next())
1076 return value
1077 items = list(tuple2)
1078 items[1:1] = list(tuple1)
1079 gen = imap(g, items)
1080 z = izip(*[gen]*len(tuple1))
1081 z.next()
1082
1083 def f(t):
1084 global T
1085 T = t
1086 first[:] = list(T)
1087
1088 first = []
1089 mutatingtuple((1,2,3), f, (4,5,6))
1090 second = list(T)
1091 self.assertEqual(first, second)
1092
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001093
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001094 def test_sf_950057(self):
1095 # Make sure that chain() and cycle() catch exceptions immediately
1096 # rather than when shifting between input sources
1097
1098 def gen1():
1099 hist.append(0)
1100 yield 1
1101 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001102 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001103 hist.append(2)
1104
1105 def gen2(x):
1106 hist.append(3)
1107 yield 2
1108 hist.append(4)
1109 if x:
1110 raise StopIteration
1111
1112 hist = []
1113 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1114 self.assertEqual(hist, [0,1])
1115
1116 hist = []
1117 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1118 self.assertEqual(hist, [0,1])
1119
1120 hist = []
1121 self.assertRaises(AssertionError, list, cycle(gen1()))
1122 self.assertEqual(hist, [0,1])
1123
Georg Brandlb84c1372007-01-21 10:28:43 +00001124class SubclassWithKwargsTest(unittest.TestCase):
1125 def test_keywords_in_subclass(self):
1126 # count is not subclassable...
1127 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
1128 starmap, islice, takewhile, dropwhile, cycle):
1129 class Subclass(cls):
1130 def __init__(self, newarg=None, *args):
1131 cls.__init__(self, *args)
1132 try:
1133 Subclass(newarg=1)
1134 except TypeError, err:
1135 # we expect type errors because of wrong argument count
1136 self.failIf("does not take keyword arguments" in err.args[0])
1137
1138
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001139libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001140
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001141
1142>>> amounts = [120.15, 764.05, 823.14]
1143>>> for checknum, amount in izip(count(1200), amounts):
1144... print 'Check %d is for $%.2f' % (checknum, amount)
1145...
1146Check 1200 is for $120.15
1147Check 1201 is for $764.05
1148Check 1202 is for $823.14
1149
1150>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001151>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1152... print cube
1153...
11541
11558
115627
1157
1158>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001159>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001160... print name.title()
1161...
1162Alex
1163Laura
1164Martin
1165Walter
1166Samuele
1167
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001168>>> from operator import itemgetter
1169>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Armin Rigoa3f09272006-05-28 19:13:17 +00001170>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001171>>> for k, g in groupby(di, itemgetter(1)):
1172... print k, map(itemgetter(0), g)
1173...
11741 ['a', 'c', 'e']
11752 ['b', 'd', 'f']
11763 ['g']
1177
Raymond Hettinger734fb572004-01-20 20:04:40 +00001178# Find runs of consecutive numbers using groupby. The key to the solution
1179# is differencing with a range so that consecutive numbers all appear in
1180# same group.
1181>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
1182>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
1183... print map(operator.itemgetter(1), g)
1184...
1185[1]
1186[4, 5, 6]
1187[10]
1188[15, 16, 17, 18]
1189[22]
1190[25, 26, 27, 28]
1191
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001192>>> def take(n, iterable):
1193... "Return first n items of the iterable as a list"
1194... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001195
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001196>>> def enumerate(iterable, start=0):
1197... return izip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001198
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001199>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001200... "Return function(0), function(1), ..."
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001201... return imap(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001202
1203>>> def nth(iterable, n):
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001204... "Returns the nth item or empty list"
Raymond Hettinger60eca932003-02-09 06:40:58 +00001205... return list(islice(iterable, n, n+1))
1206
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001207>>> def quantify(iterable, pred=bool):
1208... "Count how many times the predicate is true"
1209... return sum(imap(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001210
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001211>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001212... "Returns the sequence elements and then returns None indefinitely"
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001213... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001214
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001215>>> def ncycles(iterable, n):
1216... "Returns the seqeuence elements n times"
1217... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001218
1219>>> def dotproduct(vec1, vec2):
1220... return sum(imap(operator.mul, vec1, vec2))
1221
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001222>>> def flatten(listOfLists):
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001223... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001224
1225>>> def repeatfunc(func, times=None, *args):
1226... "Repeat calls to func with specified arguments."
1227... " Example: repeatfunc(random.random)"
1228... if times is None:
1229... return starmap(func, repeat(args))
1230... else:
1231... return starmap(func, repeat(args, times))
1232
Raymond Hettingerd591f662003-10-26 15:34:50 +00001233>>> def pairwise(iterable):
1234... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1235... a, b = tee(iterable)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001236... for elem in b:
1237... break
Raymond Hettingerad983e72003-11-12 14:32:26 +00001238... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001239
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001240>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001241... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001242... args = [iter(iterable)] * n
Raymond Hettingerf080e6d2008-07-31 01:19:50 +00001243... return izip_longest(fillvalue=fillvalue, *args)
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001244
1245>>> def roundrobin(*iterables):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001246... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001247... # Recipe credited to George Sakkis
1248... pending = len(iterables)
1249... nexts = cycle(iter(it).next for it in iterables)
1250... while pending:
1251... try:
1252... for next in nexts:
1253... yield next()
1254... except StopIteration:
1255... pending -= 1
1256... nexts = cycle(islice(nexts, pending))
1257
1258>>> def powerset(iterable):
Raymond Hettinger5cd012b2009-01-26 02:17:52 +00001259... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1260... s = list(iterable)
1261... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001262
Raymond Hettingere8b4b602008-03-11 00:19:07 +00001263>>> def compress(data, selectors):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001264... "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
1265... return (d for d, s in izip(data, selectors) if s)
Raymond Hettingere8b4b602008-03-11 00:19:07 +00001266
Raymond Hettinger33691672008-07-19 00:43:00 +00001267>>> def combinations_with_replacement(iterable, r):
1268... "combinations_with_replacement('ABC', 3) --> AA AB AC BB BC CC"
1269... pool = tuple(iterable)
1270... n = len(pool)
1271... indices = [0] * r
1272... yield tuple(pool[i] for i in indices)
1273... while 1:
1274... for i in reversed(range(r)):
1275... if indices[i] != n - 1:
1276... break
1277... else:
1278... return
1279... indices[i:] = [indices[i] + 1] * (r - i)
1280... yield tuple(pool[i] for i in indices)
1281
Raymond Hettingerd8461652009-01-02 21:20:38 +00001282>>> def unique_everseen(iterable, key=None):
1283... "List unique elements, preserving order. Remember all elements ever seen."
1284... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1285... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1286... seen = set()
1287... seen_add = seen.add
1288... if key is None:
1289... for element in iterable:
1290... if element not in seen:
1291... seen_add(element)
1292... yield element
1293... else:
1294... for element in iterable:
1295... k = key(element)
1296... if k not in seen:
1297... seen_add(k)
1298... yield element
1299
1300>>> def unique_justseen(iterable, key=None):
1301... "List unique elements, preserving order. Remember only the element just seen."
1302... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1303... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1304... return imap(next, imap(itemgetter(1), groupby(iterable, key)))
1305
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001306This is not part of the examples but it tests to make sure the definitions
1307perform as purported.
1308
Raymond Hettingera098b332003-09-08 23:58:40 +00001309>>> take(10, count())
1310[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1311
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001312>>> list(enumerate('abc'))
1313[(0, 'a'), (1, 'b'), (2, 'c')]
1314
1315>>> list(islice(tabulate(lambda x: 2*x), 4))
1316[0, 2, 4, 6]
1317
1318>>> nth('abcde', 3)
1319['d']
1320
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001321>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000132250
1323
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001324>>> a = [[1, 2, 3], [4, 5, 6]]
1325>>> flatten(a)
1326[1, 2, 3, 4, 5, 6]
1327
1328>>> list(repeatfunc(pow, 5, 2, 3))
1329[8, 8, 8, 8, 8]
1330
1331>>> import random
1332>>> take(5, imap(int, repeatfunc(random.random)))
1333[0, 0, 0, 0, 0]
1334
Raymond Hettingerd591f662003-10-26 15:34:50 +00001335>>> list(pairwise('abcd'))
1336[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001337
Raymond Hettingerd591f662003-10-26 15:34:50 +00001338>>> list(pairwise([]))
1339[]
1340
1341>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001342[]
1343
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001344>>> list(islice(padnone('abc'), 0, 6))
1345['a', 'b', 'c', None, None, None]
1346
1347>>> list(ncycles('abc', 3))
1348['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1349
1350>>> dotproduct([1,2,3], [4,5,6])
135132
1352
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001353>>> list(grouper(3, 'abcdefg', 'x'))
1354[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1355
1356>>> list(roundrobin('abc', 'd', 'ef'))
1357['a', 'd', 'e', 'b', 'f', 'c']
1358
Raymond Hettinger5cd012b2009-01-26 02:17:52 +00001359>>> list(powerset([1,2,3]))
1360[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001361
Raymond Hettingere8b4b602008-03-11 00:19:07 +00001362>>> list(compress('abcdef', [1,0,1,0,1,1]))
1363['a', 'c', 'e', 'f']
1364
Raymond Hettinger33691672008-07-19 00:43:00 +00001365>>> list(combinations_with_replacement('abc', 2))
1366[('a', 'a'), ('a', 'b'), ('a', 'c'), ('b', 'b'), ('b', 'c'), ('c', 'c')]
1367
Raymond Hettinger825758c2009-01-08 05:20:19 +00001368>>> list(combinations_with_replacement('01', 3))
1369[('0', '0', '0'), ('0', '0', '1'), ('0', '1', '1'), ('1', '1', '1')]
1370
1371>>> def combinations_with_replacement2(iterable, r):
1372... 'Alternate version that filters from product()'
1373... pool = tuple(iterable)
1374... n = len(pool)
1375... for indices in product(range(n), repeat=r):
1376... if sorted(indices) == list(indices):
1377... yield tuple(pool[i] for i in indices)
1378
1379>>> list(combinations_with_replacement('abc', 2)) == list(combinations_with_replacement2('abc', 2))
1380True
1381
1382>>> list(combinations_with_replacement('01', 3)) == list(combinations_with_replacement2('01', 3))
1383True
1384
1385>>> list(combinations_with_replacement('2310', 6)) == list(combinations_with_replacement2('2310', 6))
1386True
1387
Raymond Hettingerd8461652009-01-02 21:20:38 +00001388>>> list(unique_everseen('AAAABBBCCDAABBB'))
1389['A', 'B', 'C', 'D']
1390
1391>>> list(unique_everseen('ABBCcAD', str.lower))
1392['A', 'B', 'C', 'D']
1393
1394>>> list(unique_justseen('AAAABBBCCDAABBB'))
1395['A', 'B', 'C', 'D', 'A', 'B']
1396
1397>>> list(unique_justseen('ABBCcAD', str.lower))
1398['A', 'B', 'C', 'A', 'D']
1399
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001400"""
1401
1402__test__ = {'libreftest' : libreftest}
1403
1404def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001405 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Georg Brandlb84c1372007-01-21 10:28:43 +00001406 RegressionTests, LengthTransparency,
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001407 SubclassWithKwargsTest, TestExamples)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001408 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001409
1410 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001411 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001412 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001413 counts = [None] * 5
1414 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001415 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001416 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001417 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001418 print counts
1419
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001420 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001421 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001422
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001423if __name__ == "__main__":
1424 test_main(verbose=True)