blob: ce03b1aefb6b5e708074f59762c06fdeea37c3d0 [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):
203 n = 100
204 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
220 self.assertEqual(a.next(), 0)
221 self.assertEqual(a.next(), 1)
222 del a
223 self.assertEqual(list(b), range(n))
224
225 a, b = tee(irange(n)) # test dealloc of trailing iterator
226 self.assertEqual(a.next(), 0)
227 self.assertEqual(a.next(), 1)
228 del b
229 self.assertEqual(list(a), range(2, n))
230
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
242 self.assertRaises(TypeError, tee)
243 self.assertRaises(TypeError, tee, 3)
244 self.assertRaises(TypeError, tee, [1,2], 'x')
245
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000246 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000247 self.assertRaises(StopIteration, izip().next)
248
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000249 for f in (chain, cycle, izip):
250 self.assertRaises(StopIteration, f([]).next)
251 self.assertRaises(StopIteration, f(StopNow()).next)
252
253 self.assertRaises(StopIteration, islice([], None).next)
254 self.assertRaises(StopIteration, islice(StopNow(), None).next)
255
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000256 p, q = tee([])
257 self.assertRaises(StopIteration, p.next)
258 self.assertRaises(StopIteration, q.next)
259 p, q = tee(StopNow())
260 self.assertRaises(StopIteration, p.next)
261 self.assertRaises(StopIteration, q.next)
262
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000263 self.assertRaises(StopIteration, repeat(None, 0).next)
264
265 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
266 self.assertRaises(StopIteration, f(lambda x:x, []).next)
267 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
268
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000269class TestGC(unittest.TestCase):
270
271 def makecycle(self, iterator, container):
272 container.append(iterator)
273 iterator.next()
274 del container, iterator
275
276 def test_chain(self):
277 a = []
278 self.makecycle(chain(a), a)
279
280 def test_cycle(self):
281 a = []
282 self.makecycle(cycle([a]*2), a)
283
284 def test_ifilter(self):
285 a = []
286 self.makecycle(ifilter(lambda x:True, [a]*2), a)
287
288 def test_ifilterfalse(self):
289 a = []
290 self.makecycle(ifilterfalse(lambda x:False, a), a)
291
292 def test_izip(self):
293 a = []
294 self.makecycle(izip([a]*2, [a]*3), a)
295
296 def test_imap(self):
297 a = []
298 self.makecycle(imap(lambda x:x, [a]*2), a)
299
300 def test_islice(self):
301 a = []
302 self.makecycle(islice([a]*2, None), a)
303
304 def test_starmap(self):
305 a = []
306 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
307
308 def test_tee(self):
309 a = []
310 p, q = t = tee([a]*2)
311 a += [a, p, q, t]
312 p.next()
313 del a, p, q, t
314
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000315def R(seqn):
316 'Regular generator'
317 for i in seqn:
318 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000319
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000320class G:
321 'Sequence using __getitem__'
322 def __init__(self, seqn):
323 self.seqn = seqn
324 def __getitem__(self, i):
325 return self.seqn[i]
326
327class I:
328 'Sequence using iterator protocol'
329 def __init__(self, seqn):
330 self.seqn = seqn
331 self.i = 0
332 def __iter__(self):
333 return self
334 def next(self):
335 if self.i >= len(self.seqn): raise StopIteration
336 v = self.seqn[self.i]
337 self.i += 1
338 return v
339
340class Ig:
341 'Sequence using iterator protocol defined with a generator'
342 def __init__(self, seqn):
343 self.seqn = seqn
344 self.i = 0
345 def __iter__(self):
346 for val in self.seqn:
347 yield val
348
349class X:
350 'Missing __getitem__ and __iter__'
351 def __init__(self, seqn):
352 self.seqn = seqn
353 self.i = 0
354 def next(self):
355 if self.i >= len(self.seqn): raise StopIteration
356 v = self.seqn[self.i]
357 self.i += 1
358 return v
359
360class N:
361 'Iterator missing next()'
362 def __init__(self, seqn):
363 self.seqn = seqn
364 self.i = 0
365 def __iter__(self):
366 return self
367
368class E:
369 'Test propagation of exceptions'
370 def __init__(self, seqn):
371 self.seqn = seqn
372 self.i = 0
373 def __iter__(self):
374 return self
375 def next(self):
376 3/0
377
378class S:
379 'Test immediate stop'
380 def __init__(self, seqn):
381 pass
382 def __iter__(self):
383 return self
384 def next(self):
385 raise StopIteration
386
387def L(seqn):
388 'Test multiple tiers of iterators'
389 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
390
391
392class TestVariousIteratorArgs(unittest.TestCase):
393
394 def test_chain(self):
395 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
396 for g in (G, I, Ig, S, L, R):
397 self.assertEqual(list(chain(g(s))), list(g(s)))
398 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
399 self.assertRaises(TypeError, chain, X(s))
400 self.assertRaises(TypeError, list, chain(N(s)))
401 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
402
403 def test_cycle(self):
404 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
405 for g in (G, I, Ig, S, L, R):
406 tgtlen = len(s) * 3
407 expected = list(g(s))*3
408 actual = list(islice(cycle(g(s)), tgtlen))
409 self.assertEqual(actual, expected)
410 self.assertRaises(TypeError, cycle, X(s))
411 self.assertRaises(TypeError, list, cycle(N(s)))
412 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
413
414 def test_ifilter(self):
415 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
416 for g in (G, I, Ig, S, L, R):
417 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
418 self.assertRaises(TypeError, ifilter, isEven, X(s))
419 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
420 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
421
422 def test_ifilterfalse(self):
423 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
424 for g in (G, I, Ig, S, L, R):
425 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
426 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
427 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
428 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
429
430 def test_izip(self):
431 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
432 for g in (G, I, Ig, S, L, R):
433 self.assertEqual(list(izip(g(s))), zip(g(s)))
434 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
435 self.assertRaises(TypeError, izip, X(s))
436 self.assertRaises(TypeError, list, izip(N(s)))
437 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
438
439 def test_imap(self):
440 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
441 for g in (G, I, Ig, S, L, R):
442 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
443 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
444 self.assertRaises(TypeError, imap, onearg, X(s))
445 self.assertRaises(TypeError, list, imap(onearg, N(s)))
446 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
447
448 def test_islice(self):
449 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
450 for g in (G, I, Ig, S, L, R):
451 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
452 self.assertRaises(TypeError, islice, X(s), 10)
453 self.assertRaises(TypeError, list, islice(N(s), 10))
454 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
455
456 def test_starmap(self):
457 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
458 for g in (G, I, Ig, S, L, R):
459 ss = zip(s, s)
460 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
461 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
462 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
463 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
464
465 def test_takewhile(self):
466 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
467 for g in (G, I, Ig, S, L, R):
468 tgt = []
469 for elem in g(s):
470 if not isEven(elem): break
471 tgt.append(elem)
472 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
473 self.assertRaises(TypeError, takewhile, isEven, X(s))
474 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
475 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
476
477 def test_dropwhile(self):
478 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
479 for g in (G, I, Ig, S, L, R):
480 tgt = []
481 for elem in g(s):
482 if not tgt and isOdd(elem): continue
483 tgt.append(elem)
484 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
485 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
486 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
487 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
488
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000489 def test_tee(self):
490 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
491 for g in (G, I, Ig, S, L, R):
492 it1, it2 = tee(g(s))
493 self.assertEqual(list(it1), list(g(s)))
494 self.assertEqual(list(it2), list(g(s)))
495 self.assertRaises(TypeError, tee, X(s))
496 self.assertRaises(TypeError, list, tee(N(s))[0])
497 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
498
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000499class RegressionTests(unittest.TestCase):
500
501 def test_sf_793826(self):
502 # Fix Armin Rigo's successful efforts to wreak havoc
503
504 def mutatingtuple(tuple1, f, tuple2):
505 # this builds a tuple t which is a copy of tuple1,
506 # then calls f(t), then mutates t to be equal to tuple2
507 # (needs len(tuple1) == len(tuple2)).
508 def g(value, first=[1]):
509 if first:
510 del first[:]
511 f(z.next())
512 return value
513 items = list(tuple2)
514 items[1:1] = list(tuple1)
515 gen = imap(g, items)
516 z = izip(*[gen]*len(tuple1))
517 z.next()
518
519 def f(t):
520 global T
521 T = t
522 first[:] = list(T)
523
524 first = []
525 mutatingtuple((1,2,3), f, (4,5,6))
526 second = list(T)
527 self.assertEqual(first, second)
528
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000529
530libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000531
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000532
533>>> amounts = [120.15, 764.05, 823.14]
534>>> for checknum, amount in izip(count(1200), amounts):
535... print 'Check %d is for $%.2f' % (checknum, amount)
536...
537Check 1200 is for $120.15
538Check 1201 is for $764.05
539Check 1202 is for $823.14
540
541>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000542>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
543... print cube
544...
5451
5468
54727
548
549>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +0000550>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000551... print name.title()
552...
553Alex
554Laura
555Martin
556Walter
557Samuele
558
Raymond Hettingera098b332003-09-08 23:58:40 +0000559>>> def take(n, seq):
560... return list(islice(seq, n))
561
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000562>>> def enumerate(iterable):
563... return izip(count(), iterable)
564
565>>> def tabulate(function):
566... "Return function(0), function(1), ..."
567... return imap(function, count())
568
569>>> def iteritems(mapping):
570... return izip(mapping.iterkeys(), mapping.itervalues())
571
572>>> def nth(iterable, n):
573... "Returns the nth item"
Raymond Hettinger60eca932003-02-09 06:40:58 +0000574... return list(islice(iterable, n, n+1))
575
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000576>>> def all(seq, pred=bool):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000577... "Returns True if pred(x) is True for every element in the iterable"
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000578... return False not in imap(pred, seq)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000579
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000580>>> def any(seq, pred=bool):
Raymond Hettinger02420702003-06-29 20:36:23 +0000581... "Returns True if pred(x) is True for at least one element in the iterable"
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000582... return True in imap(pred, seq)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000583
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000584>>> def no(seq, pred=bool):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000585... "Returns True if pred(x) is False for every element in the iterable"
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000586... return True not in imap(pred, seq)
587
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000588>>> def quantify(seq, pred=bool):
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000589... "Count how many times the predicate is True in the sequence"
590... return sum(imap(pred, seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000591
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000592>>> def padnone(seq):
593... "Returns the sequence elements and then returns None indefinitely"
594... return chain(seq, repeat(None))
595
596>>> def ncycles(seq, n):
597... "Returns the sequence elements n times"
598... return chain(*repeat(seq, n))
599
600>>> def dotproduct(vec1, vec2):
601... return sum(imap(operator.mul, vec1, vec2))
602
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000603>>> def flatten(listOfLists):
604... return list(chain(*listOfLists))
605
606>>> def repeatfunc(func, times=None, *args):
607... "Repeat calls to func with specified arguments."
608... " Example: repeatfunc(random.random)"
609... if times is None:
610... return starmap(func, repeat(args))
611... else:
612... return starmap(func, repeat(args, times))
613
Raymond Hettingerbefa37d2003-06-18 19:25:37 +0000614>>> def window(seq, n=2):
615... "Returns a sliding window (of width n) over data from the iterable"
616... " s -> (s0,s1,...s[n-1]), (s1,s2,...,sn), ... "
617... it = iter(seq)
618... result = tuple(islice(it, n))
619... if len(result) == n:
620... yield result
621... for elem in it:
622... result = result[1:] + (elem,)
623... yield result
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000624
625This is not part of the examples but it tests to make sure the definitions
626perform as purported.
627
Raymond Hettingera098b332003-09-08 23:58:40 +0000628>>> take(10, count())
629[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
630
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000631>>> list(enumerate('abc'))
632[(0, 'a'), (1, 'b'), (2, 'c')]
633
634>>> list(islice(tabulate(lambda x: 2*x), 4))
635[0, 2, 4, 6]
636
637>>> nth('abcde', 3)
638['d']
639
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000640>>> all([2, 4, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000641True
642
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000643>>> all([2, 3, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000644False
645
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000646>>> any([2, 4, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000647True
648
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000649>>> any([1, 3, 5, 9], lambda x: x%2==0,)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000650False
651
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000652>>> no([1, 3, 5, 9], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000653True
654
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000655>>> no([1, 2, 5, 9], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000656False
657
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000658>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +000065950
660
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000661>>> a = [[1, 2, 3], [4, 5, 6]]
662>>> flatten(a)
663[1, 2, 3, 4, 5, 6]
664
665>>> list(repeatfunc(pow, 5, 2, 3))
666[8, 8, 8, 8, 8]
667
668>>> import random
669>>> take(5, imap(int, repeatfunc(random.random)))
670[0, 0, 0, 0, 0]
671
Raymond Hettingerbefa37d2003-06-18 19:25:37 +0000672>>> list(window('abc'))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000673[('a', 'b'), ('b', 'c')]
674
Raymond Hettingerbefa37d2003-06-18 19:25:37 +0000675>>> list(window('abc',5))
676[]
677
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000678>>> list(islice(padnone('abc'), 0, 6))
679['a', 'b', 'c', None, None, None]
680
681>>> list(ncycles('abc', 3))
682['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
683
684>>> dotproduct([1,2,3], [4,5,6])
68532
686
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000687"""
688
689__test__ = {'libreftest' : libreftest}
690
691def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000692 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
693 RegressionTests)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000694 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000695
696 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000697 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +0000698 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000699 counts = [None] * 5
700 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000701 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +0000702 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000703 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000704 print counts
705
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000706 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +0000707 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000708
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000709if __name__ == "__main__":
710 test_main(verbose=True)