blob: 744b344a81747cd1f4a0af0fc38f668344624e9f [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001import unittest
2from test import test_support
3from itertools import *
Raymond Hettingera9f60922004-10-17 16:40:14 +00004from weakref import proxy
Raymond Hettinger2012f172003-02-07 05:32:58 +00005import sys
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00006import operator
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00007import random
Christian Heimesc3f30c42008-02-22 16:37:40 +00008from functools import reduce
Guido van Rossum360e4b82007-05-14 22:51:27 +00009maxsize = test_support.MAX_Py_ssize_t
10minsize = -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
Raymond Hettinger96ef8112003-02-01 00:10:11 +000047class TestBasicOps(unittest.TestCase):
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000048 def test_chain(self):
49 self.assertEqual(list(chain('abc', 'def')), list('abcdef'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000050 self.assertEqual(list(chain('abc')), list('abc'))
51 self.assertEqual(list(chain('')), [])
Raymond Hettinger02420702003-06-29 20:36:23 +000052 self.assertEqual(take(4, chain('abc', 'def')), list('abcd'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000053 self.assertRaises(TypeError, chain, 2, 3)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000054
Raymond Hettinger96ef8112003-02-01 00:10:11 +000055 def test_count(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +000056 self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
57 self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
58 self.assertEqual(take(2, lzip('abc',count(3))), [('a', 3), ('b', 4)])
Guido van Rossum8ce8a782007-11-01 19:42:39 +000059 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
60 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger96ef8112003-02-01 00:10:11 +000061 self.assertRaises(TypeError, count, 2, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000062 self.assertRaises(TypeError, count, 'a')
Guido van Rossum8ce8a782007-11-01 19:42:39 +000063 self.assertEqual(list(islice(count(maxsize-5), 10)),
64 list(range(maxsize-5, maxsize+5)))
65 self.assertEqual(list(islice(count(-maxsize-5), 10)),
66 list(range(-maxsize-5, -maxsize+5)))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +000067 c = count(3)
68 self.assertEqual(repr(c), 'count(3)')
Georg Brandla18af4e2007-04-21 15:47:16 +000069 next(c)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +000070 self.assertEqual(repr(c), 'count(4)')
Thomas Wouters89f507f2006-12-13 04:49:30 +000071 c = count(-9)
72 self.assertEqual(repr(c), 'count(-9)')
Georg Brandla18af4e2007-04-21 15:47:16 +000073 next(c)
74 self.assertEqual(next(c), -8)
Christian Heimesa37d4c62007-12-04 23:02:19 +000075 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 +000076 # Test repr (ignoring the L in longs)
77 r1 = repr(count(i)).replace('L', '')
78 r2 = 'count(%r)'.__mod__(i).replace('L', '')
79 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +000080
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000081 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +000082 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000083 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000084 self.assertRaises(TypeError, cycle)
85 self.assertRaises(TypeError, cycle, 5)
86 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000087
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000088 def test_groupby(self):
89 # Check whether it accepts arguments correctly
90 self.assertEqual([], list(groupby([])))
91 self.assertEqual([], list(groupby([], key=id)))
92 self.assertRaises(TypeError, list, groupby('abc', []))
93 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +000094 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000095
96 # Check normal input
97 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
98 (2,15,22), (3,16,23), (3,17,23)]
99 dup = []
100 for k, g in groupby(s, lambda r:r[0]):
101 for elem in g:
102 self.assertEqual(k, elem[0])
103 dup.append(elem)
104 self.assertEqual(s, dup)
105
106 # Check nested case
107 dup = []
108 for k, g in groupby(s, lambda r:r[0]):
109 for ik, ig in groupby(g, lambda r:r[2]):
110 for elem in ig:
111 self.assertEqual(k, elem[0])
112 self.assertEqual(ik, elem[2])
113 dup.append(elem)
114 self.assertEqual(s, dup)
115
116 # Check case where inner iterator is not used
117 keys = [k for k, g in groupby(s, lambda r:r[0])]
118 expectedkeys = set([r[0] for r in s])
119 self.assertEqual(set(keys), expectedkeys)
120 self.assertEqual(len(keys), len(expectedkeys))
121
122 # Exercise pipes and filters style
123 s = 'abracadabra'
124 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000125 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000126 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
127 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000128 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000129 self.assertEqual(r, ['a', 'b', 'r'])
130 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000131 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000132 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
133 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000134 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000135 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
136
Georg Brandla18af4e2007-04-21 15:47:16 +0000137 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000138 class ExpectedError(Exception):
139 pass
140 def delayed_raise(n=0):
141 for i in range(n):
142 yield 'yo'
143 raise ExpectedError
144 def gulp(iterable, keyp=None, func=list):
145 return [func(g) for k, g in groupby(iterable, keyp)]
146
Georg Brandla18af4e2007-04-21 15:47:16 +0000147 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000148 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000149 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000150 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
151
152 # __cmp__ failure
153 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000154 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000155 raise ExpectedError
156 s = [DummyCmp(), DummyCmp(), None]
157
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000158 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000159 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000160 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000161 self.assertRaises(ExpectedError, gulp, s)
162
163 # keyfunc failure
164 def keyfunc(obj):
165 if keyfunc.skip > 0:
166 keyfunc.skip -= 1
167 return obj
168 else:
169 raise ExpectedError
170
171 # keyfunc failure on outer object
172 keyfunc.skip = 0
173 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
174 keyfunc.skip = 1
175 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
176
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000177 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000178 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000179 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000180 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000181 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000182 self.assertRaises(TypeError, ifilter, lambda x:x)
183 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000184 self.assertRaises(TypeError, ifilter, isEven, 3)
Georg Brandla18af4e2007-04-21 15:47:16 +0000185 self.assertRaises(TypeError, next, ifilter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000186
187 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000188 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
189 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000190 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000191 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000192 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
193 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000194 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Georg Brandla18af4e2007-04-21 15:47:16 +0000195 self.assertRaises(TypeError, next, ifilterfalse(range(6), range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000196
197 def test_izip(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000198 # XXX This is rather silly now that builtin zip() calls izip()...
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000199 ans = [(x,y) for x, y in izip('abc',count())]
200 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Guido van Rossum801f0d72006-08-24 19:48:10 +0000201 self.assertEqual(list(izip('abc', range(6))), lzip('abc', range(6)))
202 self.assertEqual(list(izip('abcdef', range(3))), lzip('abcdef', range(3)))
203 self.assertEqual(take(3,izip('abcdef', count())), lzip('abcdef', range(3)))
204 self.assertEqual(list(izip('abcdef')), lzip('abcdef'))
205 self.assertEqual(list(izip()), lzip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000206 self.assertRaises(TypeError, izip, 3)
207 self.assertRaises(TypeError, izip, range(3), 3)
208 # Check tuple re-use (implementation detail)
209 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000210 lzip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000211 self.assertEqual([pair for pair in izip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000212 lzip('abc', 'def'))
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000213 ids = list(map(id, izip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000214 self.assertEqual(min(ids), max(ids))
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000215 ids = list(map(id, list(izip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000216 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000217
Thomas Wouterscf297e42007-02-23 15:07:44 +0000218 def test_iziplongest(self):
219 for args in [
220 ['abc', range(6)],
221 [range(6), 'abc'],
222 [range(1000), range(2000,2100), range(3000,3050)],
223 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
224 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
225 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000226 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
227 for i in range(max(map(len, args)))]
Thomas Wouterscf297e42007-02-23 15:07:44 +0000228 self.assertEqual(list(izip_longest(*args)), target)
229 self.assertEqual(list(izip_longest(*args, **{})), target)
230 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
231 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000232
Thomas Wouterscf297e42007-02-23 15:07:44 +0000233 self.assertEqual(take(3,izip_longest('abcdef', count())), list(zip('abcdef', range(3)))) # take 3 from infinite input
234
235 self.assertEqual(list(izip_longest()), list(zip()))
236 self.assertEqual(list(izip_longest([])), list(zip([])))
237 self.assertEqual(list(izip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000238
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000239 self.assertEqual(list(izip_longest('abc', 'defg', **{})),
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000240 list(izip(list('abc')+[None], 'defg'))) # empty keyword dict
Thomas Wouterscf297e42007-02-23 15:07:44 +0000241 self.assertRaises(TypeError, izip_longest, 3)
242 self.assertRaises(TypeError, izip_longest, range(3), 3)
243
244 for stmt in [
245 "izip_longest('abc', fv=1)",
Guido van Rossumd8faa362007-04-27 19:54:29 +0000246 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000247 ]:
248 try:
249 eval(stmt, globals(), locals())
250 except TypeError:
251 pass
252 else:
253 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000254
Thomas Wouterscf297e42007-02-23 15:07:44 +0000255 # Check tuple re-use (implementation detail)
256 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
257 list(zip('abc', 'def')))
258 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
259 list(zip('abc', 'def')))
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000260 ids = list(map(id, izip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000261 self.assertEqual(min(ids), max(ids))
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000262 ids = list(map(id, list(izip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000263 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
264
Christian Heimesc3f30c42008-02-22 16:37:40 +0000265 def test_product(self):
266 for args, result in [
267 ([], []), # zero iterables ??? is this correct
268 (['ab'], [('a',), ('b',)]), # one iterable
269 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
270 ([range(0), range(2), range(3)], []), # first iterable with zero length
271 ([range(2), range(0), range(3)], []), # middle iterable with zero length
272 ([range(2), range(3), range(0)], []), # last iterable with zero length
273 ]:
274 self.assertEqual(list(product(*args)), result)
275 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
276 self.assertRaises(TypeError, product, range(6), None)
277 argtypes = ['', 'abc', '', range(0), range(4), dict(a=1, b=2, c=3),
278 set('abcdefg'), range(11), tuple(range(13))]
279 for i in range(100):
280 args = [random.choice(argtypes) for j in range(random.randrange(5))]
281 n = reduce(operator.mul, map(len, args), 1) if args else 0
282 self.assertEqual(len(list(product(*args))), n)
283 args = map(iter, args)
284 self.assertEqual(len(list(product(*args))), n)
285
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000286 # Test implementation detail: tuple re-use
287 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
288 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Christian Heimesc3f30c42008-02-22 16:37:40 +0000289
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000290 def test_repeat(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000291 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000292 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000293 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000294 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000295 self.assertEqual(list(repeat('a', 0)), [])
296 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000297 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000298 self.assertRaises(TypeError, repeat, None, 3, 4)
299 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000300 r = repeat(1+0j)
301 self.assertEqual(repr(r), 'repeat((1+0j))')
302 r = repeat(1+0j, 5)
303 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
304 list(r)
305 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000306
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000307 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000308 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
309 [0**1, 1**2, 2**3])
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000310 def tupleize(*args):
311 return args
312 self.assertEqual(list(imap(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000313 [('a',0),('b',1),('c',2)])
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000314 self.assertEqual(list(imap(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000315 [('a',0),('b',1),('c',2)])
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000316 self.assertEqual(take(2,imap(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000317 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000318 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000319 self.assertRaises(TypeError, imap)
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000320 self.assertRaises(TypeError, list, imap(None, range(3), range(3)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000321 self.assertRaises(TypeError, imap, operator.neg)
Georg Brandla18af4e2007-04-21 15:47:16 +0000322 self.assertRaises(TypeError, next, imap(10, range(5)))
323 self.assertRaises(ValueError, next, imap(errfunc, [4], [5]))
324 self.assertRaises(TypeError, next, imap(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000325
326 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000327 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
328 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000329 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
330 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000331 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +0000332 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
333 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000334 self.assertRaises(TypeError, starmap)
335 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000336 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
337 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
338 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000339
340 def test_islice(self):
341 for args in [ # islice(args) should agree with range(args)
342 (10, 20, 3),
343 (10, 3, 20),
344 (10, 20),
345 (10, 3),
346 (20,)
347 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000348 self.assertEqual(list(islice(range(100), *args)),
349 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000350
351 for args, tgtargs in [ # Stop when seqn is exhausted
352 ((10, 110, 3), ((10, 100, 3))),
353 ((10, 110), ((10, 100))),
354 ((110,), (100,))
355 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000356 self.assertEqual(list(islice(range(100), *args)),
357 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000358
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000359 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +0000360 self.assertEqual(list(islice(range(10), None)), list(range(10)))
361 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
362 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
363 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
364 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000365
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000366 # Test number of items consumed SF #1171417
367 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +0000368 self.assertEqual(list(islice(it, 3)), list(range(3)))
369 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000370
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000371 # Test invalid arguments
Guido van Rossum805365e2007-05-07 22:24:25 +0000372 self.assertRaises(TypeError, islice, range(10))
373 self.assertRaises(TypeError, islice, range(10), 1, 2, 3, 4)
374 self.assertRaises(ValueError, islice, range(10), -5, 10, 1)
375 self.assertRaises(ValueError, islice, range(10), 1, -5, -1)
376 self.assertRaises(ValueError, islice, range(10), 1, 10, -1)
377 self.assertRaises(ValueError, islice, range(10), 1, 10, 0)
378 self.assertRaises(ValueError, islice, range(10), 'a')
379 self.assertRaises(ValueError, islice, range(10), 'a', 1)
380 self.assertRaises(ValueError, islice, range(10), 1, 'a')
381 self.assertRaises(ValueError, islice, range(10), 'a', 1, 1)
382 self.assertRaises(ValueError, islice, range(10), 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +0000383 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000384
385 def test_takewhile(self):
386 data = [1, 3, 5, 20, 2, 4, 6, 8]
387 underten = lambda x: x<10
388 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000389 self.assertEqual(list(takewhile(underten, [])), [])
390 self.assertRaises(TypeError, takewhile)
391 self.assertRaises(TypeError, takewhile, operator.pow)
392 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000393 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
394 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000395 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
396 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +0000397 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000398
399 def test_dropwhile(self):
400 data = [1, 3, 5, 20, 2, 4, 6, 8]
401 underten = lambda x: x<10
402 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000403 self.assertEqual(list(dropwhile(underten, [])), [])
404 self.assertRaises(TypeError, dropwhile)
405 self.assertRaises(TypeError, dropwhile, operator.pow)
406 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000407 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
408 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000409
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000410 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000411 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000412 def irange(n):
Guido van Rossum805365e2007-05-07 22:24:25 +0000413 for i in range(n):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000414 yield i
415
416 a, b = tee([]) # test empty iterator
417 self.assertEqual(list(a), [])
418 self.assertEqual(list(b), [])
419
420 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +0000421 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000422
423 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +0000424 self.assertEqual(list(a), list(range(n)))
425 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000426
427 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000428 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000429 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000430 del a
Guido van Rossum805365e2007-05-07 22:24:25 +0000431 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000432
433 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000434 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000435 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000436 del b
Guido van Rossum805365e2007-05-07 22:24:25 +0000437 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000438
Guido van Rossum805365e2007-05-07 22:24:25 +0000439 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000440 order = [0]*n + [1]*n
441 random.shuffle(order)
442 lists = ([], [])
443 its = tee(irange(n))
444 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +0000445 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000446 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +0000447 self.assertEqual(lists[0], list(range(n)))
448 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000449
Raymond Hettingerad983e72003-11-12 14:32:26 +0000450 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000451 self.assertRaises(TypeError, tee)
452 self.assertRaises(TypeError, tee, 3)
453 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000454 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000455
Raymond Hettingerad983e72003-11-12 14:32:26 +0000456 # tee object should be instantiable
457 a, b = tee('abc')
458 c = type(a)('def')
459 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000460
Raymond Hettingerad983e72003-11-12 14:32:26 +0000461 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +0000462 a, b, c = tee(range(2000), 3)
463 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000464 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +0000465 self.assertEqual(list(b), list(range(2000)))
466 self.assertEqual([next(c), next(c)], list(range(2)))
467 self.assertEqual(list(a), list(range(100,2000)))
468 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000469
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000470 # test values of n
471 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000472 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +0000473 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000474 result = tee('abc', n)
475 self.assertEqual(type(result), tuple)
476 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000477 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000478
Raymond Hettingerad983e72003-11-12 14:32:26 +0000479 # tee pass-through to copyable iterator
480 a, b = tee('abc')
481 c, d = tee(a)
482 self.assert_(a is c)
483
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000484 # test tee_new
485 t1, t2 = tee('abc')
486 tnew = type(t1)
487 self.assertRaises(TypeError, tnew)
488 self.assertRaises(TypeError, tnew, 10)
489 t3 = tnew(t1)
490 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000491
Raymond Hettingera9f60922004-10-17 16:40:14 +0000492 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +0000493 a, b = tee(range(10))
Raymond Hettingera9f60922004-10-17 16:40:14 +0000494 p = proxy(a)
495 self.assertEqual(getattr(p, '__class__'), type(b))
496 del a
497 self.assertRaises(ReferenceError, getattr, p, '__class__')
498
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000499 def test_StopIteration(self):
Georg Brandla18af4e2007-04-21 15:47:16 +0000500 self.assertRaises(StopIteration, next, izip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000501
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000502 for f in (chain, cycle, izip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +0000503 self.assertRaises(StopIteration, next, f([]))
504 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000505
Georg Brandla18af4e2007-04-21 15:47:16 +0000506 self.assertRaises(StopIteration, next, islice([], None))
507 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000508
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000509 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +0000510 self.assertRaises(StopIteration, next, p)
511 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000512 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +0000513 self.assertRaises(StopIteration, next, p)
514 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000515
Georg Brandla18af4e2007-04-21 15:47:16 +0000516 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000517
518 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +0000519 self.assertRaises(StopIteration, next, f(lambda x:x, []))
520 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000521
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000522class TestGC(unittest.TestCase):
523
524 def makecycle(self, iterator, container):
525 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +0000526 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000527 del container, iterator
528
529 def test_chain(self):
530 a = []
531 self.makecycle(chain(a), a)
532
533 def test_cycle(self):
534 a = []
535 self.makecycle(cycle([a]*2), a)
536
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000537 def test_dropwhile(self):
538 a = []
539 self.makecycle(dropwhile(bool, [0, a, a]), a)
540
541 def test_groupby(self):
542 a = []
543 self.makecycle(groupby([a]*2, lambda x:x), a)
544
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000545 def test_ifilter(self):
546 a = []
547 self.makecycle(ifilter(lambda x:True, [a]*2), a)
548
549 def test_ifilterfalse(self):
550 a = []
551 self.makecycle(ifilterfalse(lambda x:False, a), a)
552
553 def test_izip(self):
554 a = []
555 self.makecycle(izip([a]*2, [a]*3), a)
556
557 def test_imap(self):
558 a = []
559 self.makecycle(imap(lambda x:x, [a]*2), a)
560
561 def test_islice(self):
562 a = []
563 self.makecycle(islice([a]*2, None), a)
564
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000565 def test_repeat(self):
566 a = []
567 self.makecycle(repeat(a), a)
568
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000569 def test_starmap(self):
570 a = []
571 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
572
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000573 def test_takewhile(self):
574 a = []
575 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
576
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000577def R(seqn):
578 'Regular generator'
579 for i in seqn:
580 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000581
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000582class G:
583 'Sequence using __getitem__'
584 def __init__(self, seqn):
585 self.seqn = seqn
586 def __getitem__(self, i):
587 return self.seqn[i]
588
589class I:
590 'Sequence using iterator protocol'
591 def __init__(self, seqn):
592 self.seqn = seqn
593 self.i = 0
594 def __iter__(self):
595 return self
Georg Brandla18af4e2007-04-21 15:47:16 +0000596 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000597 if self.i >= len(self.seqn): raise StopIteration
598 v = self.seqn[self.i]
599 self.i += 1
600 return v
601
602class Ig:
603 'Sequence using iterator protocol defined with a generator'
604 def __init__(self, seqn):
605 self.seqn = seqn
606 self.i = 0
607 def __iter__(self):
608 for val in self.seqn:
609 yield val
610
611class X:
612 'Missing __getitem__ and __iter__'
613 def __init__(self, seqn):
614 self.seqn = seqn
615 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +0000616 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000617 if self.i >= len(self.seqn): raise StopIteration
618 v = self.seqn[self.i]
619 self.i += 1
620 return v
621
622class N:
Georg Brandla18af4e2007-04-21 15:47:16 +0000623 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000624 def __init__(self, seqn):
625 self.seqn = seqn
626 self.i = 0
627 def __iter__(self):
628 return self
629
630class E:
631 'Test propagation of exceptions'
632 def __init__(self, seqn):
633 self.seqn = seqn
634 self.i = 0
635 def __iter__(self):
636 return self
Georg Brandla18af4e2007-04-21 15:47:16 +0000637 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +0000638 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000639
640class S:
641 'Test immediate stop'
642 def __init__(self, seqn):
643 pass
644 def __iter__(self):
645 return self
Georg Brandla18af4e2007-04-21 15:47:16 +0000646 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000647 raise StopIteration
648
649def L(seqn):
650 'Test multiple tiers of iterators'
651 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
652
653
654class TestVariousIteratorArgs(unittest.TestCase):
655
656 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000657 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000658 for g in (G, I, Ig, S, L, R):
659 self.assertEqual(list(chain(g(s))), list(g(s)))
660 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
661 self.assertRaises(TypeError, chain, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000662 self.assertRaises(TypeError, chain, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000663 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
664
Christian Heimesc3f30c42008-02-22 16:37:40 +0000665 def test_product(self):
666 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
667 self.assertRaises(TypeError, product, X(s))
668 self.assertRaises(TypeError, product, N(s))
669 self.assertRaises(ZeroDivisionError, product, E(s))
670
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000671 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000672 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000673 for g in (G, I, Ig, S, L, R):
674 tgtlen = len(s) * 3
675 expected = list(g(s))*3
676 actual = list(islice(cycle(g(s)), tgtlen))
677 self.assertEqual(actual, expected)
678 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000679 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000680 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
681
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000682 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000683 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000684 for g in (G, I, Ig, S, L, R):
685 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
686 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000687 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000688 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
689
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000690 def test_ifilter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000691 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000692 for g in (G, I, Ig, S, L, R):
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000693 self.assertEqual(list(ifilter(isEven, g(s))),
694 [x for x in g(s) if isEven(x)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000695 self.assertRaises(TypeError, ifilter, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000696 self.assertRaises(TypeError, ifilter, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000697 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
698
699 def test_ifilterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000700 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000701 for g in (G, I, Ig, S, L, R):
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000702 self.assertEqual(list(ifilterfalse(isEven, g(s))),
703 [x for x in g(s) if isOdd(x)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000704 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000705 self.assertRaises(TypeError, ifilterfalse, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000706 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
707
708 def test_izip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000709 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000710 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000711 self.assertEqual(list(izip(g(s))), lzip(g(s)))
712 self.assertEqual(list(izip(g(s), g(s))), lzip(g(s), g(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000713 self.assertRaises(TypeError, izip, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000714 self.assertRaises(TypeError, izip, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000715 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
716
Thomas Wouterscf297e42007-02-23 15:07:44 +0000717 def test_iziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000718 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000719 for g in (G, I, Ig, S, L, R):
720 self.assertEqual(list(izip_longest(g(s))), list(zip(g(s))))
721 self.assertEqual(list(izip_longest(g(s), g(s))), list(zip(g(s), g(s))))
722 self.assertRaises(TypeError, izip_longest, X(s))
723 self.assertRaises(TypeError, izip_longest, N(s))
724 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
725
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000726 def test_imap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000727 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000728 for g in (G, I, Ig, S, L, R):
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000729 self.assertEqual(list(imap(onearg, g(s))),
730 [onearg(x) for x in g(s)])
731 self.assertEqual(list(imap(operator.pow, g(s), g(s))),
732 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000733 self.assertRaises(TypeError, imap, onearg, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000734 self.assertRaises(TypeError, imap, onearg, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000735 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
736
737 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000738 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000739 for g in (G, I, Ig, S, L, R):
740 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
741 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000742 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000743 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
744
745 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000746 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000747 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000748 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000749 self.assertEqual(list(starmap(operator.pow, g(ss))),
750 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000751 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000752 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000753 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
754
755 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000756 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000757 for g in (G, I, Ig, S, L, R):
758 tgt = []
759 for elem in g(s):
760 if not isEven(elem): break
761 tgt.append(elem)
762 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
763 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000764 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000765 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
766
767 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000768 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000769 for g in (G, I, Ig, S, L, R):
770 tgt = []
771 for elem in g(s):
772 if not tgt and isOdd(elem): continue
773 tgt.append(elem)
774 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
775 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000776 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000777 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
778
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000779 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000780 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000781 for g in (G, I, Ig, S, L, R):
782 it1, it2 = tee(g(s))
783 self.assertEqual(list(it1), list(g(s)))
784 self.assertEqual(list(it2), list(g(s)))
785 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000786 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000787 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
788
Raymond Hettinger5cab2e32004-02-10 09:25:40 +0000789class LengthTransparency(unittest.TestCase):
790
791 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +0000792 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +0000793 self.assertEqual(len(repeat(None, 50)), 50)
794 self.assertRaises(TypeError, len, repeat(None))
795
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000796class RegressionTests(unittest.TestCase):
797
798 def test_sf_793826(self):
799 # Fix Armin Rigo's successful efforts to wreak havoc
800
801 def mutatingtuple(tuple1, f, tuple2):
802 # this builds a tuple t which is a copy of tuple1,
803 # then calls f(t), then mutates t to be equal to tuple2
804 # (needs len(tuple1) == len(tuple2)).
805 def g(value, first=[1]):
806 if first:
807 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +0000808 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000809 return value
810 items = list(tuple2)
811 items[1:1] = list(tuple1)
812 gen = imap(g, items)
813 z = izip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +0000814 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000815
816 def f(t):
817 global T
818 T = t
819 first[:] = list(T)
820
821 first = []
822 mutatingtuple((1,2,3), f, (4,5,6))
823 second = list(T)
824 self.assertEqual(first, second)
825
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000826
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000827 def test_sf_950057(self):
828 # Make sure that chain() and cycle() catch exceptions immediately
829 # rather than when shifting between input sources
830
831 def gen1():
832 hist.append(0)
833 yield 1
834 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +0000835 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000836 hist.append(2)
837
838 def gen2(x):
839 hist.append(3)
840 yield 2
841 hist.append(4)
842 if x:
843 raise StopIteration
844
845 hist = []
846 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
847 self.assertEqual(hist, [0,1])
848
849 hist = []
850 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
851 self.assertEqual(hist, [0,1])
852
853 hist = []
854 self.assertRaises(AssertionError, list, cycle(gen1()))
855 self.assertEqual(hist, [0,1])
856
Thomas Woutersb2137042007-02-01 18:02:27 +0000857class SubclassWithKwargsTest(unittest.TestCase):
858 def test_keywords_in_subclass(self):
859 # count is not subclassable...
860 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
861 starmap, islice, takewhile, dropwhile, cycle):
862 class Subclass(cls):
863 def __init__(self, newarg=None, *args):
864 cls.__init__(self, *args)
865 try:
866 Subclass(newarg=1)
867 except TypeError as err:
868 # we expect type errors because of wrong argument count
869 self.failIf("does not take keyword arguments" in err.args[0])
870
871
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000872libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000873
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000874
875>>> amounts = [120.15, 764.05, 823.14]
876>>> for checknum, amount in izip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +0000877... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000878...
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000879Check 1200 is for $120.15
880Check 1201 is for $764.05
881Check 1202 is for $823.14
882
883>>> import operator
Guido van Rossum805365e2007-05-07 22:24:25 +0000884>>> for cube in imap(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +0000885... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000886...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00008871
8888
88927
890
891>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +0000892>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +0000893... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +0000894...
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000895Alex
896Laura
897Martin
898Walter
899Samuele
900
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000901>>> from operator import itemgetter
902>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +0000903>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000904>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000905... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000906...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00009071 ['a', 'c', 'e']
9082 ['b', 'd', 'f']
9093 ['g']
910
Raymond Hettinger734fb572004-01-20 20:04:40 +0000911# Find runs of consecutive numbers using groupby. The key to the solution
912# is differencing with a range so that consecutive numbers all appear in
913# same group.
914>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +0000915>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000916... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000917...
Raymond Hettinger734fb572004-01-20 20:04:40 +0000918[1]
919[4, 5, 6]
920[10]
921[15, 16, 17, 18]
922[22]
923[25, 26, 27, 28]
924
Raymond Hettingera098b332003-09-08 23:58:40 +0000925>>> def take(n, seq):
926... return list(islice(seq, n))
927
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000928>>> def enumerate(iterable):
929... return izip(count(), iterable)
930
931>>> def tabulate(function):
932... "Return function(0), function(1), ..."
933... return imap(function, count())
934
935>>> def iteritems(mapping):
Guido van Rossumcc2b0162007-02-11 06:12:03 +0000936... return izip(mapping.keys(), mapping.values())
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000937
938>>> def nth(iterable, n):
939... "Returns the nth item"
Raymond Hettinger60eca932003-02-09 06:40:58 +0000940... return list(islice(iterable, n, n+1))
941
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000942>>> def all(seq, pred=None):
943... "Returns True if pred(x) is true for every element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000944... for elem in ifilterfalse(pred, seq):
945... return False
946... return True
Raymond Hettinger60eca932003-02-09 06:40:58 +0000947
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000948>>> def any(seq, pred=None):
949... "Returns True if pred(x) is true for at least one element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000950... for elem in ifilter(pred, seq):
951... return True
952... return False
Raymond Hettinger60eca932003-02-09 06:40:58 +0000953
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000954>>> def no(seq, pred=None):
955... "Returns True if pred(x) is false for every element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000956... for elem in ifilter(pred, seq):
957... return False
958... return True
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000959
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000960>>> def quantify(seq, pred=None):
961... "Count how many times the predicate is true in the sequence"
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000962... return sum(imap(pred, seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000963
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000964>>> def padnone(seq):
965... "Returns the sequence elements and then returns None indefinitely"
966... return chain(seq, repeat(None))
967
968>>> def ncycles(seq, n):
969... "Returns the sequence elements n times"
970... return chain(*repeat(seq, n))
971
972>>> def dotproduct(vec1, vec2):
973... return sum(imap(operator.mul, vec1, vec2))
974
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000975>>> def flatten(listOfLists):
976... return list(chain(*listOfLists))
977
978>>> def repeatfunc(func, times=None, *args):
979... "Repeat calls to func with specified arguments."
980... " Example: repeatfunc(random.random)"
981... if times is None:
982... return starmap(func, repeat(args))
983... else:
984... return starmap(func, repeat(args, times))
985
Raymond Hettingerd591f662003-10-26 15:34:50 +0000986>>> def pairwise(iterable):
987... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
988... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000989... try:
Georg Brandla18af4e2007-04-21 15:47:16 +0000990... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000991... except StopIteration:
992... pass
993... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000994
995This is not part of the examples but it tests to make sure the definitions
996perform as purported.
997
Raymond Hettingera098b332003-09-08 23:58:40 +0000998>>> take(10, count())
999[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1000
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001001>>> list(enumerate('abc'))
1002[(0, 'a'), (1, 'b'), (2, 'c')]
1003
1004>>> list(islice(tabulate(lambda x: 2*x), 4))
1005[0, 2, 4, 6]
1006
1007>>> nth('abcde', 3)
1008['d']
1009
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001010>>> all([2, 4, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001011True
1012
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001013>>> all([2, 3, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001014False
1015
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001016>>> any([2, 4, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001017True
1018
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001019>>> any([1, 3, 5, 9], lambda x: x%2==0,)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001020False
1021
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001022>>> no([1, 3, 5, 9], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001023True
1024
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001025>>> no([1, 2, 5, 9], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001026False
1027
Guido van Rossum805365e2007-05-07 22:24:25 +00001028>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000102950
1030
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001031>>> a = [[1, 2, 3], [4, 5, 6]]
1032>>> flatten(a)
1033[1, 2, 3, 4, 5, 6]
1034
1035>>> list(repeatfunc(pow, 5, 2, 3))
1036[8, 8, 8, 8, 8]
1037
1038>>> import random
1039>>> take(5, imap(int, repeatfunc(random.random)))
1040[0, 0, 0, 0, 0]
1041
Raymond Hettingerd591f662003-10-26 15:34:50 +00001042>>> list(pairwise('abcd'))
1043[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001044
Raymond Hettingerd591f662003-10-26 15:34:50 +00001045>>> list(pairwise([]))
1046[]
1047
1048>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001049[]
1050
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001051>>> list(islice(padnone('abc'), 0, 6))
1052['a', 'b', 'c', None, None, None]
1053
1054>>> list(ncycles('abc', 3))
1055['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1056
1057>>> dotproduct([1,2,3], [4,5,6])
105832
1059
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001060"""
1061
1062__test__ = {'libreftest' : libreftest}
1063
1064def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001065 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00001066 RegressionTests, LengthTransparency,
1067 SubclassWithKwargsTest)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001068 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001069
1070 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001071 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001072 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001073 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00001074 for i in range(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001075 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001076 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001077 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001078 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001079
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001080 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001081 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001082
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001083if __name__ == "__main__":
1084 test_main(verbose=True)