blob: 7023b293bf047a003825fb45f9883304ea36f0e1 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001import unittest
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002from test import support
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003from 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
Christian Heimesc3f30c42008-02-22 16:37:40 +00008from functools import reduce
Benjamin Petersonee8712c2008-05-20 21:35:26 +00009maxsize = support.MAX_Py_ssize_t
Guido van Rossum360e4b82007-05-14 22:51:27 +000010minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000011
Guido van Rossum801f0d72006-08-24 19:48:10 +000012def lzip(*args):
13 return list(zip(*args))
14
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000015def onearg(x):
16 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000017 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000018
19def errfunc(*args):
20 'Test function that raises an error'
21 raise ValueError
22
23def gen3():
24 'Non-restartable source sequence'
25 for i in (0, 1, 2):
26 yield i
27
28def isEven(x):
29 'Test predicate'
30 return x%2==0
31
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000032def isOdd(x):
33 'Test predicate'
34 return x%2==1
35
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000036class StopNow:
37 'Class emulating an empty iterable.'
38 def __iter__(self):
39 return self
Georg Brandla18af4e2007-04-21 15:47:16 +000040 def __next__(self):
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000041 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000042
Raymond Hettinger02420702003-06-29 20:36:23 +000043def take(n, seq):
44 'Convenience function for partially consuming a long of infinite iterable'
45 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000046
Christian Heimes78644762008-03-04 23:39:23 +000047def prod(iterable):
48 return reduce(operator.mul, iterable, 1)
49
Christian Heimes380f7f22008-02-28 11:19:05 +000050def fact(n):
51 'Factorial'
Christian Heimes78644762008-03-04 23:39:23 +000052 return prod(range(1, n+1))
53
Raymond Hettinger96ef8112003-02-01 00:10:11 +000054class TestBasicOps(unittest.TestCase):
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000055 def test_chain(self):
Christian Heimesdd15f6c2008-03-16 00:07:10 +000056
57 def chain2(*iterables):
58 'Pure python version in the docs'
59 for it in iterables:
60 for element in it:
61 yield element
62
63 for c in (chain, chain2):
64 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
65 self.assertEqual(list(c('abc')), list('abc'))
66 self.assertEqual(list(c('')), [])
67 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
68 self.assertRaises(TypeError, list,c(2, 3))
Christian Heimesf16baeb2008-02-29 14:57:44 +000069
70 def test_chain_from_iterable(self):
71 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
72 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
73 self.assertEqual(list(chain.from_iterable([''])), [])
74 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
75 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000076
Christian Heimes380f7f22008-02-28 11:19:05 +000077 def test_combinations(self):
Raymond Hettinger5bad41e2009-01-08 21:01:54 +000078 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Christian Heimes380f7f22008-02-28 11:19:05 +000079 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Christian Heimes78644762008-03-04 23:39:23 +000080 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Christian Heimes380f7f22008-02-28 11:19:05 +000081 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Raymond Hettinger5bad41e2009-01-08 21:01:54 +000082 self.assertEqual(list(combinations('abc', 32)), []) # r > n
Christian Heimes380f7f22008-02-28 11:19:05 +000083 self.assertEqual(list(combinations(range(4), 3)),
84 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
Christian Heimes78644762008-03-04 23:39:23 +000085
86 def combinations1(iterable, r):
87 'Pure python version shown in the docs'
88 pool = tuple(iterable)
89 n = len(pool)
Raymond Hettinger5bad41e2009-01-08 21:01:54 +000090 if r > n:
91 return
Christian Heimes78644762008-03-04 23:39:23 +000092 indices = list(range(r))
93 yield tuple(pool[i] for i in indices)
94 while 1:
95 for i in reversed(range(r)):
96 if indices[i] != i + n - r:
97 break
98 else:
99 return
100 indices[i] += 1
101 for j in range(i+1, r):
102 indices[j] = indices[j-1] + 1
103 yield tuple(pool[i] for i in indices)
104
105 def combinations2(iterable, r):
106 'Pure python version shown in the docs'
107 pool = tuple(iterable)
108 n = len(pool)
109 for indices in permutations(range(n), r):
110 if sorted(indices) == list(indices):
111 yield tuple(pool[i] for i in indices)
112
113 for n in range(7):
Christian Heimes380f7f22008-02-28 11:19:05 +0000114 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000115 for r in range(n+2):
Christian Heimes380f7f22008-02-28 11:19:05 +0000116 result = list(combinations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000117 self.assertEqual(len(result), 0 if r>n else fact(n) / fact(r) / fact(n-r)) # right number of combs
Christian Heimes380f7f22008-02-28 11:19:05 +0000118 self.assertEqual(len(result), len(set(result))) # no repeats
119 self.assertEqual(result, sorted(result)) # lexicographic order
120 for c in result:
121 self.assertEqual(len(c), r) # r-length combinations
122 self.assertEqual(len(set(c)), r) # no duplicate elements
123 self.assertEqual(list(c), sorted(c)) # keep original ordering
124 self.assert_(all(e in values for e in c)) # elements taken from input iterable
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000125 self.assertEqual(list(c),
126 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Christian Heimes78644762008-03-04 23:39:23 +0000127 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000128 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000129
130 # Test implementation detail: tuple re-use
131 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
132 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
133
134 def test_permutations(self):
135 self.assertRaises(TypeError, permutations) # too few arguments
136 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000137 self.assertRaises(TypeError, permutations, None) # pool is not iterable
138 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000139 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000140 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Christian Heimes78644762008-03-04 23:39:23 +0000141 self.assertEqual(list(permutations(range(3), 2)),
142 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
143
144 def permutations1(iterable, r=None):
145 'Pure python version shown in the docs'
146 pool = tuple(iterable)
147 n = len(pool)
148 r = n if r is None else r
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000149 if r > n:
150 return
Christian Heimes78644762008-03-04 23:39:23 +0000151 indices = list(range(n))
152 cycles = list(range(n-r+1, n+1))[::-1]
153 yield tuple(pool[i] for i in indices[:r])
154 while n:
155 for i in reversed(range(r)):
156 cycles[i] -= 1
157 if cycles[i] == 0:
158 indices[i:] = indices[i+1:] + indices[i:i+1]
159 cycles[i] = n - i
160 else:
161 j = cycles[i]
162 indices[i], indices[-j] = indices[-j], indices[i]
163 yield tuple(pool[i] for i in indices[:r])
164 break
165 else:
166 return
167
168 def permutations2(iterable, r=None):
169 'Pure python version shown in the docs'
170 pool = tuple(iterable)
171 n = len(pool)
172 r = n if r is None else r
173 for indices in product(range(n), repeat=r):
174 if len(set(indices)) == r:
175 yield tuple(pool[i] for i in indices)
176
177 for n in range(7):
178 values = [5*x-12 for x in range(n)]
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000179 for r in range(n+2):
Christian Heimes78644762008-03-04 23:39:23 +0000180 result = list(permutations(values, r))
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000181 self.assertEqual(len(result), 0 if r>n else fact(n) / fact(n-r)) # right number of perms
Christian Heimes78644762008-03-04 23:39:23 +0000182 self.assertEqual(len(result), len(set(result))) # no repeats
183 self.assertEqual(result, sorted(result)) # lexicographic order
184 for p in result:
185 self.assertEqual(len(p), r) # r-length permutations
186 self.assertEqual(len(set(p)), r) # no duplicate elements
187 self.assert_(all(e in values for e in p)) # elements taken from input iterable
188 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000189 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Christian Heimes78644762008-03-04 23:39:23 +0000190 if r == n:
191 self.assertEqual(result, list(permutations(values, None))) # test r as None
192 self.assertEqual(result, list(permutations(values))) # test default r
193
194 # Test implementation detail: tuple re-use
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000195 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Christian Heimes78644762008-03-04 23:39:23 +0000196 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Christian Heimes380f7f22008-02-28 11:19:05 +0000197
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000198 def test_count(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000199 self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
200 self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
201 self.assertEqual(take(2, lzip('abc',count(3))), [('a', 3), ('b', 4)])
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000202 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
203 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000204 self.assertRaises(TypeError, count, 2, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000205 self.assertRaises(TypeError, count, 'a')
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000206 self.assertEqual(list(islice(count(maxsize-5), 10)),
207 list(range(maxsize-5, maxsize+5)))
208 self.assertEqual(list(islice(count(-maxsize-5), 10)),
209 list(range(-maxsize-5, -maxsize+5)))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000210 c = count(3)
211 self.assertEqual(repr(c), 'count(3)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000212 next(c)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000213 self.assertEqual(repr(c), 'count(4)')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000214 c = count(-9)
215 self.assertEqual(repr(c), 'count(-9)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000216 next(c)
217 self.assertEqual(next(c), -8)
Christian Heimesa37d4c62007-12-04 23:02:19 +0000218 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000219 # Test repr (ignoring the L in longs)
220 r1 = repr(count(i)).replace('L', '')
221 r2 = 'count(%r)'.__mod__(i).replace('L', '')
222 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000223
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000224 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000225 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000226 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000227 self.assertRaises(TypeError, cycle)
228 self.assertRaises(TypeError, cycle, 5)
229 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000230
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000231 def test_groupby(self):
232 # Check whether it accepts arguments correctly
233 self.assertEqual([], list(groupby([])))
234 self.assertEqual([], list(groupby([], key=id)))
235 self.assertRaises(TypeError, list, groupby('abc', []))
236 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000237 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000238
239 # Check normal input
240 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
241 (2,15,22), (3,16,23), (3,17,23)]
242 dup = []
243 for k, g in groupby(s, lambda r:r[0]):
244 for elem in g:
245 self.assertEqual(k, elem[0])
246 dup.append(elem)
247 self.assertEqual(s, dup)
248
249 # Check nested case
250 dup = []
251 for k, g in groupby(s, lambda r:r[0]):
252 for ik, ig in groupby(g, lambda r:r[2]):
253 for elem in ig:
254 self.assertEqual(k, elem[0])
255 self.assertEqual(ik, elem[2])
256 dup.append(elem)
257 self.assertEqual(s, dup)
258
259 # Check case where inner iterator is not used
260 keys = [k for k, g in groupby(s, lambda r:r[0])]
261 expectedkeys = set([r[0] for r in s])
262 self.assertEqual(set(keys), expectedkeys)
263 self.assertEqual(len(keys), len(expectedkeys))
264
265 # Exercise pipes and filters style
266 s = 'abracadabra'
267 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000268 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000269 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
270 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000271 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000272 self.assertEqual(r, ['a', 'b', 'r'])
273 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000274 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000275 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
276 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000277 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000278 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
279
Georg Brandla18af4e2007-04-21 15:47:16 +0000280 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000281 class ExpectedError(Exception):
282 pass
283 def delayed_raise(n=0):
284 for i in range(n):
285 yield 'yo'
286 raise ExpectedError
287 def gulp(iterable, keyp=None, func=list):
288 return [func(g) for k, g in groupby(iterable, keyp)]
289
Georg Brandla18af4e2007-04-21 15:47:16 +0000290 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000291 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000292 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000293 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
294
295 # __cmp__ failure
296 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000297 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000298 raise ExpectedError
299 s = [DummyCmp(), DummyCmp(), None]
300
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000301 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000302 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000303 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000304 self.assertRaises(ExpectedError, gulp, s)
305
306 # keyfunc failure
307 def keyfunc(obj):
308 if keyfunc.skip > 0:
309 keyfunc.skip -= 1
310 return obj
311 else:
312 raise ExpectedError
313
314 # keyfunc failure on outer object
315 keyfunc.skip = 0
316 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
317 keyfunc.skip = 1
318 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
319
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000320 def test_filter(self):
321 self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
322 self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])
323 self.assertEqual(list(filter(bool, [0,1,0,2,0])), [1,2])
324 self.assertEqual(take(4, filter(isEven, count())), [0,2,4,6])
325 self.assertRaises(TypeError, filter)
326 self.assertRaises(TypeError, filter, lambda x:x)
327 self.assertRaises(TypeError, filter, lambda x:x, range(6), 7)
328 self.assertRaises(TypeError, filter, isEven, 3)
329 self.assertRaises(TypeError, next, filter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000330
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000331 def test_filterfalse(self):
332 self.assertEqual(list(filterfalse(isEven, range(6))), [1,3,5])
333 self.assertEqual(list(filterfalse(None, [0,1,0,2,0])), [0,0,0])
334 self.assertEqual(list(filterfalse(bool, [0,1,0,2,0])), [0,0,0])
335 self.assertEqual(take(4, filterfalse(isEven, count())), [1,3,5,7])
336 self.assertRaises(TypeError, filterfalse)
337 self.assertRaises(TypeError, filterfalse, lambda x:x)
338 self.assertRaises(TypeError, filterfalse, lambda x:x, range(6), 7)
339 self.assertRaises(TypeError, filterfalse, isEven, 3)
340 self.assertRaises(TypeError, next, filterfalse(range(6), range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000341
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000342 def test_zip(self):
343 # XXX This is rather silly now that builtin zip() calls zip()...
344 ans = [(x,y) for x, y in zip('abc',count())]
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000345 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000346 self.assertEqual(list(zip('abc', range(6))), lzip('abc', range(6)))
347 self.assertEqual(list(zip('abcdef', range(3))), lzip('abcdef', range(3)))
348 self.assertEqual(take(3,zip('abcdef', count())), lzip('abcdef', range(3)))
349 self.assertEqual(list(zip('abcdef')), lzip('abcdef'))
350 self.assertEqual(list(zip()), lzip())
351 self.assertRaises(TypeError, zip, 3)
352 self.assertRaises(TypeError, zip, range(3), 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000353 # Check tuple re-use (implementation detail)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000354 self.assertEqual([tuple(list(pair)) for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000355 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000356 self.assertEqual([pair for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000357 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000358 ids = list(map(id, zip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000359 self.assertEqual(min(ids), max(ids))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000360 ids = list(map(id, list(zip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000361 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000362
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000363 def test_ziplongest(self):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000364 for args in [
365 ['abc', range(6)],
366 [range(6), 'abc'],
367 [range(1000), range(2000,2100), range(3000,3050)],
368 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
369 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
370 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000371 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
372 for i in range(max(map(len, args)))]
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000373 self.assertEqual(list(zip_longest(*args)), target)
374 self.assertEqual(list(zip_longest(*args, **{})), target)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000375 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000376 self.assertEqual(list(zip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000377
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000378 self.assertEqual(take(3,zip_longest('abcdef', count())), list(zip('abcdef', range(3)))) # take 3 from infinite input
Thomas Wouterscf297e42007-02-23 15:07:44 +0000379
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000380 self.assertEqual(list(zip_longest()), list(zip()))
381 self.assertEqual(list(zip_longest([])), list(zip([])))
382 self.assertEqual(list(zip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000383
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000384 self.assertEqual(list(zip_longest('abc', 'defg', **{})),
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000385 list(zip(list('abc')+[None], 'defg'))) # empty keyword dict
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000386 self.assertRaises(TypeError, zip_longest, 3)
387 self.assertRaises(TypeError, zip_longest, range(3), 3)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000388
389 for stmt in [
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000390 "zip_longest('abc', fv=1)",
391 "zip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000392 ]:
393 try:
394 eval(stmt, globals(), locals())
395 except TypeError:
396 pass
397 else:
398 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000399
Thomas Wouterscf297e42007-02-23 15:07:44 +0000400 # Check tuple re-use (implementation detail)
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000401 self.assertEqual([tuple(list(pair)) for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000402 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000403 self.assertEqual([pair for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000404 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000405 ids = list(map(id, zip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000406 self.assertEqual(min(ids), max(ids))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000407 ids = list(map(id, list(zip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000408 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
409
Christian Heimesc3f30c42008-02-22 16:37:40 +0000410 def test_product(self):
411 for args, result in [
Christian Heimes78644762008-03-04 23:39:23 +0000412 ([], [()]), # zero iterables
Christian Heimesc3f30c42008-02-22 16:37:40 +0000413 (['ab'], [('a',), ('b',)]), # one iterable
414 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
415 ([range(0), range(2), range(3)], []), # first iterable with zero length
416 ([range(2), range(0), range(3)], []), # middle iterable with zero length
417 ([range(2), range(3), range(0)], []), # last iterable with zero length
418 ]:
419 self.assertEqual(list(product(*args)), result)
Christian Heimesf16baeb2008-02-29 14:57:44 +0000420 for r in range(4):
421 self.assertEqual(list(product(*(args*r))),
422 list(product(*args, **dict(repeat=r))))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000423 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
424 self.assertRaises(TypeError, product, range(6), None)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000425
426 def product1(*args, **kwds):
427 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
428 n = len(pools)
429 if n == 0:
430 yield ()
431 return
432 if any(len(pool) == 0 for pool in pools):
433 return
434 indices = [0] * n
435 yield tuple(pool[i] for pool, i in zip(pools, indices))
436 while 1:
437 for i in reversed(range(n)): # right to left
438 if indices[i] == len(pools[i]) - 1:
439 continue
440 indices[i] += 1
441 for j in range(i+1, n):
442 indices[j] = 0
443 yield tuple(pool[i] for pool, i in zip(pools, indices))
444 break
445 else:
446 return
447
448 def product2(*args, **kwds):
449 'Pure python version used in docs'
450 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
451 result = [[]]
452 for pool in pools:
453 result = [x+[y] for x in result for y in pool]
454 for prod in result:
455 yield tuple(prod)
456
Christian Heimesc3f30c42008-02-22 16:37:40 +0000457 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
458 set('abcdefg'), range(11), tuple(range(13))]
459 for i in range(100):
460 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Christian Heimes78644762008-03-04 23:39:23 +0000461 expected_len = prod(map(len, args))
462 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000463 self.assertEqual(list(product(*args)), list(product1(*args)))
464 self.assertEqual(list(product(*args)), list(product2(*args)))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000465 args = map(iter, args)
Christian Heimes78644762008-03-04 23:39:23 +0000466 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000467
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000468 # Test implementation detail: tuple re-use
469 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
470 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000471
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000472 def test_repeat(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000473 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000474 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000475 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000476 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000477 self.assertEqual(list(repeat('a', 0)), [])
478 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000479 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000480 self.assertRaises(TypeError, repeat, None, 3, 4)
481 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000482 r = repeat(1+0j)
483 self.assertEqual(repr(r), 'repeat((1+0j))')
484 r = repeat(1+0j, 5)
485 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
486 list(r)
487 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000488
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000489 def test_map(self):
490 self.assertEqual(list(map(operator.pow, range(3), range(1,7))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000491 [0**1, 1**2, 2**3])
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000492 def tupleize(*args):
493 return args
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000494 self.assertEqual(list(map(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000495 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000496 self.assertEqual(list(map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000497 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000498 self.assertEqual(take(2,map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000499 [('a',0),('b',1)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000500 self.assertEqual(list(map(operator.pow, [])), [])
501 self.assertRaises(TypeError, map)
502 self.assertRaises(TypeError, list, map(None, range(3), range(3)))
503 self.assertRaises(TypeError, map, operator.neg)
504 self.assertRaises(TypeError, next, map(10, range(5)))
505 self.assertRaises(ValueError, next, map(errfunc, [4], [5]))
506 self.assertRaises(TypeError, next, map(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000507
508 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000509 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
510 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000511 self.assertEqual(take(3, starmap(operator.pow, zip(count(), count(1)))),
Raymond Hettinger02420702003-06-29 20:36:23 +0000512 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000513 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +0000514 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
515 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000516 self.assertRaises(TypeError, starmap)
517 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000518 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
519 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
520 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000521
522 def test_islice(self):
523 for args in [ # islice(args) should agree with range(args)
524 (10, 20, 3),
525 (10, 3, 20),
526 (10, 20),
527 (10, 3),
528 (20,)
529 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000530 self.assertEqual(list(islice(range(100), *args)),
531 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000532
533 for args, tgtargs in [ # Stop when seqn is exhausted
534 ((10, 110, 3), ((10, 100, 3))),
535 ((10, 110), ((10, 100))),
536 ((110,), (100,))
537 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000538 self.assertEqual(list(islice(range(100), *args)),
539 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000540
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000541 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +0000542 self.assertEqual(list(islice(range(10), None)), list(range(10)))
543 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
544 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
545 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
546 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000547
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000548 # Test number of items consumed SF #1171417
549 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +0000550 self.assertEqual(list(islice(it, 3)), list(range(3)))
551 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000552
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000553 # Test invalid arguments
Guido van Rossum805365e2007-05-07 22:24:25 +0000554 self.assertRaises(TypeError, islice, range(10))
555 self.assertRaises(TypeError, islice, range(10), 1, 2, 3, 4)
556 self.assertRaises(ValueError, islice, range(10), -5, 10, 1)
557 self.assertRaises(ValueError, islice, range(10), 1, -5, -1)
558 self.assertRaises(ValueError, islice, range(10), 1, 10, -1)
559 self.assertRaises(ValueError, islice, range(10), 1, 10, 0)
560 self.assertRaises(ValueError, islice, range(10), 'a')
561 self.assertRaises(ValueError, islice, range(10), 'a', 1)
562 self.assertRaises(ValueError, islice, range(10), 1, 'a')
563 self.assertRaises(ValueError, islice, range(10), 'a', 1, 1)
564 self.assertRaises(ValueError, islice, range(10), 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +0000565 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000566
567 def test_takewhile(self):
568 data = [1, 3, 5, 20, 2, 4, 6, 8]
569 underten = lambda x: x<10
570 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000571 self.assertEqual(list(takewhile(underten, [])), [])
572 self.assertRaises(TypeError, takewhile)
573 self.assertRaises(TypeError, takewhile, operator.pow)
574 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000575 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
576 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000577 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
578 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +0000579 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000580
581 def test_dropwhile(self):
582 data = [1, 3, 5, 20, 2, 4, 6, 8]
583 underten = lambda x: x<10
584 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000585 self.assertEqual(list(dropwhile(underten, [])), [])
586 self.assertRaises(TypeError, dropwhile)
587 self.assertRaises(TypeError, dropwhile, operator.pow)
588 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000589 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
590 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000591
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000592 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000593 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000594 def irange(n):
Guido van Rossum805365e2007-05-07 22:24:25 +0000595 for i in range(n):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000596 yield i
597
598 a, b = tee([]) # test empty iterator
599 self.assertEqual(list(a), [])
600 self.assertEqual(list(b), [])
601
602 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +0000603 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000604
605 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +0000606 self.assertEqual(list(a), list(range(n)))
607 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000608
609 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000610 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000611 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000612 del a
Guido van Rossum805365e2007-05-07 22:24:25 +0000613 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000614
615 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000616 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000617 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000618 del b
Guido van Rossum805365e2007-05-07 22:24:25 +0000619 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000620
Guido van Rossum805365e2007-05-07 22:24:25 +0000621 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000622 order = [0]*n + [1]*n
623 random.shuffle(order)
624 lists = ([], [])
625 its = tee(irange(n))
626 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +0000627 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000628 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +0000629 self.assertEqual(lists[0], list(range(n)))
630 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000631
Raymond Hettingerad983e72003-11-12 14:32:26 +0000632 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000633 self.assertRaises(TypeError, tee)
634 self.assertRaises(TypeError, tee, 3)
635 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000636 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000637
Raymond Hettingerad983e72003-11-12 14:32:26 +0000638 # tee object should be instantiable
639 a, b = tee('abc')
640 c = type(a)('def')
641 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000642
Raymond Hettingerad983e72003-11-12 14:32:26 +0000643 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +0000644 a, b, c = tee(range(2000), 3)
645 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000646 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +0000647 self.assertEqual(list(b), list(range(2000)))
648 self.assertEqual([next(c), next(c)], list(range(2)))
649 self.assertEqual(list(a), list(range(100,2000)))
650 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000651
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000652 # test values of n
653 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000654 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +0000655 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000656 result = tee('abc', n)
657 self.assertEqual(type(result), tuple)
658 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000659 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000660
Raymond Hettingerad983e72003-11-12 14:32:26 +0000661 # tee pass-through to copyable iterator
662 a, b = tee('abc')
663 c, d = tee(a)
664 self.assert_(a is c)
665
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000666 # test tee_new
667 t1, t2 = tee('abc')
668 tnew = type(t1)
669 self.assertRaises(TypeError, tnew)
670 self.assertRaises(TypeError, tnew, 10)
671 t3 = tnew(t1)
672 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000673
Raymond Hettingera9f60922004-10-17 16:40:14 +0000674 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +0000675 a, b = tee(range(10))
Raymond Hettingera9f60922004-10-17 16:40:14 +0000676 p = proxy(a)
677 self.assertEqual(getattr(p, '__class__'), type(b))
678 del a
679 self.assertRaises(ReferenceError, getattr, p, '__class__')
680
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000681 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000682 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000683
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000684 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +0000685 self.assertRaises(StopIteration, next, f([]))
686 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000687
Georg Brandla18af4e2007-04-21 15:47:16 +0000688 self.assertRaises(StopIteration, next, islice([], None))
689 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000690
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000691 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +0000692 self.assertRaises(StopIteration, next, p)
693 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000694 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +0000695 self.assertRaises(StopIteration, next, p)
696 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000697
Georg Brandla18af4e2007-04-21 15:47:16 +0000698 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000699
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000700 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +0000701 self.assertRaises(StopIteration, next, f(lambda x:x, []))
702 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000703
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000704class TestExamples(unittest.TestCase):
705
706 def test_chain(self):
707 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
708
709 def test_chain_from_iterable(self):
710 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
711
712 def test_combinations(self):
713 self.assertEqual(list(combinations('ABCD', 2)),
714 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
715 self.assertEqual(list(combinations(range(4), 3)),
716 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
717
718 def test_count(self):
719 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
720
721 def test_cycle(self):
722 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
723
724 def test_dropwhile(self):
725 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
726
727 def test_groupby(self):
728 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
729 list('ABCDAB'))
730 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
731 [list('AAAA'), list('BBB'), list('CC'), list('D')])
732
733 def test_filter(self):
734 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
735
736 def test_filterfalse(self):
737 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
738
739 def test_map(self):
740 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
741
742 def test_islice(self):
743 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
744 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
745 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
746 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
747
748 def test_zip(self):
749 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
750
751 def test_zip_longest(self):
752 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
753 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
754
755 def test_permutations(self):
756 self.assertEqual(list(permutations('ABCD', 2)),
757 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
758 self.assertEqual(list(permutations(range(3))),
759 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
760
761 def test_product(self):
762 self.assertEqual(list(product('ABCD', 'xy')),
763 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
764 self.assertEqual(list(product(range(2), repeat=3)),
765 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
766 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
767
768 def test_repeat(self):
769 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
770
771 def test_stapmap(self):
772 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
773 [32, 9, 1000])
774
775 def test_takewhile(self):
776 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
777
778
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000779class TestGC(unittest.TestCase):
780
781 def makecycle(self, iterator, container):
782 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +0000783 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000784 del container, iterator
785
786 def test_chain(self):
787 a = []
788 self.makecycle(chain(a), a)
789
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000790 def test_chain_from_iterable(self):
791 a = []
792 self.makecycle(chain.from_iterable([a]), a)
793
794 def test_combinations(self):
795 a = []
796 self.makecycle(combinations([1,2,a,3], 3), a)
797
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000798 def test_cycle(self):
799 a = []
800 self.makecycle(cycle([a]*2), a)
801
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000802 def test_dropwhile(self):
803 a = []
804 self.makecycle(dropwhile(bool, [0, a, a]), a)
805
806 def test_groupby(self):
807 a = []
808 self.makecycle(groupby([a]*2, lambda x:x), a)
809
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000810 def test_issue2246(self):
811 # Issue 2246 -- the _grouper iterator was not included in GC
812 n = 10
813 keyfunc = lambda x: x
814 for i, j in groupby(range(n), key=keyfunc):
815 keyfunc.__dict__.setdefault('x',[]).append(j)
816
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000817 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000818 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000819 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000820
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000821 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000822 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000823 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000824
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000825 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000826 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000827 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000828
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000829 def test_zip_longest(self):
830 a = []
831 self.makecycle(zip_longest([a]*2, [a]*3), a)
832 b = [a, None]
833 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
834
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000835 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000836 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000837 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000838
839 def test_islice(self):
840 a = []
841 self.makecycle(islice([a]*2, None), a)
842
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000843 def test_permutations(self):
844 a = []
845 self.makecycle(permutations([1,2,a,3], 3), a)
846
847 def test_product(self):
848 a = []
849 self.makecycle(product([1,2,a,3], repeat=3), a)
850
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000851 def test_repeat(self):
852 a = []
853 self.makecycle(repeat(a), a)
854
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000855 def test_starmap(self):
856 a = []
857 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
858
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000859 def test_takewhile(self):
860 a = []
861 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
862
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000863def R(seqn):
864 'Regular generator'
865 for i in seqn:
866 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000867
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000868class G:
869 'Sequence using __getitem__'
870 def __init__(self, seqn):
871 self.seqn = seqn
872 def __getitem__(self, i):
873 return self.seqn[i]
874
875class I:
876 'Sequence using iterator protocol'
877 def __init__(self, seqn):
878 self.seqn = seqn
879 self.i = 0
880 def __iter__(self):
881 return self
Georg Brandla18af4e2007-04-21 15:47:16 +0000882 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000883 if self.i >= len(self.seqn): raise StopIteration
884 v = self.seqn[self.i]
885 self.i += 1
886 return v
887
888class Ig:
889 'Sequence using iterator protocol defined with a generator'
890 def __init__(self, seqn):
891 self.seqn = seqn
892 self.i = 0
893 def __iter__(self):
894 for val in self.seqn:
895 yield val
896
897class X:
898 'Missing __getitem__ and __iter__'
899 def __init__(self, seqn):
900 self.seqn = seqn
901 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +0000902 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000903 if self.i >= len(self.seqn): raise StopIteration
904 v = self.seqn[self.i]
905 self.i += 1
906 return v
907
908class N:
Georg Brandla18af4e2007-04-21 15:47:16 +0000909 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000910 def __init__(self, seqn):
911 self.seqn = seqn
912 self.i = 0
913 def __iter__(self):
914 return self
915
916class E:
917 'Test propagation of exceptions'
918 def __init__(self, seqn):
919 self.seqn = seqn
920 self.i = 0
921 def __iter__(self):
922 return self
Georg Brandla18af4e2007-04-21 15:47:16 +0000923 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +0000924 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000925
926class S:
927 'Test immediate stop'
928 def __init__(self, seqn):
929 pass
930 def __iter__(self):
931 return self
Georg Brandla18af4e2007-04-21 15:47:16 +0000932 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000933 raise StopIteration
934
935def L(seqn):
936 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000937 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000938
939
940class TestVariousIteratorArgs(unittest.TestCase):
941
942 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000943 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000944 for g in (G, I, Ig, S, L, R):
945 self.assertEqual(list(chain(g(s))), list(g(s)))
946 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +0000947 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +0000948 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000949 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
950
Christian Heimesc3f30c42008-02-22 16:37:40 +0000951 def test_product(self):
952 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
953 self.assertRaises(TypeError, product, X(s))
954 self.assertRaises(TypeError, product, N(s))
955 self.assertRaises(ZeroDivisionError, product, E(s))
956
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000957 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000958 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000959 for g in (G, I, Ig, S, L, R):
960 tgtlen = len(s) * 3
961 expected = list(g(s))*3
962 actual = list(islice(cycle(g(s)), tgtlen))
963 self.assertEqual(actual, expected)
964 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000965 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000966 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
967
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000968 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000969 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000970 for g in (G, I, Ig, S, L, R):
971 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
972 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000973 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000974 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
975
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000976 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000977 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000978 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000979 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000980 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000981 self.assertRaises(TypeError, filter, isEven, X(s))
982 self.assertRaises(TypeError, filter, isEven, N(s))
983 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000984
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000985 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000986 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000987 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000988 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000989 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000990 self.assertRaises(TypeError, filterfalse, isEven, X(s))
991 self.assertRaises(TypeError, filterfalse, isEven, N(s))
992 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000993
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000994 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000995 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000996 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000997 self.assertEqual(list(zip(g(s))), lzip(g(s)))
998 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
999 self.assertRaises(TypeError, zip, X(s))
1000 self.assertRaises(TypeError, zip, N(s))
1001 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001002
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001003 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001004 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001005 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001006 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1007 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1008 self.assertRaises(TypeError, zip_longest, X(s))
1009 self.assertRaises(TypeError, zip_longest, N(s))
1010 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001011
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001012 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001013 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001014 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001015 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001016 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001017 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001018 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001019 self.assertRaises(TypeError, map, onearg, X(s))
1020 self.assertRaises(TypeError, map, onearg, N(s))
1021 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001022
1023 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001024 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001025 for g in (G, I, Ig, S, L, R):
1026 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1027 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001028 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001029 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1030
1031 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001032 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001033 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001034 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001035 self.assertEqual(list(starmap(operator.pow, g(ss))),
1036 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001037 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001038 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001039 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1040
1041 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001042 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001043 for g in (G, I, Ig, S, L, R):
1044 tgt = []
1045 for elem in g(s):
1046 if not isEven(elem): break
1047 tgt.append(elem)
1048 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1049 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001050 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001051 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1052
1053 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001054 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001055 for g in (G, I, Ig, S, L, R):
1056 tgt = []
1057 for elem in g(s):
1058 if not tgt and isOdd(elem): continue
1059 tgt.append(elem)
1060 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1061 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001062 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001063 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1064
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001065 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001066 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001067 for g in (G, I, Ig, S, L, R):
1068 it1, it2 = tee(g(s))
1069 self.assertEqual(list(it1), list(g(s)))
1070 self.assertEqual(list(it2), list(g(s)))
1071 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001072 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001073 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1074
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001075class LengthTransparency(unittest.TestCase):
1076
1077 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001078 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001079 self.assertEqual(len(repeat(None, 50)), 50)
1080 self.assertRaises(TypeError, len, repeat(None))
1081
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001082class RegressionTests(unittest.TestCase):
1083
1084 def test_sf_793826(self):
1085 # Fix Armin Rigo's successful efforts to wreak havoc
1086
1087 def mutatingtuple(tuple1, f, tuple2):
1088 # this builds a tuple t which is a copy of tuple1,
1089 # then calls f(t), then mutates t to be equal to tuple2
1090 # (needs len(tuple1) == len(tuple2)).
1091 def g(value, first=[1]):
1092 if first:
1093 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00001094 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001095 return value
1096 items = list(tuple2)
1097 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001098 gen = map(g, items)
1099 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00001100 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001101
1102 def f(t):
1103 global T
1104 T = t
1105 first[:] = list(T)
1106
1107 first = []
1108 mutatingtuple((1,2,3), f, (4,5,6))
1109 second = list(T)
1110 self.assertEqual(first, second)
1111
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001112
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001113 def test_sf_950057(self):
1114 # Make sure that chain() and cycle() catch exceptions immediately
1115 # rather than when shifting between input sources
1116
1117 def gen1():
1118 hist.append(0)
1119 yield 1
1120 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001121 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001122 hist.append(2)
1123
1124 def gen2(x):
1125 hist.append(3)
1126 yield 2
1127 hist.append(4)
1128 if x:
1129 raise StopIteration
1130
1131 hist = []
1132 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1133 self.assertEqual(hist, [0,1])
1134
1135 hist = []
1136 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1137 self.assertEqual(hist, [0,1])
1138
1139 hist = []
1140 self.assertRaises(AssertionError, list, cycle(gen1()))
1141 self.assertEqual(hist, [0,1])
1142
Thomas Woutersb2137042007-02-01 18:02:27 +00001143class SubclassWithKwargsTest(unittest.TestCase):
1144 def test_keywords_in_subclass(self):
1145 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001146 for cls in (repeat, zip, filter, filterfalse, chain, map,
Thomas Woutersb2137042007-02-01 18:02:27 +00001147 starmap, islice, takewhile, dropwhile, cycle):
1148 class Subclass(cls):
1149 def __init__(self, newarg=None, *args):
1150 cls.__init__(self, *args)
1151 try:
1152 Subclass(newarg=1)
1153 except TypeError as err:
1154 # we expect type errors because of wrong argument count
1155 self.failIf("does not take keyword arguments" in err.args[0])
1156
1157
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001158libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001159
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001160
1161>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001162>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00001163... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001164...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001165Check 1200 is for $120.15
1166Check 1201 is for $764.05
1167Check 1202 is for $823.14
1168
1169>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001170>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00001171... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00001172...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000011731
11748
117527
1176
1177>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001178>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00001179... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00001180...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001181Alex
1182Laura
1183Martin
1184Walter
1185Samuele
1186
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001187>>> from operator import itemgetter
1188>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00001189>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001190>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001191... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001192...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000011931 ['a', 'c', 'e']
11942 ['b', 'd', 'f']
11953 ['g']
1196
Raymond Hettinger734fb572004-01-20 20:04:40 +00001197# Find runs of consecutive numbers using groupby. The key to the solution
1198# is differencing with a range so that consecutive numbers all appear in
1199# same group.
1200>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00001201>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001202... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001203...
Raymond Hettinger734fb572004-01-20 20:04:40 +00001204[1]
1205[4, 5, 6]
1206[10]
1207[15, 16, 17, 18]
1208[22]
1209[25, 26, 27, 28]
1210
Georg Brandl3dbca812008-07-23 16:10:53 +00001211>>> def take(n, iterable):
1212... "Return first n items of the iterable as a list"
1213... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001214
Georg Brandl3dbca812008-07-23 16:10:53 +00001215>>> def enumerate(iterable, start=0):
1216... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001217
Georg Brandl3dbca812008-07-23 16:10:53 +00001218>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001219... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00001220... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001221
1222>>> def nth(iterable, n):
Georg Brandl3dbca812008-07-23 16:10:53 +00001223... "Returns the nth item or empty list"
Raymond Hettinger60eca932003-02-09 06:40:58 +00001224... return list(islice(iterable, n, n+1))
1225
Georg Brandl3dbca812008-07-23 16:10:53 +00001226>>> def quantify(iterable, pred=bool):
1227... "Count how many times the predicate is true"
1228... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001229
Georg Brandl3dbca812008-07-23 16:10:53 +00001230>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001231... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00001232... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001233
Georg Brandl3dbca812008-07-23 16:10:53 +00001234>>> def ncycles(iterable, n):
1235... "Returns the seqeuence elements n times"
1236... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001237
1238>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001239... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001240
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001241>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001242... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001243
1244>>> def repeatfunc(func, times=None, *args):
1245... "Repeat calls to func with specified arguments."
1246... " Example: repeatfunc(random.random)"
1247... if times is None:
1248... return starmap(func, repeat(args))
1249... else:
1250... return starmap(func, repeat(args, times))
1251
Raymond Hettingerd591f662003-10-26 15:34:50 +00001252>>> def pairwise(iterable):
1253... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1254... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001255... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00001256... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001257... except StopIteration:
1258... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001259... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001260
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001261>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001262... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001263... args = [iter(iterable)] * n
Benjamin Petersond18de0e2008-07-31 20:21:46 +00001264... return zip_longest(fillvalue=fillvalue, *args)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001265
1266>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001267... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001268... # Recipe credited to George Sakkis
1269... pending = len(iterables)
1270... nexts = cycle(iter(it).__next__ for it in iterables)
1271... while pending:
1272... try:
1273... for next in nexts:
1274... yield next()
1275... except StopIteration:
1276... pending -= 1
1277... nexts = cycle(islice(nexts, pending))
1278
1279>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00001280... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1281... s = list(iterable)
1282... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001283
1284>>> def compress(data, selectors):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001285... "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
1286... return (d for d, s in zip(data, selectors) if s)
Georg Brandl3dbca812008-07-23 16:10:53 +00001287
1288>>> def combinations_with_replacement(iterable, r):
1289... "combinations_with_replacement('ABC', 3) --> AA AB AC BB BC CC"
1290... pool = tuple(iterable)
1291... n = len(pool)
1292... indices = [0] * r
1293... yield tuple(pool[i] for i in indices)
1294... while 1:
1295... for i in reversed(range(r)):
1296... if indices[i] != n - 1:
1297... break
1298... else:
1299... return
1300... indices[i:] = [indices[i] + 1] * (r - i)
1301... yield tuple(pool[i] for i in indices)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001302
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001303>>> def unique_everseen(iterable, key=None):
1304... "List unique elements, preserving order. Remember all elements ever seen."
1305... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1306... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1307... seen = set()
1308... seen_add = seen.add
1309... if key is None:
1310... for element in iterable:
1311... if element not in seen:
1312... seen_add(element)
1313... yield element
1314... else:
1315... for element in iterable:
1316... k = key(element)
1317... if k not in seen:
1318... seen_add(k)
1319... yield element
1320
1321>>> def unique_justseen(iterable, key=None):
1322... "List unique elements, preserving order. Remember only the element just seen."
1323... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1324... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1325... return map(next, map(itemgetter(1), groupby(iterable, key)))
1326
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001327This is not part of the examples but it tests to make sure the definitions
1328perform as purported.
1329
Raymond Hettingera098b332003-09-08 23:58:40 +00001330>>> take(10, count())
1331[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1332
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001333>>> list(enumerate('abc'))
1334[(0, 'a'), (1, 'b'), (2, 'c')]
1335
1336>>> list(islice(tabulate(lambda x: 2*x), 4))
1337[0, 2, 4, 6]
1338
1339>>> nth('abcde', 3)
1340['d']
1341
Guido van Rossum805365e2007-05-07 22:24:25 +00001342>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000134350
1344
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001345>>> a = [[1, 2, 3], [4, 5, 6]]
1346>>> flatten(a)
1347[1, 2, 3, 4, 5, 6]
1348
1349>>> list(repeatfunc(pow, 5, 2, 3))
1350[8, 8, 8, 8, 8]
1351
1352>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001353>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001354[0, 0, 0, 0, 0]
1355
Raymond Hettingerd591f662003-10-26 15:34:50 +00001356>>> list(pairwise('abcd'))
1357[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001358
Raymond Hettingerd591f662003-10-26 15:34:50 +00001359>>> list(pairwise([]))
1360[]
1361
1362>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001363[]
1364
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001365>>> list(islice(padnone('abc'), 0, 6))
1366['a', 'b', 'c', None, None, None]
1367
1368>>> list(ncycles('abc', 3))
1369['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1370
1371>>> dotproduct([1,2,3], [4,5,6])
137232
1373
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001374>>> list(grouper(3, 'abcdefg', 'x'))
1375[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1376
1377>>> list(roundrobin('abc', 'd', 'ef'))
1378['a', 'd', 'e', 'b', 'f', 'c']
1379
Raymond Hettingerace67332009-01-26 02:23:50 +00001380>>> list(powerset([1,2,3]))
1381[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001382
1383>>> list(compress('abcdef', [1,0,1,0,1,1]))
1384['a', 'c', 'e', 'f']
1385
Georg Brandl3dbca812008-07-23 16:10:53 +00001386>>> list(combinations_with_replacement('abc', 2))
1387[('a', 'a'), ('a', 'b'), ('a', 'c'), ('b', 'b'), ('b', 'c'), ('c', 'c')]
1388
Raymond Hettinger5bad41e2009-01-08 21:01:54 +00001389>>> list(combinations_with_replacement('01', 3))
1390[('0', '0', '0'), ('0', '0', '1'), ('0', '1', '1'), ('1', '1', '1')]
1391
1392>>> def combinations_with_replacement2(iterable, r):
1393... 'Alternate version that filters from product()'
1394... pool = tuple(iterable)
1395... n = len(pool)
1396... for indices in product(range(n), repeat=r):
1397... if sorted(indices) == list(indices):
1398... yield tuple(pool[i] for i in indices)
1399
1400>>> list(combinations_with_replacement('abc', 2)) == list(combinations_with_replacement2('abc', 2))
1401True
1402
1403>>> list(combinations_with_replacement('01', 3)) == list(combinations_with_replacement2('01', 3))
1404True
1405
1406>>> list(combinations_with_replacement('2310', 6)) == list(combinations_with_replacement2('2310', 6))
1407True
1408
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001409>>> list(unique_everseen('AAAABBBCCDAABBB'))
1410['A', 'B', 'C', 'D']
1411
1412>>> list(unique_everseen('ABBCcAD', str.lower))
1413['A', 'B', 'C', 'D']
1414
1415>>> list(unique_justseen('AAAABBBCCDAABBB'))
1416['A', 'B', 'C', 'D', 'A', 'B']
1417
1418>>> list(unique_justseen('ABBCcAD', str.lower))
1419['A', 'B', 'C', 'A', 'D']
1420
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001421"""
1422
1423__test__ = {'libreftest' : libreftest}
1424
1425def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001426 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00001427 RegressionTests, LengthTransparency,
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001428 SubclassWithKwargsTest, TestExamples)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001429 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001430
1431 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001432 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001433 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001434 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00001435 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001436 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001437 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001438 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001439 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001440
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001441 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001442 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001443
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001444if __name__ == "__main__":
1445 test_main(verbose=True)