blob: 543acc199c25252672e5cf2fbcd127a7dd7c65d3 [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 Hettinger96ef8112003-02-01 00:10:11 +000064 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +000065 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +000066 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +000067 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +000068 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000069 self.assertRaises(TypeError, ifilter, lambda x:x)
70 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +000071 self.assertRaises(TypeError, ifilter, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000072 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
Raymond Hettinger60eca932003-02-09 06:40:58 +000073
74 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +000075 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
76 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +000077 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +000078 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000079 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
80 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +000081 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000082 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +000083
84 def test_izip(self):
85 ans = [(x,y) for x, y in izip('abc',count())]
86 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000087 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
88 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
Raymond Hettinger02420702003-06-29 20:36:23 +000089 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000090 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
Raymond Hettingerb5a42082003-08-08 05:10:41 +000091 self.assertEqual(list(izip()), zip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000092 self.assertRaises(TypeError, izip, 3)
93 self.assertRaises(TypeError, izip, range(3), 3)
94 # Check tuple re-use (implementation detail)
95 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
96 zip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000097 self.assertEqual([pair for pair in izip('abc', 'def')],
98 zip('abc', 'def'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000099 ids = map(id, izip('abc', 'def'))
100 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000101 ids = map(id, list(izip('abc', 'def')))
102 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000103
104 def test_repeat(self):
105 self.assertEqual(zip(xrange(3),repeat('a')),
106 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000107 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000108 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000109 self.assertEqual(list(repeat('a', 0)), [])
110 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000111 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000112 self.assertRaises(TypeError, repeat, None, 3, 4)
113 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000114
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000115 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000116 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
117 [0**1, 1**2, 2**3])
118 self.assertEqual(list(imap(None, 'abc', range(5))),
119 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000120 self.assertEqual(list(imap(None, 'abc', count())),
121 [('a',0),('b',1),('c',2)])
122 self.assertEqual(take(2,imap(None, 'abc', count())),
123 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000124 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000125 self.assertRaises(TypeError, imap)
126 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000127 self.assertRaises(TypeError, imap(10, range(5)).next)
128 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
129 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000130
131 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000132 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
133 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000134 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
135 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000136 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger2012f172003-02-07 05:32:58 +0000137 self.assertRaises(TypeError, list, starmap(operator.pow, [[4,5]]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000138 self.assertRaises(TypeError, starmap)
139 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
140 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
141 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
142 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000143
144 def test_islice(self):
145 for args in [ # islice(args) should agree with range(args)
146 (10, 20, 3),
147 (10, 3, 20),
148 (10, 20),
149 (10, 3),
150 (20,)
151 ]:
152 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
153
154 for args, tgtargs in [ # Stop when seqn is exhausted
155 ((10, 110, 3), ((10, 100, 3))),
156 ((10, 110), ((10, 100))),
157 ((110,), (100,))
158 ]:
159 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
160
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000161 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000162 self.assertEqual(list(islice(xrange(10), None)), range(10))
163 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
164 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
165
166 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000167 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000168 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
169 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
170 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
171 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
172 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000173 self.assertRaises(ValueError, islice, xrange(10), 'a')
174 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
175 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
176 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
177 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Raymond Hettinger2012f172003-02-07 05:32:58 +0000178 self.assertEqual(len(list(islice(count(), 1, 10, sys.maxint))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000179
180 def test_takewhile(self):
181 data = [1, 3, 5, 20, 2, 4, 6, 8]
182 underten = lambda x: x<10
183 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000184 self.assertEqual(list(takewhile(underten, [])), [])
185 self.assertRaises(TypeError, takewhile)
186 self.assertRaises(TypeError, takewhile, operator.pow)
187 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
188 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
189 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000190
191 def test_dropwhile(self):
192 data = [1, 3, 5, 20, 2, 4, 6, 8]
193 underten = lambda x: x<10
194 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000195 self.assertEqual(list(dropwhile(underten, [])), [])
196 self.assertRaises(TypeError, dropwhile)
197 self.assertRaises(TypeError, dropwhile, operator.pow)
198 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
199 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
200 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000201
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000202 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000203 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000204 def irange(n):
205 for i in xrange(n):
206 yield i
207
208 a, b = tee([]) # test empty iterator
209 self.assertEqual(list(a), [])
210 self.assertEqual(list(b), [])
211
212 a, b = tee(irange(n)) # test 100% interleaved
213 self.assertEqual(zip(a,b), zip(range(n),range(n)))
214
215 a, b = tee(irange(n)) # test 0% interleaved
216 self.assertEqual(list(a), range(n))
217 self.assertEqual(list(b), range(n))
218
219 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000220 for i in xrange(100):
221 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000222 del a
223 self.assertEqual(list(b), range(n))
224
225 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000226 for i in xrange(100):
227 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000228 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000229 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000230
231 for j in xrange(5): # test randomly interleaved
232 order = [0]*n + [1]*n
233 random.shuffle(order)
234 lists = ([], [])
235 its = tee(irange(n))
236 for i in order:
237 value = its[i].next()
238 lists[i].append(value)
239 self.assertEqual(lists[0], range(n))
240 self.assertEqual(lists[1], range(n))
241
Raymond Hettingerad983e72003-11-12 14:32:26 +0000242 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000243 self.assertRaises(TypeError, tee)
244 self.assertRaises(TypeError, tee, 3)
245 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000246 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000247
Raymond Hettingerad983e72003-11-12 14:32:26 +0000248 # tee object should be instantiable
249 a, b = tee('abc')
250 c = type(a)('def')
251 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000252
Raymond Hettingerad983e72003-11-12 14:32:26 +0000253 # test long-lagged and multi-way split
254 a, b, c = tee(xrange(2000), 3)
255 for i in xrange(100):
256 self.assertEqual(a.next(), i)
257 self.assertEqual(list(b), range(2000))
258 self.assertEqual([c.next(), c.next()], range(2))
259 self.assertEqual(list(a), range(100,2000))
260 self.assertEqual(list(c), range(2,2000))
261
262 # tee pass-through to copyable iterator
263 a, b = tee('abc')
264 c, d = tee(a)
265 self.assert_(a is c)
266
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000267
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000268 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000269 self.assertRaises(StopIteration, izip().next)
270
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000271 for f in (chain, cycle, izip):
272 self.assertRaises(StopIteration, f([]).next)
273 self.assertRaises(StopIteration, f(StopNow()).next)
274
275 self.assertRaises(StopIteration, islice([], None).next)
276 self.assertRaises(StopIteration, islice(StopNow(), None).next)
277
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000278 p, q = tee([])
279 self.assertRaises(StopIteration, p.next)
280 self.assertRaises(StopIteration, q.next)
281 p, q = tee(StopNow())
282 self.assertRaises(StopIteration, p.next)
283 self.assertRaises(StopIteration, q.next)
284
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000285 self.assertRaises(StopIteration, repeat(None, 0).next)
286
287 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
288 self.assertRaises(StopIteration, f(lambda x:x, []).next)
289 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
290
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000291class TestGC(unittest.TestCase):
292
293 def makecycle(self, iterator, container):
294 container.append(iterator)
295 iterator.next()
296 del container, iterator
297
298 def test_chain(self):
299 a = []
300 self.makecycle(chain(a), a)
301
302 def test_cycle(self):
303 a = []
304 self.makecycle(cycle([a]*2), a)
305
306 def test_ifilter(self):
307 a = []
308 self.makecycle(ifilter(lambda x:True, [a]*2), a)
309
310 def test_ifilterfalse(self):
311 a = []
312 self.makecycle(ifilterfalse(lambda x:False, a), a)
313
314 def test_izip(self):
315 a = []
316 self.makecycle(izip([a]*2, [a]*3), a)
317
318 def test_imap(self):
319 a = []
320 self.makecycle(imap(lambda x:x, [a]*2), a)
321
322 def test_islice(self):
323 a = []
324 self.makecycle(islice([a]*2, None), a)
325
326 def test_starmap(self):
327 a = []
328 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
329
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000330def R(seqn):
331 'Regular generator'
332 for i in seqn:
333 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000334
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000335class G:
336 'Sequence using __getitem__'
337 def __init__(self, seqn):
338 self.seqn = seqn
339 def __getitem__(self, i):
340 return self.seqn[i]
341
342class I:
343 'Sequence using iterator protocol'
344 def __init__(self, seqn):
345 self.seqn = seqn
346 self.i = 0
347 def __iter__(self):
348 return self
349 def next(self):
350 if self.i >= len(self.seqn): raise StopIteration
351 v = self.seqn[self.i]
352 self.i += 1
353 return v
354
355class Ig:
356 'Sequence using iterator protocol defined with a generator'
357 def __init__(self, seqn):
358 self.seqn = seqn
359 self.i = 0
360 def __iter__(self):
361 for val in self.seqn:
362 yield val
363
364class X:
365 'Missing __getitem__ and __iter__'
366 def __init__(self, seqn):
367 self.seqn = seqn
368 self.i = 0
369 def next(self):
370 if self.i >= len(self.seqn): raise StopIteration
371 v = self.seqn[self.i]
372 self.i += 1
373 return v
374
375class N:
376 'Iterator missing next()'
377 def __init__(self, seqn):
378 self.seqn = seqn
379 self.i = 0
380 def __iter__(self):
381 return self
382
383class E:
384 'Test propagation of exceptions'
385 def __init__(self, seqn):
386 self.seqn = seqn
387 self.i = 0
388 def __iter__(self):
389 return self
390 def next(self):
391 3/0
392
393class S:
394 'Test immediate stop'
395 def __init__(self, seqn):
396 pass
397 def __iter__(self):
398 return self
399 def next(self):
400 raise StopIteration
401
402def L(seqn):
403 'Test multiple tiers of iterators'
404 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
405
406
407class TestVariousIteratorArgs(unittest.TestCase):
408
409 def test_chain(self):
410 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
411 for g in (G, I, Ig, S, L, R):
412 self.assertEqual(list(chain(g(s))), list(g(s)))
413 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
414 self.assertRaises(TypeError, chain, X(s))
415 self.assertRaises(TypeError, list, chain(N(s)))
416 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
417
418 def test_cycle(self):
419 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
420 for g in (G, I, Ig, S, L, R):
421 tgtlen = len(s) * 3
422 expected = list(g(s))*3
423 actual = list(islice(cycle(g(s)), tgtlen))
424 self.assertEqual(actual, expected)
425 self.assertRaises(TypeError, cycle, X(s))
426 self.assertRaises(TypeError, list, cycle(N(s)))
427 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
428
429 def test_ifilter(self):
430 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
431 for g in (G, I, Ig, S, L, R):
432 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
433 self.assertRaises(TypeError, ifilter, isEven, X(s))
434 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
435 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
436
437 def test_ifilterfalse(self):
438 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
439 for g in (G, I, Ig, S, L, R):
440 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
441 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
442 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
443 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
444
445 def test_izip(self):
446 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
447 for g in (G, I, Ig, S, L, R):
448 self.assertEqual(list(izip(g(s))), zip(g(s)))
449 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
450 self.assertRaises(TypeError, izip, X(s))
451 self.assertRaises(TypeError, list, izip(N(s)))
452 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
453
454 def test_imap(self):
455 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
456 for g in (G, I, Ig, S, L, R):
457 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
458 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
459 self.assertRaises(TypeError, imap, onearg, X(s))
460 self.assertRaises(TypeError, list, imap(onearg, N(s)))
461 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
462
463 def test_islice(self):
464 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
465 for g in (G, I, Ig, S, L, R):
466 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
467 self.assertRaises(TypeError, islice, X(s), 10)
468 self.assertRaises(TypeError, list, islice(N(s), 10))
469 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
470
471 def test_starmap(self):
472 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
473 for g in (G, I, Ig, S, L, R):
474 ss = zip(s, s)
475 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
476 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
477 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
478 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
479
480 def test_takewhile(self):
481 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
482 for g in (G, I, Ig, S, L, R):
483 tgt = []
484 for elem in g(s):
485 if not isEven(elem): break
486 tgt.append(elem)
487 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
488 self.assertRaises(TypeError, takewhile, isEven, X(s))
489 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
490 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
491
492 def test_dropwhile(self):
493 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
494 for g in (G, I, Ig, S, L, R):
495 tgt = []
496 for elem in g(s):
497 if not tgt and isOdd(elem): continue
498 tgt.append(elem)
499 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
500 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
501 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
502 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
503
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000504 def test_tee(self):
505 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
506 for g in (G, I, Ig, S, L, R):
507 it1, it2 = tee(g(s))
508 self.assertEqual(list(it1), list(g(s)))
509 self.assertEqual(list(it2), list(g(s)))
510 self.assertRaises(TypeError, tee, X(s))
511 self.assertRaises(TypeError, list, tee(N(s))[0])
512 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
513
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000514class RegressionTests(unittest.TestCase):
515
516 def test_sf_793826(self):
517 # Fix Armin Rigo's successful efforts to wreak havoc
518
519 def mutatingtuple(tuple1, f, tuple2):
520 # this builds a tuple t which is a copy of tuple1,
521 # then calls f(t), then mutates t to be equal to tuple2
522 # (needs len(tuple1) == len(tuple2)).
523 def g(value, first=[1]):
524 if first:
525 del first[:]
526 f(z.next())
527 return value
528 items = list(tuple2)
529 items[1:1] = list(tuple1)
530 gen = imap(g, items)
531 z = izip(*[gen]*len(tuple1))
532 z.next()
533
534 def f(t):
535 global T
536 T = t
537 first[:] = list(T)
538
539 first = []
540 mutatingtuple((1,2,3), f, (4,5,6))
541 second = list(T)
542 self.assertEqual(first, second)
543
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000544
545libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000546
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000547
548>>> amounts = [120.15, 764.05, 823.14]
549>>> for checknum, amount in izip(count(1200), amounts):
550... print 'Check %d is for $%.2f' % (checknum, amount)
551...
552Check 1200 is for $120.15
553Check 1201 is for $764.05
554Check 1202 is for $823.14
555
556>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000557>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
558... print cube
559...
5601
5618
56227
563
564>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +0000565>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000566... print name.title()
567...
568Alex
569Laura
570Martin
571Walter
572Samuele
573
Raymond Hettingera098b332003-09-08 23:58:40 +0000574>>> def take(n, seq):
575... return list(islice(seq, n))
576
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000577>>> def enumerate(iterable):
578... return izip(count(), iterable)
579
580>>> def tabulate(function):
581... "Return function(0), function(1), ..."
582... return imap(function, count())
583
584>>> def iteritems(mapping):
585... return izip(mapping.iterkeys(), mapping.itervalues())
586
587>>> def nth(iterable, n):
588... "Returns the nth item"
Raymond Hettinger60eca932003-02-09 06:40:58 +0000589... return list(islice(iterable, n, n+1))
590
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000591>>> def all(seq, pred=bool):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000592... "Returns True if pred(x) is True for every element in the iterable"
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000593... return False not in imap(pred, seq)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000594
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000595>>> def any(seq, pred=bool):
Raymond Hettinger02420702003-06-29 20:36:23 +0000596... "Returns True if pred(x) is True for at least one element in the iterable"
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000597... return True in imap(pred, seq)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000598
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000599>>> def no(seq, pred=bool):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000600... "Returns True if pred(x) is False for every element in the iterable"
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000601... return True not in imap(pred, seq)
602
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000603>>> def quantify(seq, pred=bool):
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000604... "Count how many times the predicate is True in the sequence"
605... return sum(imap(pred, seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000606
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000607>>> def padnone(seq):
608... "Returns the sequence elements and then returns None indefinitely"
609... return chain(seq, repeat(None))
610
611>>> def ncycles(seq, n):
612... "Returns the sequence elements n times"
613... return chain(*repeat(seq, n))
614
615>>> def dotproduct(vec1, vec2):
616... return sum(imap(operator.mul, vec1, vec2))
617
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000618>>> def flatten(listOfLists):
619... return list(chain(*listOfLists))
620
621>>> def repeatfunc(func, times=None, *args):
622... "Repeat calls to func with specified arguments."
623... " Example: repeatfunc(random.random)"
624... if times is None:
625... return starmap(func, repeat(args))
626... else:
627... return starmap(func, repeat(args, times))
628
Raymond Hettingerd591f662003-10-26 15:34:50 +0000629>>> def pairwise(iterable):
630... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
631... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000632... try:
633... b.next()
634... except StopIteration:
635... pass
636... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000637
638This is not part of the examples but it tests to make sure the definitions
639perform as purported.
640
Raymond Hettingera098b332003-09-08 23:58:40 +0000641>>> take(10, count())
642[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
643
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000644>>> list(enumerate('abc'))
645[(0, 'a'), (1, 'b'), (2, 'c')]
646
647>>> list(islice(tabulate(lambda x: 2*x), 4))
648[0, 2, 4, 6]
649
650>>> nth('abcde', 3)
651['d']
652
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000653>>> all([2, 4, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000654True
655
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000656>>> all([2, 3, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000657False
658
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000659>>> any([2, 4, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000660True
661
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000662>>> any([1, 3, 5, 9], lambda x: x%2==0,)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000663False
664
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000665>>> no([1, 3, 5, 9], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000666True
667
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000668>>> no([1, 2, 5, 9], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000669False
670
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000671>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +000067250
673
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000674>>> a = [[1, 2, 3], [4, 5, 6]]
675>>> flatten(a)
676[1, 2, 3, 4, 5, 6]
677
678>>> list(repeatfunc(pow, 5, 2, 3))
679[8, 8, 8, 8, 8]
680
681>>> import random
682>>> take(5, imap(int, repeatfunc(random.random)))
683[0, 0, 0, 0, 0]
684
Raymond Hettingerd591f662003-10-26 15:34:50 +0000685>>> list(pairwise('abcd'))
686[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000687
Raymond Hettingerd591f662003-10-26 15:34:50 +0000688>>> list(pairwise([]))
689[]
690
691>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +0000692[]
693
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000694>>> list(islice(padnone('abc'), 0, 6))
695['a', 'b', 'c', None, None, None]
696
697>>> list(ncycles('abc', 3))
698['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
699
700>>> dotproduct([1,2,3], [4,5,6])
70132
702
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000703"""
704
705__test__ = {'libreftest' : libreftest}
706
707def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000708 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
709 RegressionTests)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000710 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000711
712 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000713 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +0000714 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000715 counts = [None] * 5
716 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000717 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +0000718 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000719 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000720 print counts
721
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000722 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +0000723 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000724
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000725if __name__ == "__main__":
726 test_main(verbose=True)