blob: b4c0a8bff2c1ac0c2bf33c3e20f085c7570d9ce8 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001import unittest
2from test import test_support
3from itertools import *
Raymond Hettinger2012f172003-02-07 05:32:58 +00004import sys
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00005import operator
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00006import random
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00007
8def onearg(x):
9 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000010 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000011
12def errfunc(*args):
13 'Test function that raises an error'
14 raise ValueError
15
16def gen3():
17 'Non-restartable source sequence'
18 for i in (0, 1, 2):
19 yield i
20
21def isEven(x):
22 'Test predicate'
23 return x%2==0
24
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000025def isOdd(x):
26 'Test predicate'
27 return x%2==1
28
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000029class StopNow:
30 'Class emulating an empty iterable.'
31 def __iter__(self):
32 return self
33 def next(self):
34 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000035
Raymond Hettinger02420702003-06-29 20:36:23 +000036def take(n, seq):
37 'Convenience function for partially consuming a long of infinite iterable'
38 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000039
Raymond Hettinger96ef8112003-02-01 00:10:11 +000040class TestBasicOps(unittest.TestCase):
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000041 def test_chain(self):
42 self.assertEqual(list(chain('abc', 'def')), list('abcdef'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000043 self.assertEqual(list(chain('abc')), list('abc'))
44 self.assertEqual(list(chain('')), [])
Raymond Hettinger02420702003-06-29 20:36:23 +000045 self.assertEqual(take(4, chain('abc', 'def')), list('abcd'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000046 self.assertRaises(TypeError, chain, 2, 3)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000047
Raymond Hettinger96ef8112003-02-01 00:10:11 +000048 def test_count(self):
49 self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
50 self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
Raymond Hettinger02420702003-06-29 20:36:23 +000051 self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
Raymond Hettinger96ef8112003-02-01 00:10:11 +000052 self.assertRaises(TypeError, count, 2, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000053 self.assertRaises(TypeError, count, 'a')
54 c = count(sys.maxint-2) # verify that rollover doesn't crash
55 c.next(); c.next(); c.next(); c.next(); c.next()
Raymond Hettinger96ef8112003-02-01 00:10:11 +000056
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000057 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +000058 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000059 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000060 self.assertRaises(TypeError, cycle)
61 self.assertRaises(TypeError, cycle, 5)
62 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000063
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000064 def test_groupby(self):
65 # Check whether it accepts arguments correctly
66 self.assertEqual([], list(groupby([])))
67 self.assertEqual([], list(groupby([], key=id)))
68 self.assertRaises(TypeError, list, groupby('abc', []))
69 self.assertRaises(TypeError, groupby, None)
70
71 # Check normal input
72 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
73 (2,15,22), (3,16,23), (3,17,23)]
74 dup = []
75 for k, g in groupby(s, lambda r:r[0]):
76 for elem in g:
77 self.assertEqual(k, elem[0])
78 dup.append(elem)
79 self.assertEqual(s, dup)
80
81 # Check nested case
82 dup = []
83 for k, g in groupby(s, lambda r:r[0]):
84 for ik, ig in groupby(g, lambda r:r[2]):
85 for elem in ig:
86 self.assertEqual(k, elem[0])
87 self.assertEqual(ik, elem[2])
88 dup.append(elem)
89 self.assertEqual(s, dup)
90
91 # Check case where inner iterator is not used
92 keys = [k for k, g in groupby(s, lambda r:r[0])]
93 expectedkeys = set([r[0] for r in s])
94 self.assertEqual(set(keys), expectedkeys)
95 self.assertEqual(len(keys), len(expectedkeys))
96
97 # Exercise pipes and filters style
98 s = 'abracadabra'
99 # sort s | uniq
100 r = [k for k, g in groupby(list.sorted(s))]
101 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
102 # sort s | uniq -d
103 r = [k for k, g in groupby(list.sorted(s)) if list(islice(g,1,2))]
104 self.assertEqual(r, ['a', 'b', 'r'])
105 # sort s | uniq -c
106 r = [(len(list(g)), k) for k, g in groupby(list.sorted(s))]
107 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
108 # sort s | uniq -c | sort -rn | head -3
109 r = list.sorted([(len(list(g)) , k) for k, g in groupby(list.sorted(s))], reverse=True)[:3]
110 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
111
112 # iter.next failure
113 class ExpectedError(Exception):
114 pass
115 def delayed_raise(n=0):
116 for i in range(n):
117 yield 'yo'
118 raise ExpectedError
119 def gulp(iterable, keyp=None, func=list):
120 return [func(g) for k, g in groupby(iterable, keyp)]
121
122 # iter.next failure on outer object
123 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
124 # iter.next failure on inner object
125 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
126
127 # __cmp__ failure
128 class DummyCmp:
129 def __cmp__(self, dst):
130 raise ExpectedError
131 s = [DummyCmp(), DummyCmp(), None]
132
133 # __cmp__ failure on outer object
134 self.assertRaises(ExpectedError, gulp, s, func=id)
135 # __cmp__ failure on inner object
136 self.assertRaises(ExpectedError, gulp, s)
137
138 # keyfunc failure
139 def keyfunc(obj):
140 if keyfunc.skip > 0:
141 keyfunc.skip -= 1
142 return obj
143 else:
144 raise ExpectedError
145
146 # keyfunc failure on outer object
147 keyfunc.skip = 0
148 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
149 keyfunc.skip = 1
150 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
151
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000152 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000153 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000154 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000155 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000156 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000157 self.assertRaises(TypeError, ifilter, lambda x:x)
158 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000159 self.assertRaises(TypeError, ifilter, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000160 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000161
162 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000163 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
164 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000165 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000166 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000167 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
168 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000169 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000170 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000171
172 def test_izip(self):
173 ans = [(x,y) for x, y in izip('abc',count())]
174 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000175 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
176 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
Raymond Hettinger02420702003-06-29 20:36:23 +0000177 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000178 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000179 self.assertEqual(list(izip()), zip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000180 self.assertRaises(TypeError, izip, 3)
181 self.assertRaises(TypeError, izip, range(3), 3)
182 # Check tuple re-use (implementation detail)
183 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
184 zip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000185 self.assertEqual([pair for pair in izip('abc', 'def')],
186 zip('abc', 'def'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000187 ids = map(id, izip('abc', 'def'))
188 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000189 ids = map(id, list(izip('abc', 'def')))
190 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000191
192 def test_repeat(self):
193 self.assertEqual(zip(xrange(3),repeat('a')),
194 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000195 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000196 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000197 self.assertEqual(list(repeat('a', 0)), [])
198 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000199 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000200 self.assertRaises(TypeError, repeat, None, 3, 4)
201 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000202
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000203 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000204 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
205 [0**1, 1**2, 2**3])
206 self.assertEqual(list(imap(None, 'abc', range(5))),
207 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000208 self.assertEqual(list(imap(None, 'abc', count())),
209 [('a',0),('b',1),('c',2)])
210 self.assertEqual(take(2,imap(None, 'abc', count())),
211 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000212 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000213 self.assertRaises(TypeError, imap)
214 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000215 self.assertRaises(TypeError, imap(10, range(5)).next)
216 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
217 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000218
219 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000220 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
221 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000222 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
223 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000224 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger2012f172003-02-07 05:32:58 +0000225 self.assertRaises(TypeError, list, starmap(operator.pow, [[4,5]]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000226 self.assertRaises(TypeError, starmap)
227 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
228 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
229 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
230 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000231
232 def test_islice(self):
233 for args in [ # islice(args) should agree with range(args)
234 (10, 20, 3),
235 (10, 3, 20),
236 (10, 20),
237 (10, 3),
238 (20,)
239 ]:
240 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
241
242 for args, tgtargs in [ # Stop when seqn is exhausted
243 ((10, 110, 3), ((10, 100, 3))),
244 ((10, 110), ((10, 100))),
245 ((110,), (100,))
246 ]:
247 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
248
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000249 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000250 self.assertEqual(list(islice(xrange(10), None)), range(10))
251 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
252 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
253
254 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000255 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000256 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
257 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
258 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
259 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
260 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000261 self.assertRaises(ValueError, islice, xrange(10), 'a')
262 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
263 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
264 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
265 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Raymond Hettinger2012f172003-02-07 05:32:58 +0000266 self.assertEqual(len(list(islice(count(), 1, 10, sys.maxint))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000267
268 def test_takewhile(self):
269 data = [1, 3, 5, 20, 2, 4, 6, 8]
270 underten = lambda x: x<10
271 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000272 self.assertEqual(list(takewhile(underten, [])), [])
273 self.assertRaises(TypeError, takewhile)
274 self.assertRaises(TypeError, takewhile, operator.pow)
275 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
276 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
277 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000278
279 def test_dropwhile(self):
280 data = [1, 3, 5, 20, 2, 4, 6, 8]
281 underten = lambda x: x<10
282 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000283 self.assertEqual(list(dropwhile(underten, [])), [])
284 self.assertRaises(TypeError, dropwhile)
285 self.assertRaises(TypeError, dropwhile, operator.pow)
286 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
287 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
288 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000289
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000290 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000291 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000292 def irange(n):
293 for i in xrange(n):
294 yield i
295
296 a, b = tee([]) # test empty iterator
297 self.assertEqual(list(a), [])
298 self.assertEqual(list(b), [])
299
300 a, b = tee(irange(n)) # test 100% interleaved
301 self.assertEqual(zip(a,b), zip(range(n),range(n)))
302
303 a, b = tee(irange(n)) # test 0% interleaved
304 self.assertEqual(list(a), range(n))
305 self.assertEqual(list(b), range(n))
306
307 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000308 for i in xrange(100):
309 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000310 del a
311 self.assertEqual(list(b), range(n))
312
313 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000314 for i in xrange(100):
315 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000316 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000317 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000318
319 for j in xrange(5): # test randomly interleaved
320 order = [0]*n + [1]*n
321 random.shuffle(order)
322 lists = ([], [])
323 its = tee(irange(n))
324 for i in order:
325 value = its[i].next()
326 lists[i].append(value)
327 self.assertEqual(lists[0], range(n))
328 self.assertEqual(lists[1], range(n))
329
Raymond Hettingerad983e72003-11-12 14:32:26 +0000330 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000331 self.assertRaises(TypeError, tee)
332 self.assertRaises(TypeError, tee, 3)
333 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000334 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000335
Raymond Hettingerad983e72003-11-12 14:32:26 +0000336 # tee object should be instantiable
337 a, b = tee('abc')
338 c = type(a)('def')
339 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000340
Raymond Hettingerad983e72003-11-12 14:32:26 +0000341 # test long-lagged and multi-way split
342 a, b, c = tee(xrange(2000), 3)
343 for i in xrange(100):
344 self.assertEqual(a.next(), i)
345 self.assertEqual(list(b), range(2000))
346 self.assertEqual([c.next(), c.next()], range(2))
347 self.assertEqual(list(a), range(100,2000))
348 self.assertEqual(list(c), range(2,2000))
349
350 # tee pass-through to copyable iterator
351 a, b = tee('abc')
352 c, d = tee(a)
353 self.assert_(a is c)
354
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000355
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000356 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000357 self.assertRaises(StopIteration, izip().next)
358
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000359 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000360 self.assertRaises(StopIteration, f([]).next)
361 self.assertRaises(StopIteration, f(StopNow()).next)
362
363 self.assertRaises(StopIteration, islice([], None).next)
364 self.assertRaises(StopIteration, islice(StopNow(), None).next)
365
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000366 p, q = tee([])
367 self.assertRaises(StopIteration, p.next)
368 self.assertRaises(StopIteration, q.next)
369 p, q = tee(StopNow())
370 self.assertRaises(StopIteration, p.next)
371 self.assertRaises(StopIteration, q.next)
372
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000373 self.assertRaises(StopIteration, repeat(None, 0).next)
374
375 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
376 self.assertRaises(StopIteration, f(lambda x:x, []).next)
377 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
378
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000379class TestGC(unittest.TestCase):
380
381 def makecycle(self, iterator, container):
382 container.append(iterator)
383 iterator.next()
384 del container, iterator
385
386 def test_chain(self):
387 a = []
388 self.makecycle(chain(a), a)
389
390 def test_cycle(self):
391 a = []
392 self.makecycle(cycle([a]*2), a)
393
394 def test_ifilter(self):
395 a = []
396 self.makecycle(ifilter(lambda x:True, [a]*2), a)
397
398 def test_ifilterfalse(self):
399 a = []
400 self.makecycle(ifilterfalse(lambda x:False, a), a)
401
402 def test_izip(self):
403 a = []
404 self.makecycle(izip([a]*2, [a]*3), a)
405
406 def test_imap(self):
407 a = []
408 self.makecycle(imap(lambda x:x, [a]*2), a)
409
410 def test_islice(self):
411 a = []
412 self.makecycle(islice([a]*2, None), a)
413
414 def test_starmap(self):
415 a = []
416 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
417
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000418def R(seqn):
419 'Regular generator'
420 for i in seqn:
421 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000422
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000423class G:
424 'Sequence using __getitem__'
425 def __init__(self, seqn):
426 self.seqn = seqn
427 def __getitem__(self, i):
428 return self.seqn[i]
429
430class I:
431 'Sequence using iterator protocol'
432 def __init__(self, seqn):
433 self.seqn = seqn
434 self.i = 0
435 def __iter__(self):
436 return self
437 def next(self):
438 if self.i >= len(self.seqn): raise StopIteration
439 v = self.seqn[self.i]
440 self.i += 1
441 return v
442
443class Ig:
444 'Sequence using iterator protocol defined with a generator'
445 def __init__(self, seqn):
446 self.seqn = seqn
447 self.i = 0
448 def __iter__(self):
449 for val in self.seqn:
450 yield val
451
452class X:
453 'Missing __getitem__ and __iter__'
454 def __init__(self, seqn):
455 self.seqn = seqn
456 self.i = 0
457 def next(self):
458 if self.i >= len(self.seqn): raise StopIteration
459 v = self.seqn[self.i]
460 self.i += 1
461 return v
462
463class N:
464 'Iterator missing next()'
465 def __init__(self, seqn):
466 self.seqn = seqn
467 self.i = 0
468 def __iter__(self):
469 return self
470
471class E:
472 'Test propagation of exceptions'
473 def __init__(self, seqn):
474 self.seqn = seqn
475 self.i = 0
476 def __iter__(self):
477 return self
478 def next(self):
479 3/0
480
481class S:
482 'Test immediate stop'
483 def __init__(self, seqn):
484 pass
485 def __iter__(self):
486 return self
487 def next(self):
488 raise StopIteration
489
490def L(seqn):
491 'Test multiple tiers of iterators'
492 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
493
494
495class TestVariousIteratorArgs(unittest.TestCase):
496
497 def test_chain(self):
498 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
499 for g in (G, I, Ig, S, L, R):
500 self.assertEqual(list(chain(g(s))), list(g(s)))
501 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
502 self.assertRaises(TypeError, chain, X(s))
503 self.assertRaises(TypeError, list, chain(N(s)))
504 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
505
506 def test_cycle(self):
507 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
508 for g in (G, I, Ig, S, L, R):
509 tgtlen = len(s) * 3
510 expected = list(g(s))*3
511 actual = list(islice(cycle(g(s)), tgtlen))
512 self.assertEqual(actual, expected)
513 self.assertRaises(TypeError, cycle, X(s))
514 self.assertRaises(TypeError, list, cycle(N(s)))
515 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
516
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000517 def test_groupby(self):
518 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
519 for g in (G, I, Ig, S, L, R):
520 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
521 self.assertRaises(TypeError, groupby, X(s))
522 self.assertRaises(TypeError, list, groupby(N(s)))
523 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
524
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000525 def test_ifilter(self):
526 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
527 for g in (G, I, Ig, S, L, R):
528 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
529 self.assertRaises(TypeError, ifilter, isEven, X(s))
530 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
531 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
532
533 def test_ifilterfalse(self):
534 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
535 for g in (G, I, Ig, S, L, R):
536 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
537 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
538 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
539 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
540
541 def test_izip(self):
542 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
543 for g in (G, I, Ig, S, L, R):
544 self.assertEqual(list(izip(g(s))), zip(g(s)))
545 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
546 self.assertRaises(TypeError, izip, X(s))
547 self.assertRaises(TypeError, list, izip(N(s)))
548 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
549
550 def test_imap(self):
551 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
552 for g in (G, I, Ig, S, L, R):
553 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
554 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
555 self.assertRaises(TypeError, imap, onearg, X(s))
556 self.assertRaises(TypeError, list, imap(onearg, N(s)))
557 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
558
559 def test_islice(self):
560 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
561 for g in (G, I, Ig, S, L, R):
562 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
563 self.assertRaises(TypeError, islice, X(s), 10)
564 self.assertRaises(TypeError, list, islice(N(s), 10))
565 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
566
567 def test_starmap(self):
568 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
569 for g in (G, I, Ig, S, L, R):
570 ss = zip(s, s)
571 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
572 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
573 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
574 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
575
576 def test_takewhile(self):
577 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
578 for g in (G, I, Ig, S, L, R):
579 tgt = []
580 for elem in g(s):
581 if not isEven(elem): break
582 tgt.append(elem)
583 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
584 self.assertRaises(TypeError, takewhile, isEven, X(s))
585 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
586 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
587
588 def test_dropwhile(self):
589 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
590 for g in (G, I, Ig, S, L, R):
591 tgt = []
592 for elem in g(s):
593 if not tgt and isOdd(elem): continue
594 tgt.append(elem)
595 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
596 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
597 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
598 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
599
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000600 def test_tee(self):
601 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
602 for g in (G, I, Ig, S, L, R):
603 it1, it2 = tee(g(s))
604 self.assertEqual(list(it1), list(g(s)))
605 self.assertEqual(list(it2), list(g(s)))
606 self.assertRaises(TypeError, tee, X(s))
607 self.assertRaises(TypeError, list, tee(N(s))[0])
608 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
609
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000610class RegressionTests(unittest.TestCase):
611
612 def test_sf_793826(self):
613 # Fix Armin Rigo's successful efforts to wreak havoc
614
615 def mutatingtuple(tuple1, f, tuple2):
616 # this builds a tuple t which is a copy of tuple1,
617 # then calls f(t), then mutates t to be equal to tuple2
618 # (needs len(tuple1) == len(tuple2)).
619 def g(value, first=[1]):
620 if first:
621 del first[:]
622 f(z.next())
623 return value
624 items = list(tuple2)
625 items[1:1] = list(tuple1)
626 gen = imap(g, items)
627 z = izip(*[gen]*len(tuple1))
628 z.next()
629
630 def f(t):
631 global T
632 T = t
633 first[:] = list(T)
634
635 first = []
636 mutatingtuple((1,2,3), f, (4,5,6))
637 second = list(T)
638 self.assertEqual(first, second)
639
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000640
641libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000642
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000643
644>>> amounts = [120.15, 764.05, 823.14]
645>>> for checknum, amount in izip(count(1200), amounts):
646... print 'Check %d is for $%.2f' % (checknum, amount)
647...
648Check 1200 is for $120.15
649Check 1201 is for $764.05
650Check 1202 is for $823.14
651
652>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000653>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
654... print cube
655...
6561
6578
65827
659
660>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +0000661>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000662... print name.title()
663...
664Alex
665Laura
666Martin
667Walter
668Samuele
669
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000670>>> from operator import itemgetter
671>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
672>>> di = list.sorted(d.iteritems(), key=itemgetter(1))
673>>> for k, g in groupby(di, itemgetter(1)):
674... print k, map(itemgetter(0), g)
675...
6761 ['a', 'c', 'e']
6772 ['b', 'd', 'f']
6783 ['g']
679
Raymond Hettingera098b332003-09-08 23:58:40 +0000680>>> def take(n, seq):
681... return list(islice(seq, n))
682
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000683>>> def enumerate(iterable):
684... return izip(count(), iterable)
685
686>>> def tabulate(function):
687... "Return function(0), function(1), ..."
688... return imap(function, count())
689
690>>> def iteritems(mapping):
691... return izip(mapping.iterkeys(), mapping.itervalues())
692
693>>> def nth(iterable, n):
694... "Returns the nth item"
Raymond Hettinger60eca932003-02-09 06:40:58 +0000695... return list(islice(iterable, n, n+1))
696
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000697>>> def all(seq, pred=bool):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000698... "Returns True if pred(x) is True for every element in the iterable"
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000699... return False not in imap(pred, seq)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000700
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000701>>> def any(seq, pred=bool):
Raymond Hettinger02420702003-06-29 20:36:23 +0000702... "Returns True if pred(x) is True for at least one element in the iterable"
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000703... return True in imap(pred, seq)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000704
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000705>>> def no(seq, pred=bool):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000706... "Returns True if pred(x) is False for every element in the iterable"
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000707... return True not in imap(pred, seq)
708
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000709>>> def quantify(seq, pred=bool):
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000710... "Count how many times the predicate is True in the sequence"
711... return sum(imap(pred, seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000712
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000713>>> def padnone(seq):
714... "Returns the sequence elements and then returns None indefinitely"
715... return chain(seq, repeat(None))
716
717>>> def ncycles(seq, n):
718... "Returns the sequence elements n times"
719... return chain(*repeat(seq, n))
720
721>>> def dotproduct(vec1, vec2):
722... return sum(imap(operator.mul, vec1, vec2))
723
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000724>>> def flatten(listOfLists):
725... return list(chain(*listOfLists))
726
727>>> def repeatfunc(func, times=None, *args):
728... "Repeat calls to func with specified arguments."
729... " Example: repeatfunc(random.random)"
730... if times is None:
731... return starmap(func, repeat(args))
732... else:
733... return starmap(func, repeat(args, times))
734
Raymond Hettingerd591f662003-10-26 15:34:50 +0000735>>> def pairwise(iterable):
736... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
737... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000738... try:
739... b.next()
740... except StopIteration:
741... pass
742... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000743
744This is not part of the examples but it tests to make sure the definitions
745perform as purported.
746
Raymond Hettingera098b332003-09-08 23:58:40 +0000747>>> take(10, count())
748[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
749
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000750>>> list(enumerate('abc'))
751[(0, 'a'), (1, 'b'), (2, 'c')]
752
753>>> list(islice(tabulate(lambda x: 2*x), 4))
754[0, 2, 4, 6]
755
756>>> nth('abcde', 3)
757['d']
758
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000759>>> all([2, 4, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000760True
761
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000762>>> all([2, 3, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000763False
764
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000765>>> any([2, 4, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000766True
767
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000768>>> any([1, 3, 5, 9], lambda x: x%2==0,)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000769False
770
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000771>>> no([1, 3, 5, 9], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000772True
773
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000774>>> no([1, 2, 5, 9], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000775False
776
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000777>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +000077850
779
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000780>>> a = [[1, 2, 3], [4, 5, 6]]
781>>> flatten(a)
782[1, 2, 3, 4, 5, 6]
783
784>>> list(repeatfunc(pow, 5, 2, 3))
785[8, 8, 8, 8, 8]
786
787>>> import random
788>>> take(5, imap(int, repeatfunc(random.random)))
789[0, 0, 0, 0, 0]
790
Raymond Hettingerd591f662003-10-26 15:34:50 +0000791>>> list(pairwise('abcd'))
792[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000793
Raymond Hettingerd591f662003-10-26 15:34:50 +0000794>>> list(pairwise([]))
795[]
796
797>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +0000798[]
799
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000800>>> list(islice(padnone('abc'), 0, 6))
801['a', 'b', 'c', None, None, None]
802
803>>> list(ncycles('abc', 3))
804['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
805
806>>> dotproduct([1,2,3], [4,5,6])
80732
808
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000809"""
810
811__test__ = {'libreftest' : libreftest}
812
813def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000814 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
815 RegressionTests)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000816 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000817
818 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000819 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +0000820 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000821 counts = [None] * 5
822 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000823 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +0000824 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000825 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000826 print counts
827
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000828 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +0000829 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000830
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000831if __name__ == "__main__":
832 test_main(verbose=True)