blob: c2a539c30a35c692b5a2d1fc26ee14169b63311b [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
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00008
Guido van Rossum801f0d72006-08-24 19:48:10 +00009def lzip(*args):
10 return list(zip(*args))
11
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000012def onearg(x):
13 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000014 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000015
16def errfunc(*args):
17 'Test function that raises an error'
18 raise ValueError
19
20def gen3():
21 'Non-restartable source sequence'
22 for i in (0, 1, 2):
23 yield i
24
25def isEven(x):
26 'Test predicate'
27 return x%2==0
28
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000029def isOdd(x):
30 'Test predicate'
31 return x%2==1
32
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000033class StopNow:
34 'Class emulating an empty iterable.'
35 def __iter__(self):
36 return self
Georg Brandla18af4e2007-04-21 15:47:16 +000037 def __next__(self):
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000038 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000039
Raymond Hettinger02420702003-06-29 20:36:23 +000040def take(n, seq):
41 'Convenience function for partially consuming a long of infinite iterable'
42 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000043
Raymond Hettinger96ef8112003-02-01 00:10:11 +000044class TestBasicOps(unittest.TestCase):
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000045 def test_chain(self):
46 self.assertEqual(list(chain('abc', 'def')), list('abcdef'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000047 self.assertEqual(list(chain('abc')), list('abc'))
48 self.assertEqual(list(chain('')), [])
Raymond Hettinger02420702003-06-29 20:36:23 +000049 self.assertEqual(take(4, chain('abc', 'def')), list('abcd'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000050 self.assertRaises(TypeError, chain, 2, 3)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000051
Raymond Hettinger96ef8112003-02-01 00:10:11 +000052 def test_count(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +000053 self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
54 self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
55 self.assertEqual(take(2, lzip('abc',count(3))), [('a', 3), ('b', 4)])
Raymond Hettinger96ef8112003-02-01 00:10:11 +000056 self.assertRaises(TypeError, count, 2, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000057 self.assertRaises(TypeError, count, 'a')
Thomas Wouterscf297e42007-02-23 15:07:44 +000058 self.assertRaises(OverflowError, list, islice(count(sys.maxint-5), 10))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +000059 c = count(3)
60 self.assertEqual(repr(c), 'count(3)')
Georg Brandla18af4e2007-04-21 15:47:16 +000061 next(c)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +000062 self.assertEqual(repr(c), 'count(4)')
Thomas Wouters89f507f2006-12-13 04:49:30 +000063 c = count(-9)
64 self.assertEqual(repr(c), 'count(-9)')
Georg Brandla18af4e2007-04-21 15:47:16 +000065 next(c)
66 self.assertEqual(next(c), -8)
Raymond Hettinger96ef8112003-02-01 00:10:11 +000067
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000068 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +000069 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000070 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000071 self.assertRaises(TypeError, cycle)
72 self.assertRaises(TypeError, cycle, 5)
73 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000074
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000075 def test_groupby(self):
76 # Check whether it accepts arguments correctly
77 self.assertEqual([], list(groupby([])))
78 self.assertEqual([], list(groupby([], key=id)))
79 self.assertRaises(TypeError, list, groupby('abc', []))
80 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +000081 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000082
83 # Check normal input
84 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
85 (2,15,22), (3,16,23), (3,17,23)]
86 dup = []
87 for k, g in groupby(s, lambda r:r[0]):
88 for elem in g:
89 self.assertEqual(k, elem[0])
90 dup.append(elem)
91 self.assertEqual(s, dup)
92
93 # Check nested case
94 dup = []
95 for k, g in groupby(s, lambda r:r[0]):
96 for ik, ig in groupby(g, lambda r:r[2]):
97 for elem in ig:
98 self.assertEqual(k, elem[0])
99 self.assertEqual(ik, elem[2])
100 dup.append(elem)
101 self.assertEqual(s, dup)
102
103 # Check case where inner iterator is not used
104 keys = [k for k, g in groupby(s, lambda r:r[0])]
105 expectedkeys = set([r[0] for r in s])
106 self.assertEqual(set(keys), expectedkeys)
107 self.assertEqual(len(keys), len(expectedkeys))
108
109 # Exercise pipes and filters style
110 s = 'abracadabra'
111 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000112 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000113 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
114 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000115 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000116 self.assertEqual(r, ['a', 'b', 'r'])
117 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000118 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000119 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
120 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000121 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000122 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
123
Georg Brandla18af4e2007-04-21 15:47:16 +0000124 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000125 class ExpectedError(Exception):
126 pass
127 def delayed_raise(n=0):
128 for i in range(n):
129 yield 'yo'
130 raise ExpectedError
131 def gulp(iterable, keyp=None, func=list):
132 return [func(g) for k, g in groupby(iterable, keyp)]
133
Georg Brandla18af4e2007-04-21 15:47:16 +0000134 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000135 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000136 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000137 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
138
139 # __cmp__ failure
140 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000141 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000142 raise ExpectedError
143 s = [DummyCmp(), DummyCmp(), None]
144
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000145 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000146 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000147 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000148 self.assertRaises(ExpectedError, gulp, s)
149
150 # keyfunc failure
151 def keyfunc(obj):
152 if keyfunc.skip > 0:
153 keyfunc.skip -= 1
154 return obj
155 else:
156 raise ExpectedError
157
158 # keyfunc failure on outer object
159 keyfunc.skip = 0
160 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
161 keyfunc.skip = 1
162 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
163
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000164 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000165 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000166 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000167 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000168 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000169 self.assertRaises(TypeError, ifilter, lambda x:x)
170 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000171 self.assertRaises(TypeError, ifilter, isEven, 3)
Georg Brandla18af4e2007-04-21 15:47:16 +0000172 self.assertRaises(TypeError, next, ifilter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000173
174 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000175 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
176 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000177 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000178 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000179 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
180 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000181 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Georg Brandla18af4e2007-04-21 15:47:16 +0000182 self.assertRaises(TypeError, next, ifilterfalse(range(6), range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000183
184 def test_izip(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000185 # XXX This is rather silly now that builtin zip() calls izip()...
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000186 ans = [(x,y) for x, y in izip('abc',count())]
187 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Guido van Rossum801f0d72006-08-24 19:48:10 +0000188 self.assertEqual(list(izip('abc', range(6))), lzip('abc', range(6)))
189 self.assertEqual(list(izip('abcdef', range(3))), lzip('abcdef', range(3)))
190 self.assertEqual(take(3,izip('abcdef', count())), lzip('abcdef', range(3)))
191 self.assertEqual(list(izip('abcdef')), lzip('abcdef'))
192 self.assertEqual(list(izip()), lzip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000193 self.assertRaises(TypeError, izip, 3)
194 self.assertRaises(TypeError, izip, range(3), 3)
195 # Check tuple re-use (implementation detail)
196 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000197 lzip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000198 self.assertEqual([pair for pair in izip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000199 lzip('abc', 'def'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000200 ids = map(id, izip('abc', 'def'))
201 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000202 ids = map(id, list(izip('abc', 'def')))
203 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000204
Thomas Wouterscf297e42007-02-23 15:07:44 +0000205 def test_iziplongest(self):
206 for args in [
207 ['abc', range(6)],
208 [range(6), 'abc'],
209 [range(1000), range(2000,2100), range(3000,3050)],
210 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
211 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
212 ]:
213 target = map(None, *args)
214 self.assertEqual(list(izip_longest(*args)), target)
215 self.assertEqual(list(izip_longest(*args, **{})), target)
216 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
217 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000218
Thomas Wouterscf297e42007-02-23 15:07:44 +0000219 self.assertEqual(take(3,izip_longest('abcdef', count())), list(zip('abcdef', range(3)))) # take 3 from infinite input
220
221 self.assertEqual(list(izip_longest()), list(zip()))
222 self.assertEqual(list(izip_longest([])), list(zip([])))
223 self.assertEqual(list(izip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000224
Thomas Wouterscf297e42007-02-23 15:07:44 +0000225 self.assertEqual(list(izip_longest('abc', 'defg', **{})), map(None, 'abc', 'defg')) # empty keyword dict
226 self.assertRaises(TypeError, izip_longest, 3)
227 self.assertRaises(TypeError, izip_longest, range(3), 3)
228
229 for stmt in [
230 "izip_longest('abc', fv=1)",
Guido van Rossumd8faa362007-04-27 19:54:29 +0000231 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000232 ]:
233 try:
234 eval(stmt, globals(), locals())
235 except TypeError:
236 pass
237 else:
238 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000239
Thomas Wouterscf297e42007-02-23 15:07:44 +0000240 # Check tuple re-use (implementation detail)
241 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
242 list(zip('abc', 'def')))
243 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
244 list(zip('abc', 'def')))
245 ids = map(id, izip_longest('abc', 'def'))
246 self.assertEqual(min(ids), max(ids))
247 ids = map(id, list(izip_longest('abc', 'def')))
248 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
249
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000250 def test_repeat(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000251 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000252 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000253 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000254 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000255 self.assertEqual(list(repeat('a', 0)), [])
256 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000257 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000258 self.assertRaises(TypeError, repeat, None, 3, 4)
259 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000260 r = repeat(1+0j)
261 self.assertEqual(repr(r), 'repeat((1+0j))')
262 r = repeat(1+0j, 5)
263 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
264 list(r)
265 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000266
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000267 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000268 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
269 [0**1, 1**2, 2**3])
270 self.assertEqual(list(imap(None, 'abc', range(5))),
271 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000272 self.assertEqual(list(imap(None, 'abc', count())),
273 [('a',0),('b',1),('c',2)])
274 self.assertEqual(take(2,imap(None, 'abc', count())),
275 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000276 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000277 self.assertRaises(TypeError, imap)
278 self.assertRaises(TypeError, imap, operator.neg)
Georg Brandla18af4e2007-04-21 15:47:16 +0000279 self.assertRaises(TypeError, next, imap(10, range(5)))
280 self.assertRaises(ValueError, next, imap(errfunc, [4], [5]))
281 self.assertRaises(TypeError, next, imap(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000282
283 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000284 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
285 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000286 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
287 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000288 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger2012f172003-02-07 05:32:58 +0000289 self.assertRaises(TypeError, list, starmap(operator.pow, [[4,5]]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000290 self.assertRaises(TypeError, starmap)
291 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000292 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
293 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
294 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000295
296 def test_islice(self):
297 for args in [ # islice(args) should agree with range(args)
298 (10, 20, 3),
299 (10, 3, 20),
300 (10, 20),
301 (10, 3),
302 (20,)
303 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000304 self.assertEqual(list(islice(range(100), *args)),
305 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000306
307 for args, tgtargs in [ # Stop when seqn is exhausted
308 ((10, 110, 3), ((10, 100, 3))),
309 ((10, 110), ((10, 100))),
310 ((110,), (100,))
311 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000312 self.assertEqual(list(islice(range(100), *args)),
313 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000314
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000315 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +0000316 self.assertEqual(list(islice(range(10), None)), list(range(10)))
317 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
318 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
319 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
320 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000321
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000322 # Test number of items consumed SF #1171417
323 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +0000324 self.assertEqual(list(islice(it, 3)), list(range(3)))
325 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000326
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000327 # Test invalid arguments
Guido van Rossum805365e2007-05-07 22:24:25 +0000328 self.assertRaises(TypeError, islice, range(10))
329 self.assertRaises(TypeError, islice, range(10), 1, 2, 3, 4)
330 self.assertRaises(ValueError, islice, range(10), -5, 10, 1)
331 self.assertRaises(ValueError, islice, range(10), 1, -5, -1)
332 self.assertRaises(ValueError, islice, range(10), 1, 10, -1)
333 self.assertRaises(ValueError, islice, range(10), 1, 10, 0)
334 self.assertRaises(ValueError, islice, range(10), 'a')
335 self.assertRaises(ValueError, islice, range(10), 'a', 1)
336 self.assertRaises(ValueError, islice, range(10), 1, 'a')
337 self.assertRaises(ValueError, islice, range(10), 'a', 1, 1)
338 self.assertRaises(ValueError, islice, range(10), 1, 'a', 1)
Raymond Hettinger2012f172003-02-07 05:32:58 +0000339 self.assertEqual(len(list(islice(count(), 1, 10, sys.maxint))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000340
341 def test_takewhile(self):
342 data = [1, 3, 5, 20, 2, 4, 6, 8]
343 underten = lambda x: x<10
344 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000345 self.assertEqual(list(takewhile(underten, [])), [])
346 self.assertRaises(TypeError, takewhile)
347 self.assertRaises(TypeError, takewhile, operator.pow)
348 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000349 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
350 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000351 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
352 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +0000353 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000354
355 def test_dropwhile(self):
356 data = [1, 3, 5, 20, 2, 4, 6, 8]
357 underten = lambda x: x<10
358 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000359 self.assertEqual(list(dropwhile(underten, [])), [])
360 self.assertRaises(TypeError, dropwhile)
361 self.assertRaises(TypeError, dropwhile, operator.pow)
362 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000363 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
364 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000365
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000366 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000367 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000368 def irange(n):
Guido van Rossum805365e2007-05-07 22:24:25 +0000369 for i in range(n):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000370 yield i
371
372 a, b = tee([]) # test empty iterator
373 self.assertEqual(list(a), [])
374 self.assertEqual(list(b), [])
375
376 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +0000377 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000378
379 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +0000380 self.assertEqual(list(a), list(range(n)))
381 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000382
383 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000384 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000385 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000386 del a
Guido van Rossum805365e2007-05-07 22:24:25 +0000387 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000388
389 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000390 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000391 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000392 del b
Guido van Rossum805365e2007-05-07 22:24:25 +0000393 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000394
Guido van Rossum805365e2007-05-07 22:24:25 +0000395 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000396 order = [0]*n + [1]*n
397 random.shuffle(order)
398 lists = ([], [])
399 its = tee(irange(n))
400 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +0000401 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000402 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +0000403 self.assertEqual(lists[0], list(range(n)))
404 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000405
Raymond Hettingerad983e72003-11-12 14:32:26 +0000406 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000407 self.assertRaises(TypeError, tee)
408 self.assertRaises(TypeError, tee, 3)
409 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000410 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000411
Raymond Hettingerad983e72003-11-12 14:32:26 +0000412 # tee object should be instantiable
413 a, b = tee('abc')
414 c = type(a)('def')
415 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000416
Raymond Hettingerad983e72003-11-12 14:32:26 +0000417 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +0000418 a, b, c = tee(range(2000), 3)
419 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000420 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +0000421 self.assertEqual(list(b), list(range(2000)))
422 self.assertEqual([next(c), next(c)], list(range(2)))
423 self.assertEqual(list(a), list(range(100,2000)))
424 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000425
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000426 # test values of n
427 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000428 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +0000429 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000430 result = tee('abc', n)
431 self.assertEqual(type(result), tuple)
432 self.assertEqual(len(result), n)
433 self.assertEqual(map(list, result), [list('abc')]*n)
434
Raymond Hettingerad983e72003-11-12 14:32:26 +0000435 # tee pass-through to copyable iterator
436 a, b = tee('abc')
437 c, d = tee(a)
438 self.assert_(a is c)
439
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000440 # test tee_new
441 t1, t2 = tee('abc')
442 tnew = type(t1)
443 self.assertRaises(TypeError, tnew)
444 self.assertRaises(TypeError, tnew, 10)
445 t3 = tnew(t1)
446 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000447
Raymond Hettingera9f60922004-10-17 16:40:14 +0000448 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +0000449 a, b = tee(range(10))
Raymond Hettingera9f60922004-10-17 16:40:14 +0000450 p = proxy(a)
451 self.assertEqual(getattr(p, '__class__'), type(b))
452 del a
453 self.assertRaises(ReferenceError, getattr, p, '__class__')
454
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000455 def test_StopIteration(self):
Georg Brandla18af4e2007-04-21 15:47:16 +0000456 self.assertRaises(StopIteration, next, izip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000457
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000458 for f in (chain, cycle, izip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +0000459 self.assertRaises(StopIteration, next, f([]))
460 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000461
Georg Brandla18af4e2007-04-21 15:47:16 +0000462 self.assertRaises(StopIteration, next, islice([], None))
463 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000464
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000465 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +0000466 self.assertRaises(StopIteration, next, p)
467 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000468 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +0000469 self.assertRaises(StopIteration, next, p)
470 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000471
Georg Brandla18af4e2007-04-21 15:47:16 +0000472 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000473
474 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +0000475 self.assertRaises(StopIteration, next, f(lambda x:x, []))
476 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000477
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000478class TestGC(unittest.TestCase):
479
480 def makecycle(self, iterator, container):
481 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +0000482 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000483 del container, iterator
484
485 def test_chain(self):
486 a = []
487 self.makecycle(chain(a), a)
488
489 def test_cycle(self):
490 a = []
491 self.makecycle(cycle([a]*2), a)
492
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000493 def test_dropwhile(self):
494 a = []
495 self.makecycle(dropwhile(bool, [0, a, a]), a)
496
497 def test_groupby(self):
498 a = []
499 self.makecycle(groupby([a]*2, lambda x:x), a)
500
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000501 def test_ifilter(self):
502 a = []
503 self.makecycle(ifilter(lambda x:True, [a]*2), a)
504
505 def test_ifilterfalse(self):
506 a = []
507 self.makecycle(ifilterfalse(lambda x:False, a), a)
508
509 def test_izip(self):
510 a = []
511 self.makecycle(izip([a]*2, [a]*3), a)
512
513 def test_imap(self):
514 a = []
515 self.makecycle(imap(lambda x:x, [a]*2), a)
516
517 def test_islice(self):
518 a = []
519 self.makecycle(islice([a]*2, None), a)
520
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000521 def test_repeat(self):
522 a = []
523 self.makecycle(repeat(a), a)
524
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000525 def test_starmap(self):
526 a = []
527 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
528
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000529 def test_takewhile(self):
530 a = []
531 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
532
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000533def R(seqn):
534 'Regular generator'
535 for i in seqn:
536 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000537
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000538class G:
539 'Sequence using __getitem__'
540 def __init__(self, seqn):
541 self.seqn = seqn
542 def __getitem__(self, i):
543 return self.seqn[i]
544
545class I:
546 'Sequence using iterator protocol'
547 def __init__(self, seqn):
548 self.seqn = seqn
549 self.i = 0
550 def __iter__(self):
551 return self
Georg Brandla18af4e2007-04-21 15:47:16 +0000552 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000553 if self.i >= len(self.seqn): raise StopIteration
554 v = self.seqn[self.i]
555 self.i += 1
556 return v
557
558class Ig:
559 'Sequence using iterator protocol defined with a generator'
560 def __init__(self, seqn):
561 self.seqn = seqn
562 self.i = 0
563 def __iter__(self):
564 for val in self.seqn:
565 yield val
566
567class X:
568 'Missing __getitem__ and __iter__'
569 def __init__(self, seqn):
570 self.seqn = seqn
571 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +0000572 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000573 if self.i >= len(self.seqn): raise StopIteration
574 v = self.seqn[self.i]
575 self.i += 1
576 return v
577
578class N:
Georg Brandla18af4e2007-04-21 15:47:16 +0000579 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000580 def __init__(self, seqn):
581 self.seqn = seqn
582 self.i = 0
583 def __iter__(self):
584 return self
585
586class E:
587 'Test propagation of exceptions'
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 Hettingerffdb8bb2004-09-27 15:29:05 +0000594 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000595
596class S:
597 'Test immediate stop'
598 def __init__(self, seqn):
599 pass
600 def __iter__(self):
601 return self
Georg Brandla18af4e2007-04-21 15:47:16 +0000602 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000603 raise StopIteration
604
605def L(seqn):
606 'Test multiple tiers of iterators'
607 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
608
609
610class TestVariousIteratorArgs(unittest.TestCase):
611
612 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000613 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000614 for g in (G, I, Ig, S, L, R):
615 self.assertEqual(list(chain(g(s))), list(g(s)))
616 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
617 self.assertRaises(TypeError, chain, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000618 self.assertRaises(TypeError, chain, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000619 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
620
621 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000622 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000623 for g in (G, I, Ig, S, L, R):
624 tgtlen = len(s) * 3
625 expected = list(g(s))*3
626 actual = list(islice(cycle(g(s)), tgtlen))
627 self.assertEqual(actual, expected)
628 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000629 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000630 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
631
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000632 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000633 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000634 for g in (G, I, Ig, S, L, R):
635 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
636 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000637 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000638 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
639
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000640 def test_ifilter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000641 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000642 for g in (G, I, Ig, S, L, R):
643 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
644 self.assertRaises(TypeError, ifilter, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000645 self.assertRaises(TypeError, ifilter, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000646 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
647
648 def test_ifilterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000649 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000650 for g in (G, I, Ig, S, L, R):
651 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
652 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000653 self.assertRaises(TypeError, ifilterfalse, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000654 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
655
656 def test_izip(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):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000659 self.assertEqual(list(izip(g(s))), lzip(g(s)))
660 self.assertEqual(list(izip(g(s), g(s))), lzip(g(s), g(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000661 self.assertRaises(TypeError, izip, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000662 self.assertRaises(TypeError, izip, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000663 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
664
Thomas Wouterscf297e42007-02-23 15:07:44 +0000665 def test_iziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000666 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000667 for g in (G, I, Ig, S, L, R):
668 self.assertEqual(list(izip_longest(g(s))), list(zip(g(s))))
669 self.assertEqual(list(izip_longest(g(s), g(s))), list(zip(g(s), g(s))))
670 self.assertRaises(TypeError, izip_longest, X(s))
671 self.assertRaises(TypeError, izip_longest, N(s))
672 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
673
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000674 def test_imap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000675 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000676 for g in (G, I, Ig, S, L, R):
677 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
678 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
679 self.assertRaises(TypeError, imap, onearg, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000680 self.assertRaises(TypeError, imap, onearg, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000681 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
682
683 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000684 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000685 for g in (G, I, Ig, S, L, R):
686 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
687 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000688 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000689 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
690
691 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000692 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000693 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000694 ss = lzip(s, s)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000695 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
696 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000697 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000698 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
699
700 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000701 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000702 for g in (G, I, Ig, S, L, R):
703 tgt = []
704 for elem in g(s):
705 if not isEven(elem): break
706 tgt.append(elem)
707 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
708 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000709 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000710 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
711
712 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000713 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000714 for g in (G, I, Ig, S, L, R):
715 tgt = []
716 for elem in g(s):
717 if not tgt and isOdd(elem): continue
718 tgt.append(elem)
719 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
720 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000721 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000722 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
723
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000724 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000725 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000726 for g in (G, I, Ig, S, L, R):
727 it1, it2 = tee(g(s))
728 self.assertEqual(list(it1), list(g(s)))
729 self.assertEqual(list(it2), list(g(s)))
730 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000731 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000732 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
733
Raymond Hettinger5cab2e32004-02-10 09:25:40 +0000734class LengthTransparency(unittest.TestCase):
735
736 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +0000737 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +0000738 self.assertEqual(len(repeat(None, 50)), 50)
739 self.assertRaises(TypeError, len, repeat(None))
740
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000741class RegressionTests(unittest.TestCase):
742
743 def test_sf_793826(self):
744 # Fix Armin Rigo's successful efforts to wreak havoc
745
746 def mutatingtuple(tuple1, f, tuple2):
747 # this builds a tuple t which is a copy of tuple1,
748 # then calls f(t), then mutates t to be equal to tuple2
749 # (needs len(tuple1) == len(tuple2)).
750 def g(value, first=[1]):
751 if first:
752 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +0000753 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000754 return value
755 items = list(tuple2)
756 items[1:1] = list(tuple1)
757 gen = imap(g, items)
758 z = izip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +0000759 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000760
761 def f(t):
762 global T
763 T = t
764 first[:] = list(T)
765
766 first = []
767 mutatingtuple((1,2,3), f, (4,5,6))
768 second = list(T)
769 self.assertEqual(first, second)
770
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000771
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000772 def test_sf_950057(self):
773 # Make sure that chain() and cycle() catch exceptions immediately
774 # rather than when shifting between input sources
775
776 def gen1():
777 hist.append(0)
778 yield 1
779 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +0000780 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000781 hist.append(2)
782
783 def gen2(x):
784 hist.append(3)
785 yield 2
786 hist.append(4)
787 if x:
788 raise StopIteration
789
790 hist = []
791 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
792 self.assertEqual(hist, [0,1])
793
794 hist = []
795 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
796 self.assertEqual(hist, [0,1])
797
798 hist = []
799 self.assertRaises(AssertionError, list, cycle(gen1()))
800 self.assertEqual(hist, [0,1])
801
Thomas Woutersb2137042007-02-01 18:02:27 +0000802class SubclassWithKwargsTest(unittest.TestCase):
803 def test_keywords_in_subclass(self):
804 # count is not subclassable...
805 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
806 starmap, islice, takewhile, dropwhile, cycle):
807 class Subclass(cls):
808 def __init__(self, newarg=None, *args):
809 cls.__init__(self, *args)
810 try:
811 Subclass(newarg=1)
812 except TypeError as err:
813 # we expect type errors because of wrong argument count
814 self.failIf("does not take keyword arguments" in err.args[0])
815
816
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000817libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000818
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000819
820>>> amounts = [120.15, 764.05, 823.14]
821>>> for checknum, amount in izip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +0000822... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000823...
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000824Check 1200 is for $120.15
825Check 1201 is for $764.05
826Check 1202 is for $823.14
827
828>>> import operator
Guido van Rossum805365e2007-05-07 22:24:25 +0000829>>> for cube in imap(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +0000830... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000831...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00008321
8338
83427
835
836>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +0000837>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +0000838... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +0000839...
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000840Alex
841Laura
842Martin
843Walter
844Samuele
845
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000846>>> from operator import itemgetter
847>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +0000848>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000849>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossum7131f842007-02-09 20:13:25 +0000850... print(k, map(itemgetter(0), g))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000851...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00008521 ['a', 'c', 'e']
8532 ['b', 'd', 'f']
8543 ['g']
855
Raymond Hettinger734fb572004-01-20 20:04:40 +0000856# Find runs of consecutive numbers using groupby. The key to the solution
857# is differencing with a range so that consecutive numbers all appear in
858# same group.
859>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
860>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
Guido van Rossum7131f842007-02-09 20:13:25 +0000861... print(map(operator.itemgetter(1), g))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000862...
Raymond Hettinger734fb572004-01-20 20:04:40 +0000863[1]
864[4, 5, 6]
865[10]
866[15, 16, 17, 18]
867[22]
868[25, 26, 27, 28]
869
Raymond Hettingera098b332003-09-08 23:58:40 +0000870>>> def take(n, seq):
871... return list(islice(seq, n))
872
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000873>>> def enumerate(iterable):
874... return izip(count(), iterable)
875
876>>> def tabulate(function):
877... "Return function(0), function(1), ..."
878... return imap(function, count())
879
880>>> def iteritems(mapping):
Guido van Rossumcc2b0162007-02-11 06:12:03 +0000881... return izip(mapping.keys(), mapping.values())
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000882
883>>> def nth(iterable, n):
884... "Returns the nth item"
Raymond Hettinger60eca932003-02-09 06:40:58 +0000885... return list(islice(iterable, n, n+1))
886
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000887>>> def all(seq, pred=None):
888... "Returns True if pred(x) is true for every element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000889... for elem in ifilterfalse(pred, seq):
890... return False
891... return True
Raymond Hettinger60eca932003-02-09 06:40:58 +0000892
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000893>>> def any(seq, pred=None):
894... "Returns True if pred(x) is true for at least one element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000895... for elem in ifilter(pred, seq):
896... return True
897... return False
Raymond Hettinger60eca932003-02-09 06:40:58 +0000898
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000899>>> def no(seq, pred=None):
900... "Returns True if pred(x) is false for every element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000901... for elem in ifilter(pred, seq):
902... return False
903... return True
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000904
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000905>>> def quantify(seq, pred=None):
906... "Count how many times the predicate is true in the sequence"
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000907... return sum(imap(pred, seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000908
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000909>>> def padnone(seq):
910... "Returns the sequence elements and then returns None indefinitely"
911... return chain(seq, repeat(None))
912
913>>> def ncycles(seq, n):
914... "Returns the sequence elements n times"
915... return chain(*repeat(seq, n))
916
917>>> def dotproduct(vec1, vec2):
918... return sum(imap(operator.mul, vec1, vec2))
919
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000920>>> def flatten(listOfLists):
921... return list(chain(*listOfLists))
922
923>>> def repeatfunc(func, times=None, *args):
924... "Repeat calls to func with specified arguments."
925... " Example: repeatfunc(random.random)"
926... if times is None:
927... return starmap(func, repeat(args))
928... else:
929... return starmap(func, repeat(args, times))
930
Raymond Hettingerd591f662003-10-26 15:34:50 +0000931>>> def pairwise(iterable):
932... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
933... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000934... try:
Georg Brandla18af4e2007-04-21 15:47:16 +0000935... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000936... except StopIteration:
937... pass
938... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000939
940This is not part of the examples but it tests to make sure the definitions
941perform as purported.
942
Raymond Hettingera098b332003-09-08 23:58:40 +0000943>>> take(10, count())
944[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
945
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000946>>> list(enumerate('abc'))
947[(0, 'a'), (1, 'b'), (2, 'c')]
948
949>>> list(islice(tabulate(lambda x: 2*x), 4))
950[0, 2, 4, 6]
951
952>>> nth('abcde', 3)
953['d']
954
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000955>>> all([2, 4, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000956True
957
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000958>>> all([2, 3, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000959False
960
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000961>>> any([2, 4, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000962True
963
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000964>>> any([1, 3, 5, 9], lambda x: x%2==0,)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000965False
966
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000967>>> no([1, 3, 5, 9], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000968True
969
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000970>>> no([1, 2, 5, 9], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000971False
972
Guido van Rossum805365e2007-05-07 22:24:25 +0000973>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +000097450
975
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000976>>> a = [[1, 2, 3], [4, 5, 6]]
977>>> flatten(a)
978[1, 2, 3, 4, 5, 6]
979
980>>> list(repeatfunc(pow, 5, 2, 3))
981[8, 8, 8, 8, 8]
982
983>>> import random
984>>> take(5, imap(int, repeatfunc(random.random)))
985[0, 0, 0, 0, 0]
986
Raymond Hettingerd591f662003-10-26 15:34:50 +0000987>>> list(pairwise('abcd'))
988[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000989
Raymond Hettingerd591f662003-10-26 15:34:50 +0000990>>> list(pairwise([]))
991[]
992
993>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +0000994[]
995
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000996>>> list(islice(padnone('abc'), 0, 6))
997['a', 'b', 'c', None, None, None]
998
999>>> list(ncycles('abc', 3))
1000['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1001
1002>>> dotproduct([1,2,3], [4,5,6])
100332
1004
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001005"""
1006
1007__test__ = {'libreftest' : libreftest}
1008
1009def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001010 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00001011 RegressionTests, LengthTransparency,
1012 SubclassWithKwargsTest)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001013 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001014
1015 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001016 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001017 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001018 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00001019 for i in range(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001020 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001021 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001022 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001023 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001024
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001025 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001026 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001027
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001028if __name__ == "__main__":
1029 test_main(verbose=True)