blob: 380d1217d586e5f749d07f24d63de3c820e47193 [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
286
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000287 def test_repeat(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000288 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000289 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000290 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000291 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000292 self.assertEqual(list(repeat('a', 0)), [])
293 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000294 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000295 self.assertRaises(TypeError, repeat, None, 3, 4)
296 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000297 r = repeat(1+0j)
298 self.assertEqual(repr(r), 'repeat((1+0j))')
299 r = repeat(1+0j, 5)
300 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
301 list(r)
302 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000303
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000304 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000305 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
306 [0**1, 1**2, 2**3])
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000307 def tupleize(*args):
308 return args
309 self.assertEqual(list(imap(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000310 [('a',0),('b',1),('c',2)])
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000311 self.assertEqual(list(imap(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000312 [('a',0),('b',1),('c',2)])
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000313 self.assertEqual(take(2,imap(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000314 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000315 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000316 self.assertRaises(TypeError, imap)
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000317 self.assertRaises(TypeError, list, imap(None, range(3), range(3)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000318 self.assertRaises(TypeError, imap, operator.neg)
Georg Brandla18af4e2007-04-21 15:47:16 +0000319 self.assertRaises(TypeError, next, imap(10, range(5)))
320 self.assertRaises(ValueError, next, imap(errfunc, [4], [5]))
321 self.assertRaises(TypeError, next, imap(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000322
323 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000324 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
325 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000326 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
327 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000328 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +0000329 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
330 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000331 self.assertRaises(TypeError, starmap)
332 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000333 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
334 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
335 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000336
337 def test_islice(self):
338 for args in [ # islice(args) should agree with range(args)
339 (10, 20, 3),
340 (10, 3, 20),
341 (10, 20),
342 (10, 3),
343 (20,)
344 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000345 self.assertEqual(list(islice(range(100), *args)),
346 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000347
348 for args, tgtargs in [ # Stop when seqn is exhausted
349 ((10, 110, 3), ((10, 100, 3))),
350 ((10, 110), ((10, 100))),
351 ((110,), (100,))
352 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000353 self.assertEqual(list(islice(range(100), *args)),
354 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000355
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000356 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +0000357 self.assertEqual(list(islice(range(10), None)), list(range(10)))
358 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
359 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
360 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
361 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000362
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000363 # Test number of items consumed SF #1171417
364 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +0000365 self.assertEqual(list(islice(it, 3)), list(range(3)))
366 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000367
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000368 # Test invalid arguments
Guido van Rossum805365e2007-05-07 22:24:25 +0000369 self.assertRaises(TypeError, islice, range(10))
370 self.assertRaises(TypeError, islice, range(10), 1, 2, 3, 4)
371 self.assertRaises(ValueError, islice, range(10), -5, 10, 1)
372 self.assertRaises(ValueError, islice, range(10), 1, -5, -1)
373 self.assertRaises(ValueError, islice, range(10), 1, 10, -1)
374 self.assertRaises(ValueError, islice, range(10), 1, 10, 0)
375 self.assertRaises(ValueError, islice, range(10), 'a')
376 self.assertRaises(ValueError, islice, range(10), 'a', 1)
377 self.assertRaises(ValueError, islice, range(10), 1, 'a')
378 self.assertRaises(ValueError, islice, range(10), 'a', 1, 1)
379 self.assertRaises(ValueError, islice, range(10), 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +0000380 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000381
382 def test_takewhile(self):
383 data = [1, 3, 5, 20, 2, 4, 6, 8]
384 underten = lambda x: x<10
385 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000386 self.assertEqual(list(takewhile(underten, [])), [])
387 self.assertRaises(TypeError, takewhile)
388 self.assertRaises(TypeError, takewhile, operator.pow)
389 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000390 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
391 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000392 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
393 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +0000394 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000395
396 def test_dropwhile(self):
397 data = [1, 3, 5, 20, 2, 4, 6, 8]
398 underten = lambda x: x<10
399 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000400 self.assertEqual(list(dropwhile(underten, [])), [])
401 self.assertRaises(TypeError, dropwhile)
402 self.assertRaises(TypeError, dropwhile, operator.pow)
403 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000404 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
405 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000406
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000407 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000408 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000409 def irange(n):
Guido van Rossum805365e2007-05-07 22:24:25 +0000410 for i in range(n):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000411 yield i
412
413 a, b = tee([]) # test empty iterator
414 self.assertEqual(list(a), [])
415 self.assertEqual(list(b), [])
416
417 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +0000418 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000419
420 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +0000421 self.assertEqual(list(a), list(range(n)))
422 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000423
424 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000425 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000426 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000427 del a
Guido van Rossum805365e2007-05-07 22:24:25 +0000428 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000429
430 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000431 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000432 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000433 del b
Guido van Rossum805365e2007-05-07 22:24:25 +0000434 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000435
Guido van Rossum805365e2007-05-07 22:24:25 +0000436 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000437 order = [0]*n + [1]*n
438 random.shuffle(order)
439 lists = ([], [])
440 its = tee(irange(n))
441 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +0000442 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000443 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +0000444 self.assertEqual(lists[0], list(range(n)))
445 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000446
Raymond Hettingerad983e72003-11-12 14:32:26 +0000447 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000448 self.assertRaises(TypeError, tee)
449 self.assertRaises(TypeError, tee, 3)
450 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000451 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000452
Raymond Hettingerad983e72003-11-12 14:32:26 +0000453 # tee object should be instantiable
454 a, b = tee('abc')
455 c = type(a)('def')
456 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000457
Raymond Hettingerad983e72003-11-12 14:32:26 +0000458 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +0000459 a, b, c = tee(range(2000), 3)
460 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000461 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +0000462 self.assertEqual(list(b), list(range(2000)))
463 self.assertEqual([next(c), next(c)], list(range(2)))
464 self.assertEqual(list(a), list(range(100,2000)))
465 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000466
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000467 # test values of n
468 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000469 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +0000470 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000471 result = tee('abc', n)
472 self.assertEqual(type(result), tuple)
473 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000474 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000475
Raymond Hettingerad983e72003-11-12 14:32:26 +0000476 # tee pass-through to copyable iterator
477 a, b = tee('abc')
478 c, d = tee(a)
479 self.assert_(a is c)
480
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000481 # test tee_new
482 t1, t2 = tee('abc')
483 tnew = type(t1)
484 self.assertRaises(TypeError, tnew)
485 self.assertRaises(TypeError, tnew, 10)
486 t3 = tnew(t1)
487 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000488
Raymond Hettingera9f60922004-10-17 16:40:14 +0000489 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +0000490 a, b = tee(range(10))
Raymond Hettingera9f60922004-10-17 16:40:14 +0000491 p = proxy(a)
492 self.assertEqual(getattr(p, '__class__'), type(b))
493 del a
494 self.assertRaises(ReferenceError, getattr, p, '__class__')
495
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000496 def test_StopIteration(self):
Georg Brandla18af4e2007-04-21 15:47:16 +0000497 self.assertRaises(StopIteration, next, izip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000498
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000499 for f in (chain, cycle, izip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +0000500 self.assertRaises(StopIteration, next, f([]))
501 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000502
Georg Brandla18af4e2007-04-21 15:47:16 +0000503 self.assertRaises(StopIteration, next, islice([], None))
504 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000505
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000506 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +0000507 self.assertRaises(StopIteration, next, p)
508 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000509 p, q = tee(StopNow())
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
Georg Brandla18af4e2007-04-21 15:47:16 +0000513 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000514
515 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +0000516 self.assertRaises(StopIteration, next, f(lambda x:x, []))
517 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000518
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000519class TestGC(unittest.TestCase):
520
521 def makecycle(self, iterator, container):
522 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +0000523 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000524 del container, iterator
525
526 def test_chain(self):
527 a = []
528 self.makecycle(chain(a), a)
529
530 def test_cycle(self):
531 a = []
532 self.makecycle(cycle([a]*2), a)
533
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000534 def test_dropwhile(self):
535 a = []
536 self.makecycle(dropwhile(bool, [0, a, a]), a)
537
538 def test_groupby(self):
539 a = []
540 self.makecycle(groupby([a]*2, lambda x:x), a)
541
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000542 def test_ifilter(self):
543 a = []
544 self.makecycle(ifilter(lambda x:True, [a]*2), a)
545
546 def test_ifilterfalse(self):
547 a = []
548 self.makecycle(ifilterfalse(lambda x:False, a), a)
549
550 def test_izip(self):
551 a = []
552 self.makecycle(izip([a]*2, [a]*3), a)
553
554 def test_imap(self):
555 a = []
556 self.makecycle(imap(lambda x:x, [a]*2), a)
557
558 def test_islice(self):
559 a = []
560 self.makecycle(islice([a]*2, None), a)
561
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000562 def test_repeat(self):
563 a = []
564 self.makecycle(repeat(a), a)
565
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000566 def test_starmap(self):
567 a = []
568 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
569
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000570 def test_takewhile(self):
571 a = []
572 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
573
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000574def R(seqn):
575 'Regular generator'
576 for i in seqn:
577 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000578
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000579class G:
580 'Sequence using __getitem__'
581 def __init__(self, seqn):
582 self.seqn = seqn
583 def __getitem__(self, i):
584 return self.seqn[i]
585
586class I:
587 'Sequence using iterator protocol'
588 def __init__(self, seqn):
589 self.seqn = seqn
590 self.i = 0
591 def __iter__(self):
592 return self
Georg Brandla18af4e2007-04-21 15:47:16 +0000593 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000594 if self.i >= len(self.seqn): raise StopIteration
595 v = self.seqn[self.i]
596 self.i += 1
597 return v
598
599class Ig:
600 'Sequence using iterator protocol defined with a generator'
601 def __init__(self, seqn):
602 self.seqn = seqn
603 self.i = 0
604 def __iter__(self):
605 for val in self.seqn:
606 yield val
607
608class X:
609 'Missing __getitem__ and __iter__'
610 def __init__(self, seqn):
611 self.seqn = seqn
612 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +0000613 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000614 if self.i >= len(self.seqn): raise StopIteration
615 v = self.seqn[self.i]
616 self.i += 1
617 return v
618
619class N:
Georg Brandla18af4e2007-04-21 15:47:16 +0000620 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000621 def __init__(self, seqn):
622 self.seqn = seqn
623 self.i = 0
624 def __iter__(self):
625 return self
626
627class E:
628 'Test propagation of exceptions'
629 def __init__(self, seqn):
630 self.seqn = seqn
631 self.i = 0
632 def __iter__(self):
633 return self
Georg Brandla18af4e2007-04-21 15:47:16 +0000634 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +0000635 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000636
637class S:
638 'Test immediate stop'
639 def __init__(self, seqn):
640 pass
641 def __iter__(self):
642 return self
Georg Brandla18af4e2007-04-21 15:47:16 +0000643 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000644 raise StopIteration
645
646def L(seqn):
647 'Test multiple tiers of iterators'
648 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
649
650
651class TestVariousIteratorArgs(unittest.TestCase):
652
653 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000654 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000655 for g in (G, I, Ig, S, L, R):
656 self.assertEqual(list(chain(g(s))), list(g(s)))
657 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
658 self.assertRaises(TypeError, chain, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000659 self.assertRaises(TypeError, chain, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000660 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
661
Christian Heimesc3f30c42008-02-22 16:37:40 +0000662 def test_product(self):
663 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
664 self.assertRaises(TypeError, product, X(s))
665 self.assertRaises(TypeError, product, N(s))
666 self.assertRaises(ZeroDivisionError, product, E(s))
667
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000668 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000669 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000670 for g in (G, I, Ig, S, L, R):
671 tgtlen = len(s) * 3
672 expected = list(g(s))*3
673 actual = list(islice(cycle(g(s)), tgtlen))
674 self.assertEqual(actual, expected)
675 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000676 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000677 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
678
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000679 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000680 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000681 for g in (G, I, Ig, S, L, R):
682 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
683 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000684 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000685 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
686
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000687 def test_ifilter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000688 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000689 for g in (G, I, Ig, S, L, R):
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000690 self.assertEqual(list(ifilter(isEven, g(s))),
691 [x for x in g(s) if isEven(x)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000692 self.assertRaises(TypeError, ifilter, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000693 self.assertRaises(TypeError, ifilter, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000694 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
695
696 def test_ifilterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000697 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000698 for g in (G, I, Ig, S, L, R):
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000699 self.assertEqual(list(ifilterfalse(isEven, g(s))),
700 [x for x in g(s) if isOdd(x)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000701 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000702 self.assertRaises(TypeError, ifilterfalse, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000703 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
704
705 def test_izip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000706 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000707 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000708 self.assertEqual(list(izip(g(s))), lzip(g(s)))
709 self.assertEqual(list(izip(g(s), g(s))), lzip(g(s), g(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000710 self.assertRaises(TypeError, izip, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000711 self.assertRaises(TypeError, izip, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000712 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
713
Thomas Wouterscf297e42007-02-23 15:07:44 +0000714 def test_iziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000715 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000716 for g in (G, I, Ig, S, L, R):
717 self.assertEqual(list(izip_longest(g(s))), list(zip(g(s))))
718 self.assertEqual(list(izip_longest(g(s), g(s))), list(zip(g(s), g(s))))
719 self.assertRaises(TypeError, izip_longest, X(s))
720 self.assertRaises(TypeError, izip_longest, N(s))
721 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
722
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000723 def test_imap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000724 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000725 for g in (G, I, Ig, S, L, R):
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000726 self.assertEqual(list(imap(onearg, g(s))),
727 [onearg(x) for x in g(s)])
728 self.assertEqual(list(imap(operator.pow, g(s), g(s))),
729 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000730 self.assertRaises(TypeError, imap, onearg, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000731 self.assertRaises(TypeError, imap, onearg, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000732 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
733
734 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000735 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000736 for g in (G, I, Ig, S, L, R):
737 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
738 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000739 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000740 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
741
742 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000743 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000744 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000745 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000746 self.assertEqual(list(starmap(operator.pow, g(ss))),
747 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000748 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000749 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000750 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
751
752 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000753 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000754 for g in (G, I, Ig, S, L, R):
755 tgt = []
756 for elem in g(s):
757 if not isEven(elem): break
758 tgt.append(elem)
759 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
760 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000761 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000762 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
763
764 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000765 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000766 for g in (G, I, Ig, S, L, R):
767 tgt = []
768 for elem in g(s):
769 if not tgt and isOdd(elem): continue
770 tgt.append(elem)
771 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
772 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000773 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000774 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
775
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000776 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000777 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000778 for g in (G, I, Ig, S, L, R):
779 it1, it2 = tee(g(s))
780 self.assertEqual(list(it1), list(g(s)))
781 self.assertEqual(list(it2), list(g(s)))
782 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000783 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000784 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
785
Raymond Hettinger5cab2e32004-02-10 09:25:40 +0000786class LengthTransparency(unittest.TestCase):
787
788 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +0000789 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +0000790 self.assertEqual(len(repeat(None, 50)), 50)
791 self.assertRaises(TypeError, len, repeat(None))
792
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000793class RegressionTests(unittest.TestCase):
794
795 def test_sf_793826(self):
796 # Fix Armin Rigo's successful efforts to wreak havoc
797
798 def mutatingtuple(tuple1, f, tuple2):
799 # this builds a tuple t which is a copy of tuple1,
800 # then calls f(t), then mutates t to be equal to tuple2
801 # (needs len(tuple1) == len(tuple2)).
802 def g(value, first=[1]):
803 if first:
804 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +0000805 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000806 return value
807 items = list(tuple2)
808 items[1:1] = list(tuple1)
809 gen = imap(g, items)
810 z = izip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +0000811 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000812
813 def f(t):
814 global T
815 T = t
816 first[:] = list(T)
817
818 first = []
819 mutatingtuple((1,2,3), f, (4,5,6))
820 second = list(T)
821 self.assertEqual(first, second)
822
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000823
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000824 def test_sf_950057(self):
825 # Make sure that chain() and cycle() catch exceptions immediately
826 # rather than when shifting between input sources
827
828 def gen1():
829 hist.append(0)
830 yield 1
831 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +0000832 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000833 hist.append(2)
834
835 def gen2(x):
836 hist.append(3)
837 yield 2
838 hist.append(4)
839 if x:
840 raise StopIteration
841
842 hist = []
843 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
844 self.assertEqual(hist, [0,1])
845
846 hist = []
847 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
848 self.assertEqual(hist, [0,1])
849
850 hist = []
851 self.assertRaises(AssertionError, list, cycle(gen1()))
852 self.assertEqual(hist, [0,1])
853
Thomas Woutersb2137042007-02-01 18:02:27 +0000854class SubclassWithKwargsTest(unittest.TestCase):
855 def test_keywords_in_subclass(self):
856 # count is not subclassable...
857 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
858 starmap, islice, takewhile, dropwhile, cycle):
859 class Subclass(cls):
860 def __init__(self, newarg=None, *args):
861 cls.__init__(self, *args)
862 try:
863 Subclass(newarg=1)
864 except TypeError as err:
865 # we expect type errors because of wrong argument count
866 self.failIf("does not take keyword arguments" in err.args[0])
867
868
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000869libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000870
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000871
872>>> amounts = [120.15, 764.05, 823.14]
873>>> for checknum, amount in izip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +0000874... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000875...
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000876Check 1200 is for $120.15
877Check 1201 is for $764.05
878Check 1202 is for $823.14
879
880>>> import operator
Guido van Rossum805365e2007-05-07 22:24:25 +0000881>>> for cube in imap(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +0000882... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000883...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00008841
8858
88627
887
888>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +0000889>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +0000890... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +0000891...
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000892Alex
893Laura
894Martin
895Walter
896Samuele
897
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000898>>> from operator import itemgetter
899>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +0000900>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000901>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000902... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000903...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00009041 ['a', 'c', 'e']
9052 ['b', 'd', 'f']
9063 ['g']
907
Raymond Hettinger734fb572004-01-20 20:04:40 +0000908# Find runs of consecutive numbers using groupby. The key to the solution
909# is differencing with a range so that consecutive numbers all appear in
910# same group.
911>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +0000912>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000913... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000914...
Raymond Hettinger734fb572004-01-20 20:04:40 +0000915[1]
916[4, 5, 6]
917[10]
918[15, 16, 17, 18]
919[22]
920[25, 26, 27, 28]
921
Raymond Hettingera098b332003-09-08 23:58:40 +0000922>>> def take(n, seq):
923... return list(islice(seq, n))
924
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000925>>> def enumerate(iterable):
926... return izip(count(), iterable)
927
928>>> def tabulate(function):
929... "Return function(0), function(1), ..."
930... return imap(function, count())
931
932>>> def iteritems(mapping):
Guido van Rossumcc2b0162007-02-11 06:12:03 +0000933... return izip(mapping.keys(), mapping.values())
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000934
935>>> def nth(iterable, n):
936... "Returns the nth item"
Raymond Hettinger60eca932003-02-09 06:40:58 +0000937... return list(islice(iterable, n, n+1))
938
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000939>>> def all(seq, pred=None):
940... "Returns True if pred(x) is true for every element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000941... for elem in ifilterfalse(pred, seq):
942... return False
943... return True
Raymond Hettinger60eca932003-02-09 06:40:58 +0000944
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000945>>> def any(seq, pred=None):
946... "Returns True if pred(x) is true for at least one element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000947... for elem in ifilter(pred, seq):
948... return True
949... return False
Raymond Hettinger60eca932003-02-09 06:40:58 +0000950
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000951>>> def no(seq, pred=None):
952... "Returns True if pred(x) is false for every element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000953... for elem in ifilter(pred, seq):
954... return False
955... return True
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000956
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000957>>> def quantify(seq, pred=None):
958... "Count how many times the predicate is true in the sequence"
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000959... return sum(imap(pred, seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000960
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000961>>> def padnone(seq):
962... "Returns the sequence elements and then returns None indefinitely"
963... return chain(seq, repeat(None))
964
965>>> def ncycles(seq, n):
966... "Returns the sequence elements n times"
967... return chain(*repeat(seq, n))
968
969>>> def dotproduct(vec1, vec2):
970... return sum(imap(operator.mul, vec1, vec2))
971
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000972>>> def flatten(listOfLists):
973... return list(chain(*listOfLists))
974
975>>> def repeatfunc(func, times=None, *args):
976... "Repeat calls to func with specified arguments."
977... " Example: repeatfunc(random.random)"
978... if times is None:
979... return starmap(func, repeat(args))
980... else:
981... return starmap(func, repeat(args, times))
982
Raymond Hettingerd591f662003-10-26 15:34:50 +0000983>>> def pairwise(iterable):
984... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
985... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000986... try:
Georg Brandla18af4e2007-04-21 15:47:16 +0000987... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000988... except StopIteration:
989... pass
990... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000991
992This is not part of the examples but it tests to make sure the definitions
993perform as purported.
994
Raymond Hettingera098b332003-09-08 23:58:40 +0000995>>> take(10, count())
996[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
997
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000998>>> list(enumerate('abc'))
999[(0, 'a'), (1, 'b'), (2, 'c')]
1000
1001>>> list(islice(tabulate(lambda x: 2*x), 4))
1002[0, 2, 4, 6]
1003
1004>>> nth('abcde', 3)
1005['d']
1006
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001007>>> all([2, 4, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001008True
1009
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001010>>> all([2, 3, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001011False
1012
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001013>>> any([2, 4, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001014True
1015
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001016>>> any([1, 3, 5, 9], lambda x: x%2==0,)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001017False
1018
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001019>>> no([1, 3, 5, 9], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001020True
1021
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001022>>> no([1, 2, 5, 9], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001023False
1024
Guido van Rossum805365e2007-05-07 22:24:25 +00001025>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000102650
1027
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001028>>> a = [[1, 2, 3], [4, 5, 6]]
1029>>> flatten(a)
1030[1, 2, 3, 4, 5, 6]
1031
1032>>> list(repeatfunc(pow, 5, 2, 3))
1033[8, 8, 8, 8, 8]
1034
1035>>> import random
1036>>> take(5, imap(int, repeatfunc(random.random)))
1037[0, 0, 0, 0, 0]
1038
Raymond Hettingerd591f662003-10-26 15:34:50 +00001039>>> list(pairwise('abcd'))
1040[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001041
Raymond Hettingerd591f662003-10-26 15:34:50 +00001042>>> list(pairwise([]))
1043[]
1044
1045>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001046[]
1047
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001048>>> list(islice(padnone('abc'), 0, 6))
1049['a', 'b', 'c', None, None, None]
1050
1051>>> list(ncycles('abc', 3))
1052['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1053
1054>>> dotproduct([1,2,3], [4,5,6])
105532
1056
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001057"""
1058
1059__test__ = {'libreftest' : libreftest}
1060
1061def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001062 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00001063 RegressionTests, LengthTransparency,
1064 SubclassWithKwargsTest)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001065 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001066
1067 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001068 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001069 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001070 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00001071 for i in range(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001072 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001073 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001074 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001075 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001076
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001077 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001078 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001079
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001080if __name__ == "__main__":
1081 test_main(verbose=True)