blob: 16789d8f9f3775ef244dd65db1f56bdba8352238 [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 Hettinger6b3b0fc2009-01-26 02:56:58 +0000198 def test_compress(self):
199 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
200 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
201 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
202 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
203 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
204 n = 10000
205 data = chain.from_iterable(repeat(range(6), n))
206 selectors = chain.from_iterable(repeat((0, 1)))
207 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
208 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
209 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
210 self.assertRaises(TypeError, compress, range(6)) # too few args
211 self.assertRaises(TypeError, compress, range(6), None) # too many args
212
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000213 def test_count(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000214 self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
215 self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
216 self.assertEqual(take(2, lzip('abc',count(3))), [('a', 3), ('b', 4)])
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000217 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
218 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000219 self.assertRaises(TypeError, count, 2, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000220 self.assertRaises(TypeError, count, 'a')
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000221 self.assertEqual(list(islice(count(maxsize-5), 10)),
222 list(range(maxsize-5, maxsize+5)))
223 self.assertEqual(list(islice(count(-maxsize-5), 10)),
224 list(range(-maxsize-5, -maxsize+5)))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000225 c = count(3)
226 self.assertEqual(repr(c), 'count(3)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000227 next(c)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000228 self.assertEqual(repr(c), 'count(4)')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000229 c = count(-9)
230 self.assertEqual(repr(c), 'count(-9)')
Georg Brandla18af4e2007-04-21 15:47:16 +0000231 next(c)
232 self.assertEqual(next(c), -8)
Christian Heimesa37d4c62007-12-04 23:02:19 +0000233 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 +0000234 # Test repr (ignoring the L in longs)
235 r1 = repr(count(i)).replace('L', '')
236 r2 = 'count(%r)'.__mod__(i).replace('L', '')
237 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000238
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000239 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000240 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000241 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000242 self.assertRaises(TypeError, cycle)
243 self.assertRaises(TypeError, cycle, 5)
244 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000245
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000246 def test_groupby(self):
247 # Check whether it accepts arguments correctly
248 self.assertEqual([], list(groupby([])))
249 self.assertEqual([], list(groupby([], key=id)))
250 self.assertRaises(TypeError, list, groupby('abc', []))
251 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000252 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000253
254 # Check normal input
255 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
256 (2,15,22), (3,16,23), (3,17,23)]
257 dup = []
258 for k, g in groupby(s, lambda r:r[0]):
259 for elem in g:
260 self.assertEqual(k, elem[0])
261 dup.append(elem)
262 self.assertEqual(s, dup)
263
264 # Check nested case
265 dup = []
266 for k, g in groupby(s, lambda r:r[0]):
267 for ik, ig in groupby(g, lambda r:r[2]):
268 for elem in ig:
269 self.assertEqual(k, elem[0])
270 self.assertEqual(ik, elem[2])
271 dup.append(elem)
272 self.assertEqual(s, dup)
273
274 # Check case where inner iterator is not used
275 keys = [k for k, g in groupby(s, lambda r:r[0])]
276 expectedkeys = set([r[0] for r in s])
277 self.assertEqual(set(keys), expectedkeys)
278 self.assertEqual(len(keys), len(expectedkeys))
279
280 # Exercise pipes and filters style
281 s = 'abracadabra'
282 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000283 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000284 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
285 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000286 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000287 self.assertEqual(r, ['a', 'b', 'r'])
288 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000289 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000290 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
291 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000292 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000293 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
294
Georg Brandla18af4e2007-04-21 15:47:16 +0000295 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000296 class ExpectedError(Exception):
297 pass
298 def delayed_raise(n=0):
299 for i in range(n):
300 yield 'yo'
301 raise ExpectedError
302 def gulp(iterable, keyp=None, func=list):
303 return [func(g) for k, g in groupby(iterable, keyp)]
304
Georg Brandla18af4e2007-04-21 15:47:16 +0000305 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000306 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000307 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000308 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
309
310 # __cmp__ failure
311 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000312 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000313 raise ExpectedError
314 s = [DummyCmp(), DummyCmp(), None]
315
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000316 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000317 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000318 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000319 self.assertRaises(ExpectedError, gulp, s)
320
321 # keyfunc failure
322 def keyfunc(obj):
323 if keyfunc.skip > 0:
324 keyfunc.skip -= 1
325 return obj
326 else:
327 raise ExpectedError
328
329 # keyfunc failure on outer object
330 keyfunc.skip = 0
331 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
332 keyfunc.skip = 1
333 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
334
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000335 def test_filter(self):
336 self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
337 self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])
338 self.assertEqual(list(filter(bool, [0,1,0,2,0])), [1,2])
339 self.assertEqual(take(4, filter(isEven, count())), [0,2,4,6])
340 self.assertRaises(TypeError, filter)
341 self.assertRaises(TypeError, filter, lambda x:x)
342 self.assertRaises(TypeError, filter, lambda x:x, range(6), 7)
343 self.assertRaises(TypeError, filter, isEven, 3)
344 self.assertRaises(TypeError, next, filter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000345
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000346 def test_filterfalse(self):
347 self.assertEqual(list(filterfalse(isEven, range(6))), [1,3,5])
348 self.assertEqual(list(filterfalse(None, [0,1,0,2,0])), [0,0,0])
349 self.assertEqual(list(filterfalse(bool, [0,1,0,2,0])), [0,0,0])
350 self.assertEqual(take(4, filterfalse(isEven, count())), [1,3,5,7])
351 self.assertRaises(TypeError, filterfalse)
352 self.assertRaises(TypeError, filterfalse, lambda x:x)
353 self.assertRaises(TypeError, filterfalse, lambda x:x, range(6), 7)
354 self.assertRaises(TypeError, filterfalse, isEven, 3)
355 self.assertRaises(TypeError, next, filterfalse(range(6), range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000356
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000357 def test_zip(self):
358 # XXX This is rather silly now that builtin zip() calls zip()...
359 ans = [(x,y) for x, y in zip('abc',count())]
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000360 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000361 self.assertEqual(list(zip('abc', range(6))), lzip('abc', range(6)))
362 self.assertEqual(list(zip('abcdef', range(3))), lzip('abcdef', range(3)))
363 self.assertEqual(take(3,zip('abcdef', count())), lzip('abcdef', range(3)))
364 self.assertEqual(list(zip('abcdef')), lzip('abcdef'))
365 self.assertEqual(list(zip()), lzip())
366 self.assertRaises(TypeError, zip, 3)
367 self.assertRaises(TypeError, zip, range(3), 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000368 # Check tuple re-use (implementation detail)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000369 self.assertEqual([tuple(list(pair)) for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000370 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000371 self.assertEqual([pair for pair in zip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000372 lzip('abc', 'def'))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000373 ids = list(map(id, zip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000374 self.assertEqual(min(ids), max(ids))
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000375 ids = list(map(id, list(zip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000376 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000377
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000378 def test_ziplongest(self):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000379 for args in [
380 ['abc', range(6)],
381 [range(6), 'abc'],
382 [range(1000), range(2000,2100), range(3000,3050)],
383 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
384 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
385 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000386 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
387 for i in range(max(map(len, args)))]
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000388 self.assertEqual(list(zip_longest(*args)), target)
389 self.assertEqual(list(zip_longest(*args, **{})), target)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000390 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 +0000391 self.assertEqual(list(zip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000392
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000393 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 +0000394
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000395 self.assertEqual(list(zip_longest()), list(zip()))
396 self.assertEqual(list(zip_longest([])), list(zip([])))
397 self.assertEqual(list(zip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000398
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000399 self.assertEqual(list(zip_longest('abc', 'defg', **{})),
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000400 list(zip(list('abc')+[None], 'defg'))) # empty keyword dict
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000401 self.assertRaises(TypeError, zip_longest, 3)
402 self.assertRaises(TypeError, zip_longest, range(3), 3)
Thomas Wouterscf297e42007-02-23 15:07:44 +0000403
404 for stmt in [
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000405 "zip_longest('abc', fv=1)",
406 "zip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000407 ]:
408 try:
409 eval(stmt, globals(), locals())
410 except TypeError:
411 pass
412 else:
413 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000414
Thomas Wouterscf297e42007-02-23 15:07:44 +0000415 # Check tuple re-use (implementation detail)
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000416 self.assertEqual([tuple(list(pair)) for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000417 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000418 self.assertEqual([pair for pair in zip_longest('abc', 'def')],
Thomas Wouterscf297e42007-02-23 15:07:44 +0000419 list(zip('abc', 'def')))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000420 ids = list(map(id, zip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000421 self.assertEqual(min(ids), max(ids))
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000422 ids = list(map(id, list(zip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000423 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
424
Christian Heimesc3f30c42008-02-22 16:37:40 +0000425 def test_product(self):
426 for args, result in [
Christian Heimes78644762008-03-04 23:39:23 +0000427 ([], [()]), # zero iterables
Christian Heimesc3f30c42008-02-22 16:37:40 +0000428 (['ab'], [('a',), ('b',)]), # one iterable
429 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
430 ([range(0), range(2), range(3)], []), # first iterable with zero length
431 ([range(2), range(0), range(3)], []), # middle iterable with zero length
432 ([range(2), range(3), range(0)], []), # last iterable with zero length
433 ]:
434 self.assertEqual(list(product(*args)), result)
Christian Heimesf16baeb2008-02-29 14:57:44 +0000435 for r in range(4):
436 self.assertEqual(list(product(*(args*r))),
437 list(product(*args, **dict(repeat=r))))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000438 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
439 self.assertRaises(TypeError, product, range(6), None)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000440
441 def product1(*args, **kwds):
442 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
443 n = len(pools)
444 if n == 0:
445 yield ()
446 return
447 if any(len(pool) == 0 for pool in pools):
448 return
449 indices = [0] * n
450 yield tuple(pool[i] for pool, i in zip(pools, indices))
451 while 1:
452 for i in reversed(range(n)): # right to left
453 if indices[i] == len(pools[i]) - 1:
454 continue
455 indices[i] += 1
456 for j in range(i+1, n):
457 indices[j] = 0
458 yield tuple(pool[i] for pool, i in zip(pools, indices))
459 break
460 else:
461 return
462
463 def product2(*args, **kwds):
464 'Pure python version used in docs'
465 pools = list(map(tuple, args)) * kwds.get('repeat', 1)
466 result = [[]]
467 for pool in pools:
468 result = [x+[y] for x in result for y in pool]
469 for prod in result:
470 yield tuple(prod)
471
Christian Heimesc3f30c42008-02-22 16:37:40 +0000472 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
473 set('abcdefg'), range(11), tuple(range(13))]
474 for i in range(100):
475 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Christian Heimes78644762008-03-04 23:39:23 +0000476 expected_len = prod(map(len, args))
477 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000478 self.assertEqual(list(product(*args)), list(product1(*args)))
479 self.assertEqual(list(product(*args)), list(product2(*args)))
Christian Heimesc3f30c42008-02-22 16:37:40 +0000480 args = map(iter, args)
Christian Heimes78644762008-03-04 23:39:23 +0000481 self.assertEqual(len(list(product(*args))), expected_len)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000482
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000483 # Test implementation detail: tuple re-use
484 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
485 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000486
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000487 def test_repeat(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000488 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000489 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000490 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000491 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000492 self.assertEqual(list(repeat('a', 0)), [])
493 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000494 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000495 self.assertRaises(TypeError, repeat, None, 3, 4)
496 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000497 r = repeat(1+0j)
498 self.assertEqual(repr(r), 'repeat((1+0j))')
499 r = repeat(1+0j, 5)
500 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
501 list(r)
502 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000503
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000504 def test_map(self):
505 self.assertEqual(list(map(operator.pow, range(3), range(1,7))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000506 [0**1, 1**2, 2**3])
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000507 def tupleize(*args):
508 return args
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000509 self.assertEqual(list(map(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000510 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000511 self.assertEqual(list(map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000512 [('a',0),('b',1),('c',2)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000513 self.assertEqual(take(2,map(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000514 [('a',0),('b',1)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000515 self.assertEqual(list(map(operator.pow, [])), [])
516 self.assertRaises(TypeError, map)
517 self.assertRaises(TypeError, list, map(None, range(3), range(3)))
518 self.assertRaises(TypeError, map, operator.neg)
519 self.assertRaises(TypeError, next, map(10, range(5)))
520 self.assertRaises(ValueError, next, map(errfunc, [4], [5]))
521 self.assertRaises(TypeError, next, map(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000522
523 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000524 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
525 [0**1, 1**2, 2**3])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000526 self.assertEqual(take(3, starmap(operator.pow, zip(count(), count(1)))),
Raymond Hettinger02420702003-06-29 20:36:23 +0000527 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000528 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +0000529 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
530 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000531 self.assertRaises(TypeError, starmap)
532 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000533 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
534 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
535 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000536
537 def test_islice(self):
538 for args in [ # islice(args) should agree with range(args)
539 (10, 20, 3),
540 (10, 3, 20),
541 (10, 20),
542 (10, 3),
543 (20,)
544 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000545 self.assertEqual(list(islice(range(100), *args)),
546 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000547
548 for args, tgtargs in [ # Stop when seqn is exhausted
549 ((10, 110, 3), ((10, 100, 3))),
550 ((10, 110), ((10, 100))),
551 ((110,), (100,))
552 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000553 self.assertEqual(list(islice(range(100), *args)),
554 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000555
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000556 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +0000557 self.assertEqual(list(islice(range(10), None)), list(range(10)))
558 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
559 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
560 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
561 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000562
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000563 # Test number of items consumed SF #1171417
564 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +0000565 self.assertEqual(list(islice(it, 3)), list(range(3)))
566 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000567
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000568 # Test invalid arguments
Guido van Rossum805365e2007-05-07 22:24:25 +0000569 self.assertRaises(TypeError, islice, range(10))
570 self.assertRaises(TypeError, islice, range(10), 1, 2, 3, 4)
571 self.assertRaises(ValueError, islice, range(10), -5, 10, 1)
572 self.assertRaises(ValueError, islice, range(10), 1, -5, -1)
573 self.assertRaises(ValueError, islice, range(10), 1, 10, -1)
574 self.assertRaises(ValueError, islice, range(10), 1, 10, 0)
575 self.assertRaises(ValueError, islice, range(10), 'a')
576 self.assertRaises(ValueError, islice, range(10), 'a', 1)
577 self.assertRaises(ValueError, islice, range(10), 1, 'a')
578 self.assertRaises(ValueError, islice, range(10), 'a', 1, 1)
579 self.assertRaises(ValueError, islice, range(10), 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +0000580 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000581
582 def test_takewhile(self):
583 data = [1, 3, 5, 20, 2, 4, 6, 8]
584 underten = lambda x: x<10
585 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000586 self.assertEqual(list(takewhile(underten, [])), [])
587 self.assertRaises(TypeError, takewhile)
588 self.assertRaises(TypeError, takewhile, operator.pow)
589 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000590 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
591 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000592 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
593 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +0000594 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000595
596 def test_dropwhile(self):
597 data = [1, 3, 5, 20, 2, 4, 6, 8]
598 underten = lambda x: x<10
599 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000600 self.assertEqual(list(dropwhile(underten, [])), [])
601 self.assertRaises(TypeError, dropwhile)
602 self.assertRaises(TypeError, dropwhile, operator.pow)
603 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000604 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
605 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000606
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000607 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000608 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000609 def irange(n):
Guido van Rossum805365e2007-05-07 22:24:25 +0000610 for i in range(n):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000611 yield i
612
613 a, b = tee([]) # test empty iterator
614 self.assertEqual(list(a), [])
615 self.assertEqual(list(b), [])
616
617 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +0000618 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000619
620 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +0000621 self.assertEqual(list(a), list(range(n)))
622 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000623
624 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000625 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000626 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000627 del a
Guido van Rossum805365e2007-05-07 22:24:25 +0000628 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000629
630 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000631 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000632 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000633 del b
Guido van Rossum805365e2007-05-07 22:24:25 +0000634 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000635
Guido van Rossum805365e2007-05-07 22:24:25 +0000636 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000637 order = [0]*n + [1]*n
638 random.shuffle(order)
639 lists = ([], [])
640 its = tee(irange(n))
641 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +0000642 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000643 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +0000644 self.assertEqual(lists[0], list(range(n)))
645 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000646
Raymond Hettingerad983e72003-11-12 14:32:26 +0000647 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000648 self.assertRaises(TypeError, tee)
649 self.assertRaises(TypeError, tee, 3)
650 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000651 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000652
Raymond Hettingerad983e72003-11-12 14:32:26 +0000653 # tee object should be instantiable
654 a, b = tee('abc')
655 c = type(a)('def')
656 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000657
Raymond Hettingerad983e72003-11-12 14:32:26 +0000658 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +0000659 a, b, c = tee(range(2000), 3)
660 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000661 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +0000662 self.assertEqual(list(b), list(range(2000)))
663 self.assertEqual([next(c), next(c)], list(range(2)))
664 self.assertEqual(list(a), list(range(100,2000)))
665 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000666
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000667 # test values of n
668 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000669 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +0000670 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000671 result = tee('abc', n)
672 self.assertEqual(type(result), tuple)
673 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000674 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000675
Raymond Hettingerad983e72003-11-12 14:32:26 +0000676 # tee pass-through to copyable iterator
677 a, b = tee('abc')
678 c, d = tee(a)
679 self.assert_(a is c)
680
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000681 # test tee_new
682 t1, t2 = tee('abc')
683 tnew = type(t1)
684 self.assertRaises(TypeError, tnew)
685 self.assertRaises(TypeError, tnew, 10)
686 t3 = tnew(t1)
687 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000688
Raymond Hettingera9f60922004-10-17 16:40:14 +0000689 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +0000690 a, b = tee(range(10))
Raymond Hettingera9f60922004-10-17 16:40:14 +0000691 p = proxy(a)
692 self.assertEqual(getattr(p, '__class__'), type(b))
693 del a
694 self.assertRaises(ReferenceError, getattr, p, '__class__')
695
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000696 def test_StopIteration(self):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000697 self.assertRaises(StopIteration, next, zip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000698
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000699 for f in (chain, cycle, zip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +0000700 self.assertRaises(StopIteration, next, f([]))
701 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000702
Georg Brandla18af4e2007-04-21 15:47:16 +0000703 self.assertRaises(StopIteration, next, islice([], None))
704 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000705
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000706 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +0000707 self.assertRaises(StopIteration, next, p)
708 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000709 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +0000710 self.assertRaises(StopIteration, next, p)
711 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000712
Georg Brandla18af4e2007-04-21 15:47:16 +0000713 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000714
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000715 for f in (filter, filterfalse, map, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +0000716 self.assertRaises(StopIteration, next, f(lambda x:x, []))
717 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000718
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000719class TestExamples(unittest.TestCase):
720
721 def test_chain(self):
722 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
723
724 def test_chain_from_iterable(self):
725 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
726
727 def test_combinations(self):
728 self.assertEqual(list(combinations('ABCD', 2)),
729 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
730 self.assertEqual(list(combinations(range(4), 3)),
731 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
732
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000733 def test_compress(self):
734 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
735
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000736 def test_count(self):
737 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
738
739 def test_cycle(self):
740 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
741
742 def test_dropwhile(self):
743 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
744
745 def test_groupby(self):
746 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
747 list('ABCDAB'))
748 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
749 [list('AAAA'), list('BBB'), list('CC'), list('D')])
750
751 def test_filter(self):
752 self.assertEqual(list(filter(lambda x: x%2, range(10))), [1,3,5,7,9])
753
754 def test_filterfalse(self):
755 self.assertEqual(list(filterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
756
757 def test_map(self):
758 self.assertEqual(list(map(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
759
760 def test_islice(self):
761 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
762 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
763 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
764 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
765
766 def test_zip(self):
767 self.assertEqual(list(zip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
768
769 def test_zip_longest(self):
770 self.assertEqual(list(zip_longest('ABCD', 'xy', fillvalue='-')),
771 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
772
773 def test_permutations(self):
774 self.assertEqual(list(permutations('ABCD', 2)),
775 list(map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split())))
776 self.assertEqual(list(permutations(range(3))),
777 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
778
779 def test_product(self):
780 self.assertEqual(list(product('ABCD', 'xy')),
781 list(map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split())))
782 self.assertEqual(list(product(range(2), repeat=3)),
783 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
784 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
785
786 def test_repeat(self):
787 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
788
789 def test_stapmap(self):
790 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
791 [32, 9, 1000])
792
793 def test_takewhile(self):
794 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
795
796
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000797class TestGC(unittest.TestCase):
798
799 def makecycle(self, iterator, container):
800 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +0000801 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000802 del container, iterator
803
804 def test_chain(self):
805 a = []
806 self.makecycle(chain(a), a)
807
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000808 def test_chain_from_iterable(self):
809 a = []
810 self.makecycle(chain.from_iterable([a]), a)
811
812 def test_combinations(self):
813 a = []
814 self.makecycle(combinations([1,2,a,3], 3), a)
815
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000816 def test_compress(self):
817 a = []
818 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
819
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000820 def test_cycle(self):
821 a = []
822 self.makecycle(cycle([a]*2), a)
823
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000824 def test_dropwhile(self):
825 a = []
826 self.makecycle(dropwhile(bool, [0, a, a]), a)
827
828 def test_groupby(self):
829 a = []
830 self.makecycle(groupby([a]*2, lambda x:x), a)
831
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000832 def test_issue2246(self):
833 # Issue 2246 -- the _grouper iterator was not included in GC
834 n = 10
835 keyfunc = lambda x: x
836 for i, j in groupby(range(n), key=keyfunc):
837 keyfunc.__dict__.setdefault('x',[]).append(j)
838
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000839 def test_filter(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000840 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000841 self.makecycle(filter(lambda x:True, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000842
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000843 def test_filterfalse(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000844 a = []
Raymond Hettingerb0002d22008-03-13 01:41:43 +0000845 self.makecycle(filterfalse(lambda x:False, a), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000846
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000847 def test_zip(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000848 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000849 self.makecycle(zip([a]*2, [a]*3), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000850
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000851 def test_zip_longest(self):
852 a = []
853 self.makecycle(zip_longest([a]*2, [a]*3), a)
854 b = [a, None]
855 self.makecycle(zip_longest([a]*2, [a]*3, fillvalue=b), a)
856
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000857 def test_map(self):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000858 a = []
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000859 self.makecycle(map(lambda x:x, [a]*2), a)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000860
861 def test_islice(self):
862 a = []
863 self.makecycle(islice([a]*2, None), a)
864
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000865 def test_permutations(self):
866 a = []
867 self.makecycle(permutations([1,2,a,3], 3), a)
868
869 def test_product(self):
870 a = []
871 self.makecycle(product([1,2,a,3], repeat=3), a)
872
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000873 def test_repeat(self):
874 a = []
875 self.makecycle(repeat(a), a)
876
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000877 def test_starmap(self):
878 a = []
879 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
880
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000881 def test_takewhile(self):
882 a = []
883 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
884
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000885def R(seqn):
886 'Regular generator'
887 for i in seqn:
888 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000889
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000890class G:
891 'Sequence using __getitem__'
892 def __init__(self, seqn):
893 self.seqn = seqn
894 def __getitem__(self, i):
895 return self.seqn[i]
896
897class I:
898 'Sequence using iterator protocol'
899 def __init__(self, seqn):
900 self.seqn = seqn
901 self.i = 0
902 def __iter__(self):
903 return self
Georg Brandla18af4e2007-04-21 15:47:16 +0000904 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000905 if self.i >= len(self.seqn): raise StopIteration
906 v = self.seqn[self.i]
907 self.i += 1
908 return v
909
910class Ig:
911 'Sequence using iterator protocol defined with a generator'
912 def __init__(self, seqn):
913 self.seqn = seqn
914 self.i = 0
915 def __iter__(self):
916 for val in self.seqn:
917 yield val
918
919class X:
920 'Missing __getitem__ and __iter__'
921 def __init__(self, seqn):
922 self.seqn = seqn
923 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +0000924 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000925 if self.i >= len(self.seqn): raise StopIteration
926 v = self.seqn[self.i]
927 self.i += 1
928 return v
929
930class N:
Georg Brandla18af4e2007-04-21 15:47:16 +0000931 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000932 def __init__(self, seqn):
933 self.seqn = seqn
934 self.i = 0
935 def __iter__(self):
936 return self
937
938class E:
939 'Test propagation of exceptions'
940 def __init__(self, seqn):
941 self.seqn = seqn
942 self.i = 0
943 def __iter__(self):
944 return self
Georg Brandla18af4e2007-04-21 15:47:16 +0000945 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +0000946 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000947
948class S:
949 'Test immediate stop'
950 def __init__(self, seqn):
951 pass
952 def __iter__(self):
953 return self
Georg Brandla18af4e2007-04-21 15:47:16 +0000954 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000955 raise StopIteration
956
957def L(seqn):
958 'Test multiple tiers of iterators'
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000959 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000960
961
962class TestVariousIteratorArgs(unittest.TestCase):
963
964 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000965 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000966 for g in (G, I, Ig, S, L, R):
967 self.assertEqual(list(chain(g(s))), list(g(s)))
968 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Christian Heimesf16baeb2008-02-29 14:57:44 +0000969 self.assertRaises(TypeError, list, chain(X(s)))
Mark Dickinson26a96fa2008-03-01 02:27:46 +0000970 self.assertRaises(TypeError, list, chain(N(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000971 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
972
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000973 def test_compress(self):
974 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
975 n = len(s)
976 for g in (G, I, Ig, S, L, R):
977 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
978 self.assertRaises(TypeError, compress, X(s), repeat(1))
979 self.assertRaises(TypeError, compress, N(s), repeat(1))
980 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
981
Christian Heimesc3f30c42008-02-22 16:37:40 +0000982 def test_product(self):
983 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
984 self.assertRaises(TypeError, product, X(s))
985 self.assertRaises(TypeError, product, N(s))
986 self.assertRaises(ZeroDivisionError, product, E(s))
987
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000988 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000989 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000990 for g in (G, I, Ig, S, L, R):
991 tgtlen = len(s) * 3
992 expected = list(g(s))*3
993 actual = list(islice(cycle(g(s)), tgtlen))
994 self.assertEqual(actual, expected)
995 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000996 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000997 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
998
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000999 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001000 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001001 for g in (G, I, Ig, S, L, R):
1002 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1003 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001004 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001005 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1006
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001007 def test_filter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001008 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001009 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001010 self.assertEqual(list(filter(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001011 [x for x in g(s) if isEven(x)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001012 self.assertRaises(TypeError, filter, isEven, X(s))
1013 self.assertRaises(TypeError, filter, isEven, N(s))
1014 self.assertRaises(ZeroDivisionError, list, filter(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001015
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001016 def test_filterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001017 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001018 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001019 self.assertEqual(list(filterfalse(isEven, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001020 [x for x in g(s) if isOdd(x)])
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001021 self.assertRaises(TypeError, filterfalse, isEven, X(s))
1022 self.assertRaises(TypeError, filterfalse, isEven, N(s))
1023 self.assertRaises(ZeroDivisionError, list, filterfalse(isEven, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001024
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001025 def test_zip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001026 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001027 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001028 self.assertEqual(list(zip(g(s))), lzip(g(s)))
1029 self.assertEqual(list(zip(g(s), g(s))), lzip(g(s), g(s)))
1030 self.assertRaises(TypeError, zip, X(s))
1031 self.assertRaises(TypeError, zip, N(s))
1032 self.assertRaises(ZeroDivisionError, list, zip(E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001033
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001034 def test_ziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001035 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +00001036 for g in (G, I, Ig, S, L, R):
Raymond Hettingerb0002d22008-03-13 01:41:43 +00001037 self.assertEqual(list(zip_longest(g(s))), list(zip(g(s))))
1038 self.assertEqual(list(zip_longest(g(s), g(s))), list(zip(g(s), g(s))))
1039 self.assertRaises(TypeError, zip_longest, X(s))
1040 self.assertRaises(TypeError, zip_longest, N(s))
1041 self.assertRaises(ZeroDivisionError, list, zip_longest(E(s)))
Thomas Wouterscf297e42007-02-23 15:07:44 +00001042
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001043 def test_map(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001044 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001045 for g in (G, I, Ig, S, L, R):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001046 self.assertEqual(list(map(onearg, g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001047 [onearg(x) for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001048 self.assertEqual(list(map(operator.pow, g(s), g(s))),
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001049 [x**x for x in g(s)])
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001050 self.assertRaises(TypeError, map, onearg, X(s))
1051 self.assertRaises(TypeError, map, onearg, N(s))
1052 self.assertRaises(ZeroDivisionError, list, map(onearg, E(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001053
1054 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001055 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001056 for g in (G, I, Ig, S, L, R):
1057 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1058 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001059 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001060 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1061
1062 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001063 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001064 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +00001065 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001066 self.assertEqual(list(starmap(operator.pow, g(ss))),
1067 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001068 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001069 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001070 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1071
1072 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001073 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001074 for g in (G, I, Ig, S, L, R):
1075 tgt = []
1076 for elem in g(s):
1077 if not isEven(elem): break
1078 tgt.append(elem)
1079 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1080 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001081 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001082 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1083
1084 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001085 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001086 for g in (G, I, Ig, S, L, R):
1087 tgt = []
1088 for elem in g(s):
1089 if not tgt and isOdd(elem): continue
1090 tgt.append(elem)
1091 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1092 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001093 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001094 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1095
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001096 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +00001097 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001098 for g in (G, I, Ig, S, L, R):
1099 it1, it2 = tee(g(s))
1100 self.assertEqual(list(it1), list(g(s)))
1101 self.assertEqual(list(it2), list(g(s)))
1102 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +00001103 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001104 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1105
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001106class LengthTransparency(unittest.TestCase):
1107
1108 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001109 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001110 self.assertEqual(len(repeat(None, 50)), 50)
1111 self.assertRaises(TypeError, len, repeat(None))
1112
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001113class RegressionTests(unittest.TestCase):
1114
1115 def test_sf_793826(self):
1116 # Fix Armin Rigo's successful efforts to wreak havoc
1117
1118 def mutatingtuple(tuple1, f, tuple2):
1119 # this builds a tuple t which is a copy of tuple1,
1120 # then calls f(t), then mutates t to be equal to tuple2
1121 # (needs len(tuple1) == len(tuple2)).
1122 def g(value, first=[1]):
1123 if first:
1124 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +00001125 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001126 return value
1127 items = list(tuple2)
1128 items[1:1] = list(tuple1)
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001129 gen = map(g, items)
1130 z = zip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +00001131 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001132
1133 def f(t):
1134 global T
1135 T = t
1136 first[:] = list(T)
1137
1138 first = []
1139 mutatingtuple((1,2,3), f, (4,5,6))
1140 second = list(T)
1141 self.assertEqual(first, second)
1142
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001143
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001144 def test_sf_950057(self):
1145 # Make sure that chain() and cycle() catch exceptions immediately
1146 # rather than when shifting between input sources
1147
1148 def gen1():
1149 hist.append(0)
1150 yield 1
1151 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001152 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001153 hist.append(2)
1154
1155 def gen2(x):
1156 hist.append(3)
1157 yield 2
1158 hist.append(4)
1159 if x:
1160 raise StopIteration
1161
1162 hist = []
1163 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1164 self.assertEqual(hist, [0,1])
1165
1166 hist = []
1167 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1168 self.assertEqual(hist, [0,1])
1169
1170 hist = []
1171 self.assertRaises(AssertionError, list, cycle(gen1()))
1172 self.assertEqual(hist, [0,1])
1173
Thomas Woutersb2137042007-02-01 18:02:27 +00001174class SubclassWithKwargsTest(unittest.TestCase):
1175 def test_keywords_in_subclass(self):
1176 # count is not subclassable...
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001177 for cls in (repeat, zip, filter, filterfalse, chain, map,
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00001178 starmap, islice, takewhile, dropwhile, cycle, compress):
Thomas Woutersb2137042007-02-01 18:02:27 +00001179 class Subclass(cls):
1180 def __init__(self, newarg=None, *args):
1181 cls.__init__(self, *args)
1182 try:
1183 Subclass(newarg=1)
1184 except TypeError as err:
1185 # we expect type errors because of wrong argument count
1186 self.failIf("does not take keyword arguments" in err.args[0])
1187
1188
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001189libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001190
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001191
1192>>> amounts = [120.15, 764.05, 823.14]
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001193>>> for checknum, amount in zip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +00001194... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001195...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001196Check 1200 is for $120.15
1197Check 1201 is for $764.05
1198Check 1202 is for $823.14
1199
1200>>> import operator
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001201>>> for cube in map(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +00001202... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +00001203...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000012041
12058
120627
1207
1208>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001209>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +00001210... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +00001211...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001212Alex
1213Laura
1214Martin
1215Walter
1216Samuele
1217
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001218>>> from operator import itemgetter
1219>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +00001220>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001221>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001222... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001223...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000012241 ['a', 'c', 'e']
12252 ['b', 'd', 'f']
12263 ['g']
1227
Raymond Hettinger734fb572004-01-20 20:04:40 +00001228# Find runs of consecutive numbers using groupby. The key to the solution
1229# is differencing with a range so that consecutive numbers all appear in
1230# same group.
1231>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +00001232>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +00001233... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +00001234...
Raymond Hettinger734fb572004-01-20 20:04:40 +00001235[1]
1236[4, 5, 6]
1237[10]
1238[15, 16, 17, 18]
1239[22]
1240[25, 26, 27, 28]
1241
Georg Brandl3dbca812008-07-23 16:10:53 +00001242>>> def take(n, iterable):
1243... "Return first n items of the iterable as a list"
1244... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001245
Georg Brandl3dbca812008-07-23 16:10:53 +00001246>>> def enumerate(iterable, start=0):
1247... return zip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001248
Georg Brandl3dbca812008-07-23 16:10:53 +00001249>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001250... "Return function(0), function(1), ..."
Georg Brandl3dbca812008-07-23 16:10:53 +00001251... return map(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001252
1253>>> def nth(iterable, n):
Georg Brandl3dbca812008-07-23 16:10:53 +00001254... "Returns the nth item or empty list"
Raymond Hettinger60eca932003-02-09 06:40:58 +00001255... return list(islice(iterable, n, n+1))
1256
Georg Brandl3dbca812008-07-23 16:10:53 +00001257>>> def quantify(iterable, pred=bool):
1258... "Count how many times the predicate is true"
1259... return sum(map(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001260
Georg Brandl3dbca812008-07-23 16:10:53 +00001261>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001262... "Returns the sequence elements and then returns None indefinitely"
Georg Brandl3dbca812008-07-23 16:10:53 +00001263... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001264
Georg Brandl3dbca812008-07-23 16:10:53 +00001265>>> def ncycles(iterable, n):
1266... "Returns the seqeuence elements n times"
1267... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001268
1269>>> def dotproduct(vec1, vec2):
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001270... return sum(map(operator.mul, vec1, vec2))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001271
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001272>>> def flatten(listOfLists):
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001273... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001274
1275>>> def repeatfunc(func, times=None, *args):
1276... "Repeat calls to func with specified arguments."
1277... " Example: repeatfunc(random.random)"
1278... if times is None:
1279... return starmap(func, repeat(args))
1280... else:
1281... return starmap(func, repeat(args, times))
1282
Raymond Hettingerd591f662003-10-26 15:34:50 +00001283>>> def pairwise(iterable):
1284... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1285... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001286... try:
Georg Brandla18af4e2007-04-21 15:47:16 +00001287... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001288... except StopIteration:
1289... pass
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001290... return zip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001291
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001292>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001293... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001294... args = [iter(iterable)] * n
Benjamin Petersond18de0e2008-07-31 20:21:46 +00001295... return zip_longest(fillvalue=fillvalue, *args)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001296
1297>>> def roundrobin(*iterables):
Raymond Hettingerf5a2e472008-07-30 07:37:37 +00001298... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001299... # Recipe credited to George Sakkis
1300... pending = len(iterables)
1301... nexts = cycle(iter(it).__next__ for it in iterables)
1302... while pending:
1303... try:
1304... for next in nexts:
1305... yield next()
1306... except StopIteration:
1307... pending -= 1
1308... nexts = cycle(islice(nexts, pending))
1309
1310>>> def powerset(iterable):
Raymond Hettingerace67332009-01-26 02:23:50 +00001311... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1312... s = list(iterable)
1313... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001314
Georg Brandl3dbca812008-07-23 16:10:53 +00001315>>> def combinations_with_replacement(iterable, r):
1316... "combinations_with_replacement('ABC', 3) --> AA AB AC BB BC CC"
1317... pool = tuple(iterable)
1318... n = len(pool)
1319... indices = [0] * r
1320... yield tuple(pool[i] for i in indices)
1321... while 1:
1322... for i in reversed(range(r)):
1323... if indices[i] != n - 1:
1324... break
1325... else:
1326... return
1327... indices[i:] = [indices[i] + 1] * (r - i)
1328... yield tuple(pool[i] for i in indices)
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001329
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001330>>> def unique_everseen(iterable, key=None):
1331... "List unique elements, preserving order. Remember all elements ever seen."
1332... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1333... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1334... seen = set()
1335... seen_add = seen.add
1336... if key is None:
1337... for element in iterable:
1338... if element not in seen:
1339... seen_add(element)
1340... yield element
1341... else:
1342... for element in iterable:
1343... k = key(element)
1344... if k not in seen:
1345... seen_add(k)
1346... yield element
1347
1348>>> def unique_justseen(iterable, key=None):
1349... "List unique elements, preserving order. Remember only the element just seen."
1350... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1351... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1352... return map(next, map(itemgetter(1), groupby(iterable, key)))
1353
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001354This is not part of the examples but it tests to make sure the definitions
1355perform as purported.
1356
Raymond Hettingera098b332003-09-08 23:58:40 +00001357>>> take(10, count())
1358[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1359
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001360>>> list(enumerate('abc'))
1361[(0, 'a'), (1, 'b'), (2, 'c')]
1362
1363>>> list(islice(tabulate(lambda x: 2*x), 4))
1364[0, 2, 4, 6]
1365
1366>>> nth('abcde', 3)
1367['d']
1368
Guido van Rossum805365e2007-05-07 22:24:25 +00001369>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000137050
1371
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001372>>> a = [[1, 2, 3], [4, 5, 6]]
1373>>> flatten(a)
1374[1, 2, 3, 4, 5, 6]
1375
1376>>> list(repeatfunc(pow, 5, 2, 3))
1377[8, 8, 8, 8, 8]
1378
1379>>> import random
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00001380>>> take(5, map(int, repeatfunc(random.random)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001381[0, 0, 0, 0, 0]
1382
Raymond Hettingerd591f662003-10-26 15:34:50 +00001383>>> list(pairwise('abcd'))
1384[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001385
Raymond Hettingerd591f662003-10-26 15:34:50 +00001386>>> list(pairwise([]))
1387[]
1388
1389>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001390[]
1391
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001392>>> list(islice(padnone('abc'), 0, 6))
1393['a', 'b', 'c', None, None, None]
1394
1395>>> list(ncycles('abc', 3))
1396['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1397
1398>>> dotproduct([1,2,3], [4,5,6])
139932
1400
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001401>>> list(grouper(3, 'abcdefg', 'x'))
1402[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1403
1404>>> list(roundrobin('abc', 'd', 'ef'))
1405['a', 'd', 'e', 'b', 'f', 'c']
1406
Raymond Hettingerace67332009-01-26 02:23:50 +00001407>>> list(powerset([1,2,3]))
1408[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001409
Georg Brandl3dbca812008-07-23 16:10:53 +00001410>>> list(combinations_with_replacement('abc', 2))
1411[('a', 'a'), ('a', 'b'), ('a', 'c'), ('b', 'b'), ('b', 'c'), ('c', 'c')]
1412
Raymond Hettinger5bad41e2009-01-08 21:01:54 +00001413>>> list(combinations_with_replacement('01', 3))
1414[('0', '0', '0'), ('0', '0', '1'), ('0', '1', '1'), ('1', '1', '1')]
1415
1416>>> def combinations_with_replacement2(iterable, r):
1417... 'Alternate version that filters from product()'
1418... pool = tuple(iterable)
1419... n = len(pool)
1420... for indices in product(range(n), repeat=r):
1421... if sorted(indices) == list(indices):
1422... yield tuple(pool[i] for i in indices)
1423
1424>>> list(combinations_with_replacement('abc', 2)) == list(combinations_with_replacement2('abc', 2))
1425True
1426
1427>>> list(combinations_with_replacement('01', 3)) == list(combinations_with_replacement2('01', 3))
1428True
1429
1430>>> list(combinations_with_replacement('2310', 6)) == list(combinations_with_replacement2('2310', 6))
1431True
1432
Raymond Hettingerad9d96b2009-01-02 21:39:07 +00001433>>> list(unique_everseen('AAAABBBCCDAABBB'))
1434['A', 'B', 'C', 'D']
1435
1436>>> list(unique_everseen('ABBCcAD', str.lower))
1437['A', 'B', 'C', 'D']
1438
1439>>> list(unique_justseen('AAAABBBCCDAABBB'))
1440['A', 'B', 'C', 'D', 'A', 'B']
1441
1442>>> list(unique_justseen('ABBCcAD', str.lower))
1443['A', 'B', 'C', 'A', 'D']
1444
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001445"""
1446
1447__test__ = {'libreftest' : libreftest}
1448
1449def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001450 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00001451 RegressionTests, LengthTransparency,
Christian Heimesdd15f6c2008-03-16 00:07:10 +00001452 SubclassWithKwargsTest, TestExamples)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001453 support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001454
1455 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001456 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001457 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001458 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00001459 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001460 support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001461 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001462 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001463 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001464
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001465 # doctest the examples in the library reference
Benjamin Petersonee8712c2008-05-20 21:35:26 +00001466 support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001467
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001468if __name__ == "__main__":
1469 test_main(verbose=True)