blob: e3728d8d3325ac6dbeb2e4d6e5d8af052addc5d0 [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
Guido van Rossum360e4b82007-05-14 22:51:27 +00008maxsize = test_support.MAX_Py_ssize_t
9minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000010
Guido van Rossum801f0d72006-08-24 19:48:10 +000011def lzip(*args):
12 return list(zip(*args))
13
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000014def onearg(x):
15 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000016 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000017
18def errfunc(*args):
19 'Test function that raises an error'
20 raise ValueError
21
22def gen3():
23 'Non-restartable source sequence'
24 for i in (0, 1, 2):
25 yield i
26
27def isEven(x):
28 'Test predicate'
29 return x%2==0
30
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000031def isOdd(x):
32 'Test predicate'
33 return x%2==1
34
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000035class StopNow:
36 'Class emulating an empty iterable.'
37 def __iter__(self):
38 return self
Georg Brandla18af4e2007-04-21 15:47:16 +000039 def __next__(self):
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000040 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000041
Raymond Hettinger02420702003-06-29 20:36:23 +000042def take(n, seq):
43 'Convenience function for partially consuming a long of infinite iterable'
44 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000045
Raymond Hettinger96ef8112003-02-01 00:10:11 +000046class TestBasicOps(unittest.TestCase):
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000047 def test_chain(self):
48 self.assertEqual(list(chain('abc', 'def')), list('abcdef'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000049 self.assertEqual(list(chain('abc')), list('abc'))
50 self.assertEqual(list(chain('')), [])
Raymond Hettinger02420702003-06-29 20:36:23 +000051 self.assertEqual(take(4, chain('abc', 'def')), list('abcd'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000052 self.assertRaises(TypeError, chain, 2, 3)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000053
Raymond Hettinger96ef8112003-02-01 00:10:11 +000054 def test_count(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +000055 self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
56 self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
57 self.assertEqual(take(2, lzip('abc',count(3))), [('a', 3), ('b', 4)])
Raymond Hettinger96ef8112003-02-01 00:10:11 +000058 self.assertRaises(TypeError, count, 2, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000059 self.assertRaises(TypeError, count, 'a')
Guido van Rossum360e4b82007-05-14 22:51:27 +000060 self.assertRaises(OverflowError, list, islice(count(maxsize-5), 10))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +000061 c = count(3)
62 self.assertEqual(repr(c), 'count(3)')
Georg Brandla18af4e2007-04-21 15:47:16 +000063 next(c)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +000064 self.assertEqual(repr(c), 'count(4)')
Thomas Wouters89f507f2006-12-13 04:49:30 +000065 c = count(-9)
66 self.assertEqual(repr(c), 'count(-9)')
Georg Brandla18af4e2007-04-21 15:47:16 +000067 next(c)
68 self.assertEqual(next(c), -8)
Raymond Hettinger96ef8112003-02-01 00:10:11 +000069
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000070 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +000071 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000072 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000073 self.assertRaises(TypeError, cycle)
74 self.assertRaises(TypeError, cycle, 5)
75 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000076
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000077 def test_groupby(self):
78 # Check whether it accepts arguments correctly
79 self.assertEqual([], list(groupby([])))
80 self.assertEqual([], list(groupby([], key=id)))
81 self.assertRaises(TypeError, list, groupby('abc', []))
82 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +000083 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000084
85 # Check normal input
86 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
87 (2,15,22), (3,16,23), (3,17,23)]
88 dup = []
89 for k, g in groupby(s, lambda r:r[0]):
90 for elem in g:
91 self.assertEqual(k, elem[0])
92 dup.append(elem)
93 self.assertEqual(s, dup)
94
95 # Check nested case
96 dup = []
97 for k, g in groupby(s, lambda r:r[0]):
98 for ik, ig in groupby(g, lambda r:r[2]):
99 for elem in ig:
100 self.assertEqual(k, elem[0])
101 self.assertEqual(ik, elem[2])
102 dup.append(elem)
103 self.assertEqual(s, dup)
104
105 # Check case where inner iterator is not used
106 keys = [k for k, g in groupby(s, lambda r:r[0])]
107 expectedkeys = set([r[0] for r in s])
108 self.assertEqual(set(keys), expectedkeys)
109 self.assertEqual(len(keys), len(expectedkeys))
110
111 # Exercise pipes and filters style
112 s = 'abracadabra'
113 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000114 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000115 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
116 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000117 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000118 self.assertEqual(r, ['a', 'b', 'r'])
119 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000120 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000121 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
122 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000123 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000124 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
125
Georg Brandla18af4e2007-04-21 15:47:16 +0000126 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000127 class ExpectedError(Exception):
128 pass
129 def delayed_raise(n=0):
130 for i in range(n):
131 yield 'yo'
132 raise ExpectedError
133 def gulp(iterable, keyp=None, func=list):
134 return [func(g) for k, g in groupby(iterable, keyp)]
135
Georg Brandla18af4e2007-04-21 15:47:16 +0000136 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000137 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000138 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000139 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
140
141 # __cmp__ failure
142 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000143 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000144 raise ExpectedError
145 s = [DummyCmp(), DummyCmp(), None]
146
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000147 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000148 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000149 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000150 self.assertRaises(ExpectedError, gulp, s)
151
152 # keyfunc failure
153 def keyfunc(obj):
154 if keyfunc.skip > 0:
155 keyfunc.skip -= 1
156 return obj
157 else:
158 raise ExpectedError
159
160 # keyfunc failure on outer object
161 keyfunc.skip = 0
162 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
163 keyfunc.skip = 1
164 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
165
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000166 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000167 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000168 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000169 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000170 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000171 self.assertRaises(TypeError, ifilter, lambda x:x)
172 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000173 self.assertRaises(TypeError, ifilter, isEven, 3)
Georg Brandla18af4e2007-04-21 15:47:16 +0000174 self.assertRaises(TypeError, next, ifilter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000175
176 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000177 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
178 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000179 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000180 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000181 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
182 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000183 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Georg Brandla18af4e2007-04-21 15:47:16 +0000184 self.assertRaises(TypeError, next, ifilterfalse(range(6), range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000185
186 def test_izip(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000187 # XXX This is rather silly now that builtin zip() calls izip()...
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000188 ans = [(x,y) for x, y in izip('abc',count())]
189 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Guido van Rossum801f0d72006-08-24 19:48:10 +0000190 self.assertEqual(list(izip('abc', range(6))), lzip('abc', range(6)))
191 self.assertEqual(list(izip('abcdef', range(3))), lzip('abcdef', range(3)))
192 self.assertEqual(take(3,izip('abcdef', count())), lzip('abcdef', range(3)))
193 self.assertEqual(list(izip('abcdef')), lzip('abcdef'))
194 self.assertEqual(list(izip()), lzip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000195 self.assertRaises(TypeError, izip, 3)
196 self.assertRaises(TypeError, izip, range(3), 3)
197 # Check tuple re-use (implementation detail)
198 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000199 lzip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000200 self.assertEqual([pair for pair in izip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000201 lzip('abc', 'def'))
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000202 ids = list(map(id, izip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000203 self.assertEqual(min(ids), max(ids))
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000204 ids = list(map(id, list(izip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000205 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000206
Thomas Wouterscf297e42007-02-23 15:07:44 +0000207 def test_iziplongest(self):
208 for args in [
209 ['abc', range(6)],
210 [range(6), 'abc'],
211 [range(1000), range(2000,2100), range(3000,3050)],
212 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
213 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
214 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000215 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
216 for i in range(max(map(len, args)))]
Thomas Wouterscf297e42007-02-23 15:07:44 +0000217 self.assertEqual(list(izip_longest(*args)), target)
218 self.assertEqual(list(izip_longest(*args, **{})), target)
219 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
220 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000221
Thomas Wouterscf297e42007-02-23 15:07:44 +0000222 self.assertEqual(take(3,izip_longest('abcdef', count())), list(zip('abcdef', range(3)))) # take 3 from infinite input
223
224 self.assertEqual(list(izip_longest()), list(zip()))
225 self.assertEqual(list(izip_longest([])), list(zip([])))
226 self.assertEqual(list(izip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000227
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000228 self.assertEqual(list(izip_longest('abc', 'defg', **{})),
229 list(map(None, list('abc')+[None], 'defg'))) # empty keyword dict
Thomas Wouterscf297e42007-02-23 15:07:44 +0000230 self.assertRaises(TypeError, izip_longest, 3)
231 self.assertRaises(TypeError, izip_longest, range(3), 3)
232
233 for stmt in [
234 "izip_longest('abc', fv=1)",
Guido van Rossumd8faa362007-04-27 19:54:29 +0000235 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000236 ]:
237 try:
238 eval(stmt, globals(), locals())
239 except TypeError:
240 pass
241 else:
242 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000243
Thomas Wouterscf297e42007-02-23 15:07:44 +0000244 # Check tuple re-use (implementation detail)
245 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
246 list(zip('abc', 'def')))
247 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
248 list(zip('abc', 'def')))
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000249 ids = list(map(id, izip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000250 self.assertEqual(min(ids), max(ids))
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000251 ids = list(map(id, list(izip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000252 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
253
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000254 def test_repeat(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000255 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000256 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000257 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000258 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000259 self.assertEqual(list(repeat('a', 0)), [])
260 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000261 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000262 self.assertRaises(TypeError, repeat, None, 3, 4)
263 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000264 r = repeat(1+0j)
265 self.assertEqual(repr(r), 'repeat((1+0j))')
266 r = repeat(1+0j, 5)
267 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
268 list(r)
269 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000270
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000271 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000272 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
273 [0**1, 1**2, 2**3])
274 self.assertEqual(list(imap(None, 'abc', range(5))),
275 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000276 self.assertEqual(list(imap(None, 'abc', count())),
277 [('a',0),('b',1),('c',2)])
278 self.assertEqual(take(2,imap(None, 'abc', count())),
279 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000280 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000281 self.assertRaises(TypeError, imap)
282 self.assertRaises(TypeError, imap, operator.neg)
Georg Brandla18af4e2007-04-21 15:47:16 +0000283 self.assertRaises(TypeError, next, imap(10, range(5)))
284 self.assertRaises(ValueError, next, imap(errfunc, [4], [5]))
285 self.assertRaises(TypeError, next, imap(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000286
287 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000288 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
289 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000290 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
291 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000292 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger2012f172003-02-07 05:32:58 +0000293 self.assertRaises(TypeError, list, starmap(operator.pow, [[4,5]]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000294 self.assertRaises(TypeError, starmap)
295 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000296 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
297 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
298 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000299
300 def test_islice(self):
301 for args in [ # islice(args) should agree with range(args)
302 (10, 20, 3),
303 (10, 3, 20),
304 (10, 20),
305 (10, 3),
306 (20,)
307 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000308 self.assertEqual(list(islice(range(100), *args)),
309 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000310
311 for args, tgtargs in [ # Stop when seqn is exhausted
312 ((10, 110, 3), ((10, 100, 3))),
313 ((10, 110), ((10, 100))),
314 ((110,), (100,))
315 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000316 self.assertEqual(list(islice(range(100), *args)),
317 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000318
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000319 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +0000320 self.assertEqual(list(islice(range(10), None)), list(range(10)))
321 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
322 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
323 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
324 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000325
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000326 # Test number of items consumed SF #1171417
327 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +0000328 self.assertEqual(list(islice(it, 3)), list(range(3)))
329 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000330
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000331 # Test invalid arguments
Guido van Rossum805365e2007-05-07 22:24:25 +0000332 self.assertRaises(TypeError, islice, range(10))
333 self.assertRaises(TypeError, islice, range(10), 1, 2, 3, 4)
334 self.assertRaises(ValueError, islice, range(10), -5, 10, 1)
335 self.assertRaises(ValueError, islice, range(10), 1, -5, -1)
336 self.assertRaises(ValueError, islice, range(10), 1, 10, -1)
337 self.assertRaises(ValueError, islice, range(10), 1, 10, 0)
338 self.assertRaises(ValueError, islice, range(10), 'a')
339 self.assertRaises(ValueError, islice, range(10), 'a', 1)
340 self.assertRaises(ValueError, islice, range(10), 1, 'a')
341 self.assertRaises(ValueError, islice, range(10), 'a', 1, 1)
342 self.assertRaises(ValueError, islice, range(10), 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +0000343 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000344
345 def test_takewhile(self):
346 data = [1, 3, 5, 20, 2, 4, 6, 8]
347 underten = lambda x: x<10
348 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000349 self.assertEqual(list(takewhile(underten, [])), [])
350 self.assertRaises(TypeError, takewhile)
351 self.assertRaises(TypeError, takewhile, operator.pow)
352 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000353 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
354 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000355 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
356 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +0000357 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000358
359 def test_dropwhile(self):
360 data = [1, 3, 5, 20, 2, 4, 6, 8]
361 underten = lambda x: x<10
362 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000363 self.assertEqual(list(dropwhile(underten, [])), [])
364 self.assertRaises(TypeError, dropwhile)
365 self.assertRaises(TypeError, dropwhile, operator.pow)
366 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000367 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
368 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000369
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000370 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000371 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000372 def irange(n):
Guido van Rossum805365e2007-05-07 22:24:25 +0000373 for i in range(n):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000374 yield i
375
376 a, b = tee([]) # test empty iterator
377 self.assertEqual(list(a), [])
378 self.assertEqual(list(b), [])
379
380 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +0000381 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000382
383 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +0000384 self.assertEqual(list(a), list(range(n)))
385 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000386
387 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000388 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000389 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000390 del a
Guido van Rossum805365e2007-05-07 22:24:25 +0000391 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000392
393 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000394 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000395 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000396 del b
Guido van Rossum805365e2007-05-07 22:24:25 +0000397 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000398
Guido van Rossum805365e2007-05-07 22:24:25 +0000399 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000400 order = [0]*n + [1]*n
401 random.shuffle(order)
402 lists = ([], [])
403 its = tee(irange(n))
404 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +0000405 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000406 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +0000407 self.assertEqual(lists[0], list(range(n)))
408 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000409
Raymond Hettingerad983e72003-11-12 14:32:26 +0000410 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000411 self.assertRaises(TypeError, tee)
412 self.assertRaises(TypeError, tee, 3)
413 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000414 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000415
Raymond Hettingerad983e72003-11-12 14:32:26 +0000416 # tee object should be instantiable
417 a, b = tee('abc')
418 c = type(a)('def')
419 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000420
Raymond Hettingerad983e72003-11-12 14:32:26 +0000421 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +0000422 a, b, c = tee(range(2000), 3)
423 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000424 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +0000425 self.assertEqual(list(b), list(range(2000)))
426 self.assertEqual([next(c), next(c)], list(range(2)))
427 self.assertEqual(list(a), list(range(100,2000)))
428 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000429
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000430 # test values of n
431 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000432 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +0000433 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000434 result = tee('abc', n)
435 self.assertEqual(type(result), tuple)
436 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000437 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000438
Raymond Hettingerad983e72003-11-12 14:32:26 +0000439 # tee pass-through to copyable iterator
440 a, b = tee('abc')
441 c, d = tee(a)
442 self.assert_(a is c)
443
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000444 # test tee_new
445 t1, t2 = tee('abc')
446 tnew = type(t1)
447 self.assertRaises(TypeError, tnew)
448 self.assertRaises(TypeError, tnew, 10)
449 t3 = tnew(t1)
450 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000451
Raymond Hettingera9f60922004-10-17 16:40:14 +0000452 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +0000453 a, b = tee(range(10))
Raymond Hettingera9f60922004-10-17 16:40:14 +0000454 p = proxy(a)
455 self.assertEqual(getattr(p, '__class__'), type(b))
456 del a
457 self.assertRaises(ReferenceError, getattr, p, '__class__')
458
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000459 def test_StopIteration(self):
Georg Brandla18af4e2007-04-21 15:47:16 +0000460 self.assertRaises(StopIteration, next, izip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000461
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000462 for f in (chain, cycle, izip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +0000463 self.assertRaises(StopIteration, next, f([]))
464 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000465
Georg Brandla18af4e2007-04-21 15:47:16 +0000466 self.assertRaises(StopIteration, next, islice([], None))
467 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000468
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000469 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +0000470 self.assertRaises(StopIteration, next, p)
471 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000472 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +0000473 self.assertRaises(StopIteration, next, p)
474 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000475
Georg Brandla18af4e2007-04-21 15:47:16 +0000476 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000477
478 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +0000479 self.assertRaises(StopIteration, next, f(lambda x:x, []))
480 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000481
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000482class TestGC(unittest.TestCase):
483
484 def makecycle(self, iterator, container):
485 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +0000486 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000487 del container, iterator
488
489 def test_chain(self):
490 a = []
491 self.makecycle(chain(a), a)
492
493 def test_cycle(self):
494 a = []
495 self.makecycle(cycle([a]*2), a)
496
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000497 def test_dropwhile(self):
498 a = []
499 self.makecycle(dropwhile(bool, [0, a, a]), a)
500
501 def test_groupby(self):
502 a = []
503 self.makecycle(groupby([a]*2, lambda x:x), a)
504
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000505 def test_ifilter(self):
506 a = []
507 self.makecycle(ifilter(lambda x:True, [a]*2), a)
508
509 def test_ifilterfalse(self):
510 a = []
511 self.makecycle(ifilterfalse(lambda x:False, a), a)
512
513 def test_izip(self):
514 a = []
515 self.makecycle(izip([a]*2, [a]*3), a)
516
517 def test_imap(self):
518 a = []
519 self.makecycle(imap(lambda x:x, [a]*2), a)
520
521 def test_islice(self):
522 a = []
523 self.makecycle(islice([a]*2, None), a)
524
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000525 def test_repeat(self):
526 a = []
527 self.makecycle(repeat(a), a)
528
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000529 def test_starmap(self):
530 a = []
531 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
532
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000533 def test_takewhile(self):
534 a = []
535 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
536
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000537def R(seqn):
538 'Regular generator'
539 for i in seqn:
540 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000541
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000542class G:
543 'Sequence using __getitem__'
544 def __init__(self, seqn):
545 self.seqn = seqn
546 def __getitem__(self, i):
547 return self.seqn[i]
548
549class I:
550 'Sequence using iterator protocol'
551 def __init__(self, seqn):
552 self.seqn = seqn
553 self.i = 0
554 def __iter__(self):
555 return self
Georg Brandla18af4e2007-04-21 15:47:16 +0000556 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000557 if self.i >= len(self.seqn): raise StopIteration
558 v = self.seqn[self.i]
559 self.i += 1
560 return v
561
562class Ig:
563 'Sequence using iterator protocol defined with a generator'
564 def __init__(self, seqn):
565 self.seqn = seqn
566 self.i = 0
567 def __iter__(self):
568 for val in self.seqn:
569 yield val
570
571class X:
572 'Missing __getitem__ and __iter__'
573 def __init__(self, seqn):
574 self.seqn = seqn
575 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +0000576 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000577 if self.i >= len(self.seqn): raise StopIteration
578 v = self.seqn[self.i]
579 self.i += 1
580 return v
581
582class N:
Georg Brandla18af4e2007-04-21 15:47:16 +0000583 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000584 def __init__(self, seqn):
585 self.seqn = seqn
586 self.i = 0
587 def __iter__(self):
588 return self
589
590class E:
591 'Test propagation of exceptions'
592 def __init__(self, seqn):
593 self.seqn = seqn
594 self.i = 0
595 def __iter__(self):
596 return self
Georg Brandla18af4e2007-04-21 15:47:16 +0000597 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +0000598 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000599
600class S:
601 'Test immediate stop'
602 def __init__(self, seqn):
603 pass
604 def __iter__(self):
605 return self
Georg Brandla18af4e2007-04-21 15:47:16 +0000606 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000607 raise StopIteration
608
609def L(seqn):
610 'Test multiple tiers of iterators'
611 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
612
613
614class TestVariousIteratorArgs(unittest.TestCase):
615
616 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000617 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000618 for g in (G, I, Ig, S, L, R):
619 self.assertEqual(list(chain(g(s))), list(g(s)))
620 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
621 self.assertRaises(TypeError, chain, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000622 self.assertRaises(TypeError, chain, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000623 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
624
625 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000626 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000627 for g in (G, I, Ig, S, L, R):
628 tgtlen = len(s) * 3
629 expected = list(g(s))*3
630 actual = list(islice(cycle(g(s)), tgtlen))
631 self.assertEqual(actual, expected)
632 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000633 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000634 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
635
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000636 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000637 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000638 for g in (G, I, Ig, S, L, R):
639 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
640 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000641 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000642 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
643
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000644 def test_ifilter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000645 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000646 for g in (G, I, Ig, S, L, R):
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000647 self.assertEqual(list(ifilter(isEven, g(s))),
648 [x for x in g(s) if isEven(x)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000649 self.assertRaises(TypeError, ifilter, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000650 self.assertRaises(TypeError, ifilter, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000651 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
652
653 def test_ifilterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000654 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000655 for g in (G, I, Ig, S, L, R):
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000656 self.assertEqual(list(ifilterfalse(isEven, g(s))),
657 [x for x in g(s) if isOdd(x)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000658 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000659 self.assertRaises(TypeError, ifilterfalse, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000660 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
661
662 def test_izip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000663 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000664 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000665 self.assertEqual(list(izip(g(s))), lzip(g(s)))
666 self.assertEqual(list(izip(g(s), g(s))), lzip(g(s), g(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000667 self.assertRaises(TypeError, izip, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000668 self.assertRaises(TypeError, izip, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000669 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
670
Thomas Wouterscf297e42007-02-23 15:07:44 +0000671 def test_iziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000672 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000673 for g in (G, I, Ig, S, L, R):
674 self.assertEqual(list(izip_longest(g(s))), list(zip(g(s))))
675 self.assertEqual(list(izip_longest(g(s), g(s))), list(zip(g(s), g(s))))
676 self.assertRaises(TypeError, izip_longest, X(s))
677 self.assertRaises(TypeError, izip_longest, N(s))
678 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
679
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000680 def test_imap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000681 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000682 for g in (G, I, Ig, S, L, R):
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000683 self.assertEqual(list(imap(onearg, g(s))),
684 [onearg(x) for x in g(s)])
685 self.assertEqual(list(imap(operator.pow, g(s), g(s))),
686 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000687 self.assertRaises(TypeError, imap, onearg, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000688 self.assertRaises(TypeError, imap, onearg, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000689 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
690
691 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000692 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000693 for g in (G, I, Ig, S, L, R):
694 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
695 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000696 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000697 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
698
699 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000700 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000701 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000702 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000703 self.assertEqual(list(starmap(operator.pow, g(ss))),
704 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000705 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000706 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000707 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
708
709 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000710 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000711 for g in (G, I, Ig, S, L, R):
712 tgt = []
713 for elem in g(s):
714 if not isEven(elem): break
715 tgt.append(elem)
716 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
717 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000718 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000719 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
720
721 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000722 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000723 for g in (G, I, Ig, S, L, R):
724 tgt = []
725 for elem in g(s):
726 if not tgt and isOdd(elem): continue
727 tgt.append(elem)
728 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
729 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000730 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000731 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
732
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000733 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000734 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000735 for g in (G, I, Ig, S, L, R):
736 it1, it2 = tee(g(s))
737 self.assertEqual(list(it1), list(g(s)))
738 self.assertEqual(list(it2), list(g(s)))
739 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000740 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000741 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
742
Raymond Hettinger5cab2e32004-02-10 09:25:40 +0000743class LengthTransparency(unittest.TestCase):
744
745 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +0000746 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +0000747 self.assertEqual(len(repeat(None, 50)), 50)
748 self.assertRaises(TypeError, len, repeat(None))
749
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000750class RegressionTests(unittest.TestCase):
751
752 def test_sf_793826(self):
753 # Fix Armin Rigo's successful efforts to wreak havoc
754
755 def mutatingtuple(tuple1, f, tuple2):
756 # this builds a tuple t which is a copy of tuple1,
757 # then calls f(t), then mutates t to be equal to tuple2
758 # (needs len(tuple1) == len(tuple2)).
759 def g(value, first=[1]):
760 if first:
761 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +0000762 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000763 return value
764 items = list(tuple2)
765 items[1:1] = list(tuple1)
766 gen = imap(g, items)
767 z = izip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +0000768 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000769
770 def f(t):
771 global T
772 T = t
773 first[:] = list(T)
774
775 first = []
776 mutatingtuple((1,2,3), f, (4,5,6))
777 second = list(T)
778 self.assertEqual(first, second)
779
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000780
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000781 def test_sf_950057(self):
782 # Make sure that chain() and cycle() catch exceptions immediately
783 # rather than when shifting between input sources
784
785 def gen1():
786 hist.append(0)
787 yield 1
788 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +0000789 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000790 hist.append(2)
791
792 def gen2(x):
793 hist.append(3)
794 yield 2
795 hist.append(4)
796 if x:
797 raise StopIteration
798
799 hist = []
800 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
801 self.assertEqual(hist, [0,1])
802
803 hist = []
804 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
805 self.assertEqual(hist, [0,1])
806
807 hist = []
808 self.assertRaises(AssertionError, list, cycle(gen1()))
809 self.assertEqual(hist, [0,1])
810
Thomas Woutersb2137042007-02-01 18:02:27 +0000811class SubclassWithKwargsTest(unittest.TestCase):
812 def test_keywords_in_subclass(self):
813 # count is not subclassable...
814 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
815 starmap, islice, takewhile, dropwhile, cycle):
816 class Subclass(cls):
817 def __init__(self, newarg=None, *args):
818 cls.__init__(self, *args)
819 try:
820 Subclass(newarg=1)
821 except TypeError as err:
822 # we expect type errors because of wrong argument count
823 self.failIf("does not take keyword arguments" in err.args[0])
824
825
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000826libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000827
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000828
829>>> amounts = [120.15, 764.05, 823.14]
830>>> for checknum, amount in izip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +0000831... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000832...
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000833Check 1200 is for $120.15
834Check 1201 is for $764.05
835Check 1202 is for $823.14
836
837>>> import operator
Guido van Rossum805365e2007-05-07 22:24:25 +0000838>>> for cube in imap(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +0000839... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000840...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00008411
8428
84327
844
845>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +0000846>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +0000847... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +0000848...
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000849Alex
850Laura
851Martin
852Walter
853Samuele
854
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000855>>> from operator import itemgetter
856>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +0000857>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000858>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000859... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000860...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00008611 ['a', 'c', 'e']
8622 ['b', 'd', 'f']
8633 ['g']
864
Raymond Hettinger734fb572004-01-20 20:04:40 +0000865# Find runs of consecutive numbers using groupby. The key to the solution
866# is differencing with a range so that consecutive numbers all appear in
867# same group.
868>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +0000869>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000870... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000871...
Raymond Hettinger734fb572004-01-20 20:04:40 +0000872[1]
873[4, 5, 6]
874[10]
875[15, 16, 17, 18]
876[22]
877[25, 26, 27, 28]
878
Raymond Hettingera098b332003-09-08 23:58:40 +0000879>>> def take(n, seq):
880... return list(islice(seq, n))
881
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000882>>> def enumerate(iterable):
883... return izip(count(), iterable)
884
885>>> def tabulate(function):
886... "Return function(0), function(1), ..."
887... return imap(function, count())
888
889>>> def iteritems(mapping):
Guido van Rossumcc2b0162007-02-11 06:12:03 +0000890... return izip(mapping.keys(), mapping.values())
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000891
892>>> def nth(iterable, n):
893... "Returns the nth item"
Raymond Hettinger60eca932003-02-09 06:40:58 +0000894... return list(islice(iterable, n, n+1))
895
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000896>>> def all(seq, pred=None):
897... "Returns True if pred(x) is true for every element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000898... for elem in ifilterfalse(pred, seq):
899... return False
900... return True
Raymond Hettinger60eca932003-02-09 06:40:58 +0000901
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000902>>> def any(seq, pred=None):
903... "Returns True if pred(x) is true for at least one element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000904... for elem in ifilter(pred, seq):
905... return True
906... return False
Raymond Hettinger60eca932003-02-09 06:40:58 +0000907
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000908>>> def no(seq, pred=None):
909... "Returns True if pred(x) is false for every element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000910... for elem in ifilter(pred, seq):
911... return False
912... return True
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000913
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000914>>> def quantify(seq, pred=None):
915... "Count how many times the predicate is true in the sequence"
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000916... return sum(imap(pred, seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000917
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000918>>> def padnone(seq):
919... "Returns the sequence elements and then returns None indefinitely"
920... return chain(seq, repeat(None))
921
922>>> def ncycles(seq, n):
923... "Returns the sequence elements n times"
924... return chain(*repeat(seq, n))
925
926>>> def dotproduct(vec1, vec2):
927... return sum(imap(operator.mul, vec1, vec2))
928
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000929>>> def flatten(listOfLists):
930... return list(chain(*listOfLists))
931
932>>> def repeatfunc(func, times=None, *args):
933... "Repeat calls to func with specified arguments."
934... " Example: repeatfunc(random.random)"
935... if times is None:
936... return starmap(func, repeat(args))
937... else:
938... return starmap(func, repeat(args, times))
939
Raymond Hettingerd591f662003-10-26 15:34:50 +0000940>>> def pairwise(iterable):
941... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
942... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000943... try:
Georg Brandla18af4e2007-04-21 15:47:16 +0000944... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000945... except StopIteration:
946... pass
947... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000948
949This is not part of the examples but it tests to make sure the definitions
950perform as purported.
951
Raymond Hettingera098b332003-09-08 23:58:40 +0000952>>> take(10, count())
953[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
954
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000955>>> list(enumerate('abc'))
956[(0, 'a'), (1, 'b'), (2, 'c')]
957
958>>> list(islice(tabulate(lambda x: 2*x), 4))
959[0, 2, 4, 6]
960
961>>> nth('abcde', 3)
962['d']
963
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000964>>> all([2, 4, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000965True
966
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000967>>> all([2, 3, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000968False
969
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000970>>> any([2, 4, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000971True
972
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000973>>> any([1, 3, 5, 9], lambda x: x%2==0,)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000974False
975
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000976>>> no([1, 3, 5, 9], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000977True
978
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000979>>> no([1, 2, 5, 9], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000980False
981
Guido van Rossum805365e2007-05-07 22:24:25 +0000982>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +000098350
984
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000985>>> a = [[1, 2, 3], [4, 5, 6]]
986>>> flatten(a)
987[1, 2, 3, 4, 5, 6]
988
989>>> list(repeatfunc(pow, 5, 2, 3))
990[8, 8, 8, 8, 8]
991
992>>> import random
993>>> take(5, imap(int, repeatfunc(random.random)))
994[0, 0, 0, 0, 0]
995
Raymond Hettingerd591f662003-10-26 15:34:50 +0000996>>> list(pairwise('abcd'))
997[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000998
Raymond Hettingerd591f662003-10-26 15:34:50 +0000999>>> list(pairwise([]))
1000[]
1001
1002>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001003[]
1004
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001005>>> list(islice(padnone('abc'), 0, 6))
1006['a', 'b', 'c', None, None, None]
1007
1008>>> list(ncycles('abc', 3))
1009['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1010
1011>>> dotproduct([1,2,3], [4,5,6])
101232
1013
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001014"""
1015
1016__test__ = {'libreftest' : libreftest}
1017
1018def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001019 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00001020 RegressionTests, LengthTransparency,
1021 SubclassWithKwargsTest)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001022 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001023
1024 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001025 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001026 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001027 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00001028 for i in range(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001029 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001030 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001031 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001032 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001033
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001034 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001035 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001036
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001037if __name__ == "__main__":
1038 test_main(verbose=True)