blob: ae5eb901a0428332f09d5a8eb4dc0d042d6106b5 [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)])
Guido van Rossum8ce8a782007-11-01 19:42:39 +000058 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
59 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger96ef8112003-02-01 00:10:11 +000060 self.assertRaises(TypeError, count, 2, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000061 self.assertRaises(TypeError, count, 'a')
Guido van Rossum8ce8a782007-11-01 19:42:39 +000062 self.assertEqual(list(islice(count(maxsize-5), 10)),
63 list(range(maxsize-5, maxsize+5)))
64 self.assertEqual(list(islice(count(-maxsize-5), 10)),
65 list(range(-maxsize-5, -maxsize+5)))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +000066 c = count(3)
67 self.assertEqual(repr(c), 'count(3)')
Georg Brandla18af4e2007-04-21 15:47:16 +000068 next(c)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +000069 self.assertEqual(repr(c), 'count(4)')
Thomas Wouters89f507f2006-12-13 04:49:30 +000070 c = count(-9)
71 self.assertEqual(repr(c), 'count(-9)')
Georg Brandla18af4e2007-04-21 15:47:16 +000072 next(c)
73 self.assertEqual(next(c), -8)
Christian Heimesa37d4c62007-12-04 23:02:19 +000074 for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5):
Guido van Rossum8ce8a782007-11-01 19:42:39 +000075 # Test repr (ignoring the L in longs)
76 r1 = repr(count(i)).replace('L', '')
77 r2 = 'count(%r)'.__mod__(i).replace('L', '')
78 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +000079
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000080 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +000081 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000082 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000083 self.assertRaises(TypeError, cycle)
84 self.assertRaises(TypeError, cycle, 5)
85 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000086
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000087 def test_groupby(self):
88 # Check whether it accepts arguments correctly
89 self.assertEqual([], list(groupby([])))
90 self.assertEqual([], list(groupby([], key=id)))
91 self.assertRaises(TypeError, list, groupby('abc', []))
92 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +000093 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000094
95 # Check normal input
96 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
97 (2,15,22), (3,16,23), (3,17,23)]
98 dup = []
99 for k, g in groupby(s, lambda r:r[0]):
100 for elem in g:
101 self.assertEqual(k, elem[0])
102 dup.append(elem)
103 self.assertEqual(s, dup)
104
105 # Check nested case
106 dup = []
107 for k, g in groupby(s, lambda r:r[0]):
108 for ik, ig in groupby(g, lambda r:r[2]):
109 for elem in ig:
110 self.assertEqual(k, elem[0])
111 self.assertEqual(ik, elem[2])
112 dup.append(elem)
113 self.assertEqual(s, dup)
114
115 # Check case where inner iterator is not used
116 keys = [k for k, g in groupby(s, lambda r:r[0])]
117 expectedkeys = set([r[0] for r in s])
118 self.assertEqual(set(keys), expectedkeys)
119 self.assertEqual(len(keys), len(expectedkeys))
120
121 # Exercise pipes and filters style
122 s = 'abracadabra'
123 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000124 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000125 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
126 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000127 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000128 self.assertEqual(r, ['a', 'b', 'r'])
129 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000130 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000131 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
132 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000133 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000134 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
135
Georg Brandla18af4e2007-04-21 15:47:16 +0000136 # iter.__next__ failure
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000137 class ExpectedError(Exception):
138 pass
139 def delayed_raise(n=0):
140 for i in range(n):
141 yield 'yo'
142 raise ExpectedError
143 def gulp(iterable, keyp=None, func=list):
144 return [func(g) for k, g in groupby(iterable, keyp)]
145
Georg Brandla18af4e2007-04-21 15:47:16 +0000146 # iter.__next__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000147 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
Georg Brandla18af4e2007-04-21 15:47:16 +0000148 # iter.__next__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000149 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
150
151 # __cmp__ failure
152 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000153 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000154 raise ExpectedError
155 s = [DummyCmp(), DummyCmp(), None]
156
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000157 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000158 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000159 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000160 self.assertRaises(ExpectedError, gulp, s)
161
162 # keyfunc failure
163 def keyfunc(obj):
164 if keyfunc.skip > 0:
165 keyfunc.skip -= 1
166 return obj
167 else:
168 raise ExpectedError
169
170 # keyfunc failure on outer object
171 keyfunc.skip = 0
172 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
173 keyfunc.skip = 1
174 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
175
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000176 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000177 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000178 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000179 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000180 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000181 self.assertRaises(TypeError, ifilter, lambda x:x)
182 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000183 self.assertRaises(TypeError, ifilter, isEven, 3)
Georg Brandla18af4e2007-04-21 15:47:16 +0000184 self.assertRaises(TypeError, next, ifilter(range(6), range(6)))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000185
186 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000187 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
188 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000189 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000190 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000191 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
192 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000193 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Georg Brandla18af4e2007-04-21 15:47:16 +0000194 self.assertRaises(TypeError, next, ifilterfalse(range(6), range(6)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000195
196 def test_izip(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000197 # XXX This is rather silly now that builtin zip() calls izip()...
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000198 ans = [(x,y) for x, y in izip('abc',count())]
199 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Guido van Rossum801f0d72006-08-24 19:48:10 +0000200 self.assertEqual(list(izip('abc', range(6))), lzip('abc', range(6)))
201 self.assertEqual(list(izip('abcdef', range(3))), lzip('abcdef', range(3)))
202 self.assertEqual(take(3,izip('abcdef', count())), lzip('abcdef', range(3)))
203 self.assertEqual(list(izip('abcdef')), lzip('abcdef'))
204 self.assertEqual(list(izip()), lzip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000205 self.assertRaises(TypeError, izip, 3)
206 self.assertRaises(TypeError, izip, range(3), 3)
207 # Check tuple re-use (implementation detail)
208 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000209 lzip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000210 self.assertEqual([pair for pair in izip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000211 lzip('abc', 'def'))
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000212 ids = list(map(id, izip('abc', 'def')))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000213 self.assertEqual(min(ids), max(ids))
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000214 ids = list(map(id, list(izip('abc', 'def'))))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000215 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000216
Thomas Wouterscf297e42007-02-23 15:07:44 +0000217 def test_iziplongest(self):
218 for args in [
219 ['abc', range(6)],
220 [range(6), 'abc'],
221 [range(1000), range(2000,2100), range(3000,3050)],
222 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
223 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
224 ]:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000225 target = [tuple([arg[i] if i < len(arg) else None for arg in args])
226 for i in range(max(map(len, args)))]
Thomas Wouterscf297e42007-02-23 15:07:44 +0000227 self.assertEqual(list(izip_longest(*args)), target)
228 self.assertEqual(list(izip_longest(*args, **{})), target)
229 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
230 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000231
Thomas Wouterscf297e42007-02-23 15:07:44 +0000232 self.assertEqual(take(3,izip_longest('abcdef', count())), list(zip('abcdef', range(3)))) # take 3 from infinite input
233
234 self.assertEqual(list(izip_longest()), list(zip()))
235 self.assertEqual(list(izip_longest([])), list(zip([])))
236 self.assertEqual(list(izip_longest('abcdef')), list(zip('abcdef')))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000237
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000238 self.assertEqual(list(izip_longest('abc', 'defg', **{})),
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000239 list(izip(list('abc')+[None], 'defg'))) # empty keyword dict
Thomas Wouterscf297e42007-02-23 15:07:44 +0000240 self.assertRaises(TypeError, izip_longest, 3)
241 self.assertRaises(TypeError, izip_longest, range(3), 3)
242
243 for stmt in [
244 "izip_longest('abc', fv=1)",
Guido van Rossumd8faa362007-04-27 19:54:29 +0000245 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Thomas Wouterscf297e42007-02-23 15:07:44 +0000246 ]:
247 try:
248 eval(stmt, globals(), locals())
249 except TypeError:
250 pass
251 else:
252 self.fail('Did not raise Type in: ' + stmt)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000253
Thomas Wouterscf297e42007-02-23 15:07:44 +0000254 # Check tuple re-use (implementation detail)
255 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
256 list(zip('abc', 'def')))
257 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
258 list(zip('abc', 'def')))
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000259 ids = list(map(id, izip_longest('abc', 'def')))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000260 self.assertEqual(min(ids), max(ids))
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000261 ids = list(map(id, list(izip_longest('abc', 'def'))))
Thomas Wouterscf297e42007-02-23 15:07:44 +0000262 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
263
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000264 def test_repeat(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000265 self.assertEqual(lzip(range(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000266 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000267 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000268 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000269 self.assertEqual(list(repeat('a', 0)), [])
270 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000271 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000272 self.assertRaises(TypeError, repeat, None, 3, 4)
273 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000274 r = repeat(1+0j)
275 self.assertEqual(repr(r), 'repeat((1+0j))')
276 r = repeat(1+0j, 5)
277 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
278 list(r)
279 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000280
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000281 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000282 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
283 [0**1, 1**2, 2**3])
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000284 def tupleize(*args):
285 return args
286 self.assertEqual(list(imap(tupleize, 'abc', range(5))),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000287 [('a',0),('b',1),('c',2)])
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000288 self.assertEqual(list(imap(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000289 [('a',0),('b',1),('c',2)])
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000290 self.assertEqual(take(2,imap(tupleize, 'abc', count())),
Raymond Hettinger02420702003-06-29 20:36:23 +0000291 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000292 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000293 self.assertRaises(TypeError, imap)
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000294 self.assertRaises(TypeError, list, imap(None, range(3), range(3)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000295 self.assertRaises(TypeError, imap, operator.neg)
Georg Brandla18af4e2007-04-21 15:47:16 +0000296 self.assertRaises(TypeError, next, imap(10, range(5)))
297 self.assertRaises(ValueError, next, imap(errfunc, [4], [5]))
298 self.assertRaises(TypeError, next, imap(onearg, [4], [5]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000299
300 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000301 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
302 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000303 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
304 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000305 self.assertEqual(list(starmap(operator.pow, [])), [])
Christian Heimes679db4a2008-01-18 09:56:22 +0000306 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
307 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000308 self.assertRaises(TypeError, starmap)
309 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000310 self.assertRaises(TypeError, next, starmap(10, [(4,5)]))
311 self.assertRaises(ValueError, next, starmap(errfunc, [(4,5)]))
312 self.assertRaises(TypeError, next, starmap(onearg, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000313
314 def test_islice(self):
315 for args in [ # islice(args) should agree with range(args)
316 (10, 20, 3),
317 (10, 3, 20),
318 (10, 20),
319 (10, 3),
320 (20,)
321 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000322 self.assertEqual(list(islice(range(100), *args)),
323 list(range(*args)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000324
325 for args, tgtargs in [ # Stop when seqn is exhausted
326 ((10, 110, 3), ((10, 100, 3))),
327 ((10, 110), ((10, 100))),
328 ((110,), (100,))
329 ]:
Guido van Rossum805365e2007-05-07 22:24:25 +0000330 self.assertEqual(list(islice(range(100), *args)),
331 list(range(*tgtargs)))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000332
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000333 # Test stop=None
Guido van Rossum805365e2007-05-07 22:24:25 +0000334 self.assertEqual(list(islice(range(10), None)), list(range(10)))
335 self.assertEqual(list(islice(range(10), None, None)), list(range(10)))
336 self.assertEqual(list(islice(range(10), None, None, None)), list(range(10)))
337 self.assertEqual(list(islice(range(10), 2, None)), list(range(2, 10)))
338 self.assertEqual(list(islice(range(10), 1, None, 2)), list(range(1, 10, 2)))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000339
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000340 # Test number of items consumed SF #1171417
341 it = iter(range(10))
Guido van Rossum805365e2007-05-07 22:24:25 +0000342 self.assertEqual(list(islice(it, 3)), list(range(3)))
343 self.assertEqual(list(it), list(range(3, 10)))
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000344
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000345 # Test invalid arguments
Guido van Rossum805365e2007-05-07 22:24:25 +0000346 self.assertRaises(TypeError, islice, range(10))
347 self.assertRaises(TypeError, islice, range(10), 1, 2, 3, 4)
348 self.assertRaises(ValueError, islice, range(10), -5, 10, 1)
349 self.assertRaises(ValueError, islice, range(10), 1, -5, -1)
350 self.assertRaises(ValueError, islice, range(10), 1, 10, -1)
351 self.assertRaises(ValueError, islice, range(10), 1, 10, 0)
352 self.assertRaises(ValueError, islice, range(10), 'a')
353 self.assertRaises(ValueError, islice, range(10), 'a', 1)
354 self.assertRaises(ValueError, islice, range(10), 1, 'a')
355 self.assertRaises(ValueError, islice, range(10), 'a', 1, 1)
356 self.assertRaises(ValueError, islice, range(10), 1, 'a', 1)
Guido van Rossum360e4b82007-05-14 22:51:27 +0000357 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000358
359 def test_takewhile(self):
360 data = [1, 3, 5, 20, 2, 4, 6, 8]
361 underten = lambda x: x<10
362 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000363 self.assertEqual(list(takewhile(underten, [])), [])
364 self.assertRaises(TypeError, takewhile)
365 self.assertRaises(TypeError, takewhile, operator.pow)
366 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000367 self.assertRaises(TypeError, next, takewhile(10, [(4,5)]))
368 self.assertRaises(ValueError, next, takewhile(errfunc, [(4,5)]))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000369 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
370 self.assertEqual(list(t), [1, 1, 1])
Georg Brandla18af4e2007-04-21 15:47:16 +0000371 self.assertRaises(StopIteration, next, t)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000372
373 def test_dropwhile(self):
374 data = [1, 3, 5, 20, 2, 4, 6, 8]
375 underten = lambda x: x<10
376 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000377 self.assertEqual(list(dropwhile(underten, [])), [])
378 self.assertRaises(TypeError, dropwhile)
379 self.assertRaises(TypeError, dropwhile, operator.pow)
380 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
Georg Brandla18af4e2007-04-21 15:47:16 +0000381 self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
382 self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000383
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000384 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000385 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000386 def irange(n):
Guido van Rossum805365e2007-05-07 22:24:25 +0000387 for i in range(n):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000388 yield i
389
390 a, b = tee([]) # test empty iterator
391 self.assertEqual(list(a), [])
392 self.assertEqual(list(b), [])
393
394 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +0000395 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000396
397 a, b = tee(irange(n)) # test 0% interleaved
Guido van Rossum805365e2007-05-07 22:24:25 +0000398 self.assertEqual(list(a), list(range(n)))
399 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000400
401 a, b = tee(irange(n)) # test dealloc of leading iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000402 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000403 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000404 del a
Guido van Rossum805365e2007-05-07 22:24:25 +0000405 self.assertEqual(list(b), list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000406
407 a, b = tee(irange(n)) # test dealloc of trailing iterator
Guido van Rossum805365e2007-05-07 22:24:25 +0000408 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000409 self.assertEqual(next(a), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000410 del b
Guido van Rossum805365e2007-05-07 22:24:25 +0000411 self.assertEqual(list(a), list(range(100, n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000412
Guido van Rossum805365e2007-05-07 22:24:25 +0000413 for j in range(5): # test randomly interleaved
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000414 order = [0]*n + [1]*n
415 random.shuffle(order)
416 lists = ([], [])
417 its = tee(irange(n))
418 for i in order:
Georg Brandla18af4e2007-04-21 15:47:16 +0000419 value = next(its[i])
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000420 lists[i].append(value)
Guido van Rossum805365e2007-05-07 22:24:25 +0000421 self.assertEqual(lists[0], list(range(n)))
422 self.assertEqual(lists[1], list(range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000423
Raymond Hettingerad983e72003-11-12 14:32:26 +0000424 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000425 self.assertRaises(TypeError, tee)
426 self.assertRaises(TypeError, tee, 3)
427 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000428 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000429
Raymond Hettingerad983e72003-11-12 14:32:26 +0000430 # tee object should be instantiable
431 a, b = tee('abc')
432 c = type(a)('def')
433 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000434
Raymond Hettingerad983e72003-11-12 14:32:26 +0000435 # test long-lagged and multi-way split
Guido van Rossum805365e2007-05-07 22:24:25 +0000436 a, b, c = tee(range(2000), 3)
437 for i in range(100):
Georg Brandla18af4e2007-04-21 15:47:16 +0000438 self.assertEqual(next(a), i)
Guido van Rossum805365e2007-05-07 22:24:25 +0000439 self.assertEqual(list(b), list(range(2000)))
440 self.assertEqual([next(c), next(c)], list(range(2)))
441 self.assertEqual(list(a), list(range(100,2000)))
442 self.assertEqual(list(c), list(range(2,2000)))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000443
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000444 # test values of n
445 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Thomas Wouters89f507f2006-12-13 04:49:30 +0000446 self.assertRaises(ValueError, tee, [], -1)
Guido van Rossum805365e2007-05-07 22:24:25 +0000447 for n in range(5):
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000448 result = tee('abc', n)
449 self.assertEqual(type(result), tuple)
450 self.assertEqual(len(result), n)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000451 self.assertEqual([list(x) for x in result], [list('abc')]*n)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000452
Raymond Hettingerad983e72003-11-12 14:32:26 +0000453 # tee pass-through to copyable iterator
454 a, b = tee('abc')
455 c, d = tee(a)
456 self.assert_(a is c)
457
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000458 # test tee_new
459 t1, t2 = tee('abc')
460 tnew = type(t1)
461 self.assertRaises(TypeError, tnew)
462 self.assertRaises(TypeError, tnew, 10)
463 t3 = tnew(t1)
464 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000465
Raymond Hettingera9f60922004-10-17 16:40:14 +0000466 # test that tee objects are weak referencable
Guido van Rossum805365e2007-05-07 22:24:25 +0000467 a, b = tee(range(10))
Raymond Hettingera9f60922004-10-17 16:40:14 +0000468 p = proxy(a)
469 self.assertEqual(getattr(p, '__class__'), type(b))
470 del a
471 self.assertRaises(ReferenceError, getattr, p, '__class__')
472
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000473 def test_StopIteration(self):
Georg Brandla18af4e2007-04-21 15:47:16 +0000474 self.assertRaises(StopIteration, next, izip())
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000475
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000476 for f in (chain, cycle, izip, groupby):
Georg Brandla18af4e2007-04-21 15:47:16 +0000477 self.assertRaises(StopIteration, next, f([]))
478 self.assertRaises(StopIteration, next, f(StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000479
Georg Brandla18af4e2007-04-21 15:47:16 +0000480 self.assertRaises(StopIteration, next, islice([], None))
481 self.assertRaises(StopIteration, next, islice(StopNow(), None))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000482
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000483 p, q = tee([])
Georg Brandla18af4e2007-04-21 15:47:16 +0000484 self.assertRaises(StopIteration, next, p)
485 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000486 p, q = tee(StopNow())
Georg Brandla18af4e2007-04-21 15:47:16 +0000487 self.assertRaises(StopIteration, next, p)
488 self.assertRaises(StopIteration, next, q)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000489
Georg Brandla18af4e2007-04-21 15:47:16 +0000490 self.assertRaises(StopIteration, next, repeat(None, 0))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000491
492 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
Georg Brandla18af4e2007-04-21 15:47:16 +0000493 self.assertRaises(StopIteration, next, f(lambda x:x, []))
494 self.assertRaises(StopIteration, next, f(lambda x:x, StopNow()))
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000495
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000496class TestGC(unittest.TestCase):
497
498 def makecycle(self, iterator, container):
499 container.append(iterator)
Georg Brandla18af4e2007-04-21 15:47:16 +0000500 next(iterator)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000501 del container, iterator
502
503 def test_chain(self):
504 a = []
505 self.makecycle(chain(a), a)
506
507 def test_cycle(self):
508 a = []
509 self.makecycle(cycle([a]*2), a)
510
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000511 def test_dropwhile(self):
512 a = []
513 self.makecycle(dropwhile(bool, [0, a, a]), a)
514
515 def test_groupby(self):
516 a = []
517 self.makecycle(groupby([a]*2, lambda x:x), a)
518
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000519 def test_ifilter(self):
520 a = []
521 self.makecycle(ifilter(lambda x:True, [a]*2), a)
522
523 def test_ifilterfalse(self):
524 a = []
525 self.makecycle(ifilterfalse(lambda x:False, a), a)
526
527 def test_izip(self):
528 a = []
529 self.makecycle(izip([a]*2, [a]*3), a)
530
531 def test_imap(self):
532 a = []
533 self.makecycle(imap(lambda x:x, [a]*2), a)
534
535 def test_islice(self):
536 a = []
537 self.makecycle(islice([a]*2, None), a)
538
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000539 def test_repeat(self):
540 a = []
541 self.makecycle(repeat(a), a)
542
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000543 def test_starmap(self):
544 a = []
545 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
546
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000547 def test_takewhile(self):
548 a = []
549 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
550
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000551def R(seqn):
552 'Regular generator'
553 for i in seqn:
554 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000555
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000556class G:
557 'Sequence using __getitem__'
558 def __init__(self, seqn):
559 self.seqn = seqn
560 def __getitem__(self, i):
561 return self.seqn[i]
562
563class I:
564 'Sequence using iterator protocol'
565 def __init__(self, seqn):
566 self.seqn = seqn
567 self.i = 0
568 def __iter__(self):
569 return self
Georg Brandla18af4e2007-04-21 15:47:16 +0000570 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000571 if self.i >= len(self.seqn): raise StopIteration
572 v = self.seqn[self.i]
573 self.i += 1
574 return v
575
576class Ig:
577 'Sequence using iterator protocol defined with a generator'
578 def __init__(self, seqn):
579 self.seqn = seqn
580 self.i = 0
581 def __iter__(self):
582 for val in self.seqn:
583 yield val
584
585class X:
586 'Missing __getitem__ and __iter__'
587 def __init__(self, seqn):
588 self.seqn = seqn
589 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +0000590 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000591 if self.i >= len(self.seqn): raise StopIteration
592 v = self.seqn[self.i]
593 self.i += 1
594 return v
595
596class N:
Georg Brandla18af4e2007-04-21 15:47:16 +0000597 'Iterator missing __next__()'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000598 def __init__(self, seqn):
599 self.seqn = seqn
600 self.i = 0
601 def __iter__(self):
602 return self
603
604class E:
605 'Test propagation of exceptions'
606 def __init__(self, seqn):
607 self.seqn = seqn
608 self.i = 0
609 def __iter__(self):
610 return self
Georg Brandla18af4e2007-04-21 15:47:16 +0000611 def __next__(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +0000612 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000613
614class S:
615 'Test immediate stop'
616 def __init__(self, seqn):
617 pass
618 def __iter__(self):
619 return self
Georg Brandla18af4e2007-04-21 15:47:16 +0000620 def __next__(self):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000621 raise StopIteration
622
623def L(seqn):
624 'Test multiple tiers of iterators'
625 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
626
627
628class TestVariousIteratorArgs(unittest.TestCase):
629
630 def test_chain(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000631 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000632 for g in (G, I, Ig, S, L, R):
633 self.assertEqual(list(chain(g(s))), list(g(s)))
634 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
635 self.assertRaises(TypeError, chain, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000636 self.assertRaises(TypeError, chain, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000637 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
638
639 def test_cycle(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000640 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000641 for g in (G, I, Ig, S, L, R):
642 tgtlen = len(s) * 3
643 expected = list(g(s))*3
644 actual = list(islice(cycle(g(s)), tgtlen))
645 self.assertEqual(actual, expected)
646 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000647 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000648 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
649
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000650 def test_groupby(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000651 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000652 for g in (G, I, Ig, S, L, R):
653 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
654 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000655 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000656 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
657
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000658 def test_ifilter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000659 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000660 for g in (G, I, Ig, S, L, R):
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000661 self.assertEqual(list(ifilter(isEven, g(s))),
662 [x for x in g(s) if isEven(x)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000663 self.assertRaises(TypeError, ifilter, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000664 self.assertRaises(TypeError, ifilter, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000665 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
666
667 def test_ifilterfalse(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000668 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000669 for g in (G, I, Ig, S, L, R):
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000670 self.assertEqual(list(ifilterfalse(isEven, g(s))),
671 [x for x in g(s) if isOdd(x)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000672 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000673 self.assertRaises(TypeError, ifilterfalse, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000674 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
675
676 def test_izip(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000677 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000678 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000679 self.assertEqual(list(izip(g(s))), lzip(g(s)))
680 self.assertEqual(list(izip(g(s), g(s))), lzip(g(s), g(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000681 self.assertRaises(TypeError, izip, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000682 self.assertRaises(TypeError, izip, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000683 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
684
Thomas Wouterscf297e42007-02-23 15:07:44 +0000685 def test_iziplongest(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000686 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000687 for g in (G, I, Ig, S, L, R):
688 self.assertEqual(list(izip_longest(g(s))), list(zip(g(s))))
689 self.assertEqual(list(izip_longest(g(s), g(s))), list(zip(g(s), g(s))))
690 self.assertRaises(TypeError, izip_longest, X(s))
691 self.assertRaises(TypeError, izip_longest, N(s))
692 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
693
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000694 def test_imap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000695 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000696 for g in (G, I, Ig, S, L, R):
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000697 self.assertEqual(list(imap(onearg, g(s))),
698 [onearg(x) for x in g(s)])
699 self.assertEqual(list(imap(operator.pow, g(s), g(s))),
700 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000701 self.assertRaises(TypeError, imap, onearg, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000702 self.assertRaises(TypeError, imap, onearg, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000703 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
704
705 def test_islice(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000706 for s in ("12345", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000707 for g in (G, I, Ig, S, L, R):
708 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
709 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000710 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000711 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
712
713 def test_starmap(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000714 for s in (range(10), range(0), range(100), (7,11), range(20,50,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000715 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000716 ss = lzip(s, s)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000717 self.assertEqual(list(starmap(operator.pow, g(ss))),
718 [x**x for x in g(s)])
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000719 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000720 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000721 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
722
723 def test_takewhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000724 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000725 for g in (G, I, Ig, S, L, R):
726 tgt = []
727 for elem in g(s):
728 if not isEven(elem): break
729 tgt.append(elem)
730 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
731 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000732 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000733 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
734
735 def test_dropwhile(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000736 for s in (range(10), range(0), range(1000), (7,11), range(2000,2200,5)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000737 for g in (G, I, Ig, S, L, R):
738 tgt = []
739 for elem in g(s):
740 if not tgt and isOdd(elem): continue
741 tgt.append(elem)
742 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
743 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000744 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000745 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
746
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000747 def test_tee(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000748 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000749 for g in (G, I, Ig, S, L, R):
750 it1, it2 = tee(g(s))
751 self.assertEqual(list(it1), list(g(s)))
752 self.assertEqual(list(it2), list(g(s)))
753 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000754 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000755 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
756
Raymond Hettinger5cab2e32004-02-10 09:25:40 +0000757class LengthTransparency(unittest.TestCase):
758
759 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +0000760 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +0000761 self.assertEqual(len(repeat(None, 50)), 50)
762 self.assertRaises(TypeError, len, repeat(None))
763
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000764class RegressionTests(unittest.TestCase):
765
766 def test_sf_793826(self):
767 # Fix Armin Rigo's successful efforts to wreak havoc
768
769 def mutatingtuple(tuple1, f, tuple2):
770 # this builds a tuple t which is a copy of tuple1,
771 # then calls f(t), then mutates t to be equal to tuple2
772 # (needs len(tuple1) == len(tuple2)).
773 def g(value, first=[1]):
774 if first:
775 del first[:]
Georg Brandla18af4e2007-04-21 15:47:16 +0000776 f(next(z))
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000777 return value
778 items = list(tuple2)
779 items[1:1] = list(tuple1)
780 gen = imap(g, items)
781 z = izip(*[gen]*len(tuple1))
Georg Brandla18af4e2007-04-21 15:47:16 +0000782 next(z)
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000783
784 def f(t):
785 global T
786 T = t
787 first[:] = list(T)
788
789 first = []
790 mutatingtuple((1,2,3), f, (4,5,6))
791 second = list(T)
792 self.assertEqual(first, second)
793
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000794
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000795 def test_sf_950057(self):
796 # Make sure that chain() and cycle() catch exceptions immediately
797 # rather than when shifting between input sources
798
799 def gen1():
800 hist.append(0)
801 yield 1
802 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +0000803 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000804 hist.append(2)
805
806 def gen2(x):
807 hist.append(3)
808 yield 2
809 hist.append(4)
810 if x:
811 raise StopIteration
812
813 hist = []
814 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
815 self.assertEqual(hist, [0,1])
816
817 hist = []
818 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
819 self.assertEqual(hist, [0,1])
820
821 hist = []
822 self.assertRaises(AssertionError, list, cycle(gen1()))
823 self.assertEqual(hist, [0,1])
824
Thomas Woutersb2137042007-02-01 18:02:27 +0000825class SubclassWithKwargsTest(unittest.TestCase):
826 def test_keywords_in_subclass(self):
827 # count is not subclassable...
828 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
829 starmap, islice, takewhile, dropwhile, cycle):
830 class Subclass(cls):
831 def __init__(self, newarg=None, *args):
832 cls.__init__(self, *args)
833 try:
834 Subclass(newarg=1)
835 except TypeError as err:
836 # we expect type errors because of wrong argument count
837 self.failIf("does not take keyword arguments" in err.args[0])
838
839
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000840libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000841
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000842
843>>> amounts = [120.15, 764.05, 823.14]
844>>> for checknum, amount in izip(count(1200), amounts):
Guido van Rossum7131f842007-02-09 20:13:25 +0000845... print('Check %d is for $%.2f' % (checknum, amount))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000846...
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000847Check 1200 is for $120.15
848Check 1201 is for $764.05
849Check 1202 is for $823.14
850
851>>> import operator
Guido van Rossum805365e2007-05-07 22:24:25 +0000852>>> for cube in imap(operator.pow, range(1,4), repeat(3)):
Guido van Rossum7131f842007-02-09 20:13:25 +0000853... print(cube)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000854...
Raymond Hettinger96ef8112003-02-01 00:10:11 +00008551
8568
85727
858
859>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +0000860>>> for name in islice(reportlines, 3, None, 2):
Guido van Rossum7131f842007-02-09 20:13:25 +0000861... print(name.title())
Guido van Rossumd8faa362007-04-27 19:54:29 +0000862...
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000863Alex
864Laura
865Martin
866Walter
867Samuele
868
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000869>>> from operator import itemgetter
870>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Guido van Rossumcc2b0162007-02-11 06:12:03 +0000871>>> di = sorted(sorted(d.items()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000872>>> for k, g in groupby(di, itemgetter(1)):
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000873... print(k, list(map(itemgetter(0), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000874...
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00008751 ['a', 'c', 'e']
8762 ['b', 'd', 'f']
8773 ['g']
878
Raymond Hettinger734fb572004-01-20 20:04:40 +0000879# Find runs of consecutive numbers using groupby. The key to the solution
880# is differencing with a range so that consecutive numbers all appear in
881# same group.
882>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Guido van Rossum1bc535d2007-05-15 18:46:22 +0000883>>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000884... print(list(map(operator.itemgetter(1), g)))
Guido van Rossumd8faa362007-04-27 19:54:29 +0000885...
Raymond Hettinger734fb572004-01-20 20:04:40 +0000886[1]
887[4, 5, 6]
888[10]
889[15, 16, 17, 18]
890[22]
891[25, 26, 27, 28]
892
Raymond Hettingera098b332003-09-08 23:58:40 +0000893>>> def take(n, seq):
894... return list(islice(seq, n))
895
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000896>>> def enumerate(iterable):
897... return izip(count(), iterable)
898
899>>> def tabulate(function):
900... "Return function(0), function(1), ..."
901... return imap(function, count())
902
903>>> def iteritems(mapping):
Guido van Rossumcc2b0162007-02-11 06:12:03 +0000904... return izip(mapping.keys(), mapping.values())
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000905
906>>> def nth(iterable, n):
907... "Returns the nth item"
Raymond Hettinger60eca932003-02-09 06:40:58 +0000908... return list(islice(iterable, n, n+1))
909
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000910>>> def all(seq, pred=None):
911... "Returns True if pred(x) is true for every element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000912... for elem in ifilterfalse(pred, seq):
913... return False
914... return True
Raymond Hettinger60eca932003-02-09 06:40:58 +0000915
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000916>>> def any(seq, pred=None):
917... "Returns True if pred(x) is true for at least one element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000918... for elem in ifilter(pred, seq):
919... return True
920... return False
Raymond Hettinger60eca932003-02-09 06:40:58 +0000921
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000922>>> def no(seq, pred=None):
923... "Returns True if pred(x) is false for every element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000924... for elem in ifilter(pred, seq):
925... return False
926... return True
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000927
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000928>>> def quantify(seq, pred=None):
929... "Count how many times the predicate is true in the sequence"
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000930... return sum(imap(pred, seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000931
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000932>>> def padnone(seq):
933... "Returns the sequence elements and then returns None indefinitely"
934... return chain(seq, repeat(None))
935
936>>> def ncycles(seq, n):
937... "Returns the sequence elements n times"
938... return chain(*repeat(seq, n))
939
940>>> def dotproduct(vec1, vec2):
941... return sum(imap(operator.mul, vec1, vec2))
942
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000943>>> def flatten(listOfLists):
944... return list(chain(*listOfLists))
945
946>>> def repeatfunc(func, times=None, *args):
947... "Repeat calls to func with specified arguments."
948... " Example: repeatfunc(random.random)"
949... if times is None:
950... return starmap(func, repeat(args))
951... else:
952... return starmap(func, repeat(args, times))
953
Raymond Hettingerd591f662003-10-26 15:34:50 +0000954>>> def pairwise(iterable):
955... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
956... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000957... try:
Georg Brandla18af4e2007-04-21 15:47:16 +0000958... next(b)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000959... except StopIteration:
960... pass
961... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000962
963This is not part of the examples but it tests to make sure the definitions
964perform as purported.
965
Raymond Hettingera098b332003-09-08 23:58:40 +0000966>>> take(10, count())
967[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
968
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000969>>> list(enumerate('abc'))
970[(0, 'a'), (1, 'b'), (2, 'c')]
971
972>>> list(islice(tabulate(lambda x: 2*x), 4))
973[0, 2, 4, 6]
974
975>>> nth('abcde', 3)
976['d']
977
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000978>>> all([2, 4, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000979True
980
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000981>>> all([2, 3, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000982False
983
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000984>>> any([2, 4, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000985True
986
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000987>>> any([1, 3, 5, 9], lambda x: x%2==0,)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000988False
989
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000990>>> no([1, 3, 5, 9], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000991True
992
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000993>>> no([1, 2, 5, 9], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000994False
995
Guido van Rossum805365e2007-05-07 22:24:25 +0000996>>> quantify(range(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +000099750
998
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000999>>> a = [[1, 2, 3], [4, 5, 6]]
1000>>> flatten(a)
1001[1, 2, 3, 4, 5, 6]
1002
1003>>> list(repeatfunc(pow, 5, 2, 3))
1004[8, 8, 8, 8, 8]
1005
1006>>> import random
1007>>> take(5, imap(int, repeatfunc(random.random)))
1008[0, 0, 0, 0, 0]
1009
Raymond Hettingerd591f662003-10-26 15:34:50 +00001010>>> list(pairwise('abcd'))
1011[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001012
Raymond Hettingerd591f662003-10-26 15:34:50 +00001013>>> list(pairwise([]))
1014[]
1015
1016>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001017[]
1018
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001019>>> list(islice(padnone('abc'), 0, 6))
1020['a', 'b', 'c', None, None, None]
1021
1022>>> list(ncycles('abc', 3))
1023['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1024
1025>>> dotproduct([1,2,3], [4,5,6])
102632
1027
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001028"""
1029
1030__test__ = {'libreftest' : libreftest}
1031
1032def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001033 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Thomas Woutersb2137042007-02-01 18:02:27 +00001034 RegressionTests, LengthTransparency,
1035 SubclassWithKwargsTest)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001036 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001037
1038 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001039 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001040 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001041 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +00001042 for i in range(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001043 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001044 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001045 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001046 print(counts)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001047
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001048 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001049 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001050
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001051if __name__ == "__main__":
1052 test_main(verbose=True)