blob: 8f8fc73d93abf272d3c7ef70a7e89bbec86c46e0 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001import unittest
2from test import test_support
3from itertools import *
Raymond Hettingera9f60922004-10-17 16:40:14 +00004from weakref import proxy
Raymond Hettinger2012f172003-02-07 05:32:58 +00005import sys
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00006import operator
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00007import random
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00008
Guido van Rossum801f0d72006-08-24 19:48:10 +00009def lzip(*args):
10 return list(zip(*args))
11
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000012def onearg(x):
13 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000014 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000015
16def errfunc(*args):
17 'Test function that raises an error'
18 raise ValueError
19
20def gen3():
21 'Non-restartable source sequence'
22 for i in (0, 1, 2):
23 yield i
24
25def isEven(x):
26 'Test predicate'
27 return x%2==0
28
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000029def isOdd(x):
30 'Test predicate'
31 return x%2==1
32
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000033class StopNow:
34 'Class emulating an empty iterable.'
35 def __iter__(self):
36 return self
37 def next(self):
38 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000039
Raymond Hettinger02420702003-06-29 20:36:23 +000040def take(n, seq):
41 'Convenience function for partially consuming a long of infinite iterable'
42 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000043
Raymond Hettinger96ef8112003-02-01 00:10:11 +000044class TestBasicOps(unittest.TestCase):
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000045 def test_chain(self):
46 self.assertEqual(list(chain('abc', 'def')), list('abcdef'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000047 self.assertEqual(list(chain('abc')), list('abc'))
48 self.assertEqual(list(chain('')), [])
Raymond Hettinger02420702003-06-29 20:36:23 +000049 self.assertEqual(take(4, chain('abc', 'def')), list('abcd'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000050 self.assertRaises(TypeError, chain, 2, 3)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000051
Raymond Hettinger96ef8112003-02-01 00:10:11 +000052 def test_count(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +000053 self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
54 self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
55 self.assertEqual(take(2, lzip('abc',count(3))), [('a', 3), ('b', 4)])
Raymond Hettinger96ef8112003-02-01 00:10:11 +000056 self.assertRaises(TypeError, count, 2, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000057 self.assertRaises(TypeError, count, 'a')
58 c = count(sys.maxint-2) # verify that rollover doesn't crash
59 c.next(); c.next(); c.next(); c.next(); c.next()
Raymond Hettinger4cda01e2004-09-28 04:45:28 +000060 c = count(3)
61 self.assertEqual(repr(c), 'count(3)')
62 c.next()
63 self.assertEqual(repr(c), 'count(4)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +000064
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000065 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +000066 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000067 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000068 self.assertRaises(TypeError, cycle)
69 self.assertRaises(TypeError, cycle, 5)
70 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000071
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000072 def test_groupby(self):
73 # Check whether it accepts arguments correctly
74 self.assertEqual([], list(groupby([])))
75 self.assertEqual([], list(groupby([], key=id)))
76 self.assertRaises(TypeError, list, groupby('abc', []))
77 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +000078 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000079
80 # Check normal input
81 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
82 (2,15,22), (3,16,23), (3,17,23)]
83 dup = []
84 for k, g in groupby(s, lambda r:r[0]):
85 for elem in g:
86 self.assertEqual(k, elem[0])
87 dup.append(elem)
88 self.assertEqual(s, dup)
89
90 # Check nested case
91 dup = []
92 for k, g in groupby(s, lambda r:r[0]):
93 for ik, ig in groupby(g, lambda r:r[2]):
94 for elem in ig:
95 self.assertEqual(k, elem[0])
96 self.assertEqual(ik, elem[2])
97 dup.append(elem)
98 self.assertEqual(s, dup)
99
100 # Check case where inner iterator is not used
101 keys = [k for k, g in groupby(s, lambda r:r[0])]
102 expectedkeys = set([r[0] for r in s])
103 self.assertEqual(set(keys), expectedkeys)
104 self.assertEqual(len(keys), len(expectedkeys))
105
106 # Exercise pipes and filters style
107 s = 'abracadabra'
108 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000109 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000110 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
111 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000112 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000113 self.assertEqual(r, ['a', 'b', 'r'])
114 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000115 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000116 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
117 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000118 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000119 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
120
121 # iter.next failure
122 class ExpectedError(Exception):
123 pass
124 def delayed_raise(n=0):
125 for i in range(n):
126 yield 'yo'
127 raise ExpectedError
128 def gulp(iterable, keyp=None, func=list):
129 return [func(g) for k, g in groupby(iterable, keyp)]
130
131 # iter.next failure on outer object
132 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
133 # iter.next failure on inner object
134 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
135
136 # __cmp__ failure
137 class DummyCmp:
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000138 def __eq__(self, dst):
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000139 raise ExpectedError
140 s = [DummyCmp(), DummyCmp(), None]
141
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000142 # __eq__ failure on outer object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000143 self.assertRaises(ExpectedError, gulp, s, func=id)
Guido van Rossum47b9ff62006-08-24 00:41:19 +0000144 # __eq__ failure on inner object
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000145 self.assertRaises(ExpectedError, gulp, s)
146
147 # keyfunc failure
148 def keyfunc(obj):
149 if keyfunc.skip > 0:
150 keyfunc.skip -= 1
151 return obj
152 else:
153 raise ExpectedError
154
155 # keyfunc failure on outer object
156 keyfunc.skip = 0
157 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
158 keyfunc.skip = 1
159 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
160
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000161 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000162 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000163 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000164 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000165 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000166 self.assertRaises(TypeError, ifilter, lambda x:x)
167 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000168 self.assertRaises(TypeError, ifilter, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000169 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000170
171 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000172 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
173 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000174 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000175 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000176 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
177 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000178 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000179 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000180
181 def test_izip(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000182 # XXX This is rather silly now that builtin zip() calls izip()...
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000183 ans = [(x,y) for x, y in izip('abc',count())]
184 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Guido van Rossum801f0d72006-08-24 19:48:10 +0000185 self.assertEqual(list(izip('abc', range(6))), lzip('abc', range(6)))
186 self.assertEqual(list(izip('abcdef', range(3))), lzip('abcdef', range(3)))
187 self.assertEqual(take(3,izip('abcdef', count())), lzip('abcdef', range(3)))
188 self.assertEqual(list(izip('abcdef')), lzip('abcdef'))
189 self.assertEqual(list(izip()), lzip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000190 self.assertRaises(TypeError, izip, 3)
191 self.assertRaises(TypeError, izip, range(3), 3)
192 # Check tuple re-use (implementation detail)
193 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000194 lzip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000195 self.assertEqual([pair for pair in izip('abc', 'def')],
Guido van Rossum801f0d72006-08-24 19:48:10 +0000196 lzip('abc', 'def'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000197 ids = map(id, izip('abc', 'def'))
198 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000199 ids = map(id, list(izip('abc', 'def')))
200 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000201
202 def test_repeat(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000203 self.assertEqual(lzip(xrange(3),repeat('a')),
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000204 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000205 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000206 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000207 self.assertEqual(list(repeat('a', 0)), [])
208 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000209 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000210 self.assertRaises(TypeError, repeat, None, 3, 4)
211 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000212 r = repeat(1+0j)
213 self.assertEqual(repr(r), 'repeat((1+0j))')
214 r = repeat(1+0j, 5)
215 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
216 list(r)
217 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000218
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000219 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000220 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
221 [0**1, 1**2, 2**3])
222 self.assertEqual(list(imap(None, 'abc', range(5))),
223 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000224 self.assertEqual(list(imap(None, 'abc', count())),
225 [('a',0),('b',1),('c',2)])
226 self.assertEqual(take(2,imap(None, 'abc', count())),
227 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000228 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000229 self.assertRaises(TypeError, imap)
230 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000231 self.assertRaises(TypeError, imap(10, range(5)).next)
232 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
233 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000234
235 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000236 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
237 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000238 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
239 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000240 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger2012f172003-02-07 05:32:58 +0000241 self.assertRaises(TypeError, list, starmap(operator.pow, [[4,5]]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000242 self.assertRaises(TypeError, starmap)
243 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
244 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
245 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
246 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000247
248 def test_islice(self):
249 for args in [ # islice(args) should agree with range(args)
250 (10, 20, 3),
251 (10, 3, 20),
252 (10, 20),
253 (10, 3),
254 (20,)
255 ]:
256 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
257
258 for args, tgtargs in [ # Stop when seqn is exhausted
259 ((10, 110, 3), ((10, 100, 3))),
260 ((10, 110), ((10, 100))),
261 ((110,), (100,))
262 ]:
263 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
264
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000265 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000266 self.assertEqual(list(islice(xrange(10), None)), range(10))
Raymond Hettingerb2594052004-12-05 09:25:51 +0000267 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
268 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000269 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
270 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
271
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000272 # Test number of items consumed SF #1171417
273 it = iter(range(10))
274 self.assertEqual(list(islice(it, 3)), range(3))
275 self.assertEqual(list(it), range(3, 10))
276
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000277 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000278 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000279 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
280 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
281 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
282 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
283 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000284 self.assertRaises(ValueError, islice, xrange(10), 'a')
285 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
286 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
287 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
288 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Raymond Hettinger2012f172003-02-07 05:32:58 +0000289 self.assertEqual(len(list(islice(count(), 1, 10, sys.maxint))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000290
291 def test_takewhile(self):
292 data = [1, 3, 5, 20, 2, 4, 6, 8]
293 underten = lambda x: x<10
294 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000295 self.assertEqual(list(takewhile(underten, [])), [])
296 self.assertRaises(TypeError, takewhile)
297 self.assertRaises(TypeError, takewhile, operator.pow)
298 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
299 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
300 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000301 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
302 self.assertEqual(list(t), [1, 1, 1])
303 self.assertRaises(StopIteration, t.next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000304
305 def test_dropwhile(self):
306 data = [1, 3, 5, 20, 2, 4, 6, 8]
307 underten = lambda x: x<10
308 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000309 self.assertEqual(list(dropwhile(underten, [])), [])
310 self.assertRaises(TypeError, dropwhile)
311 self.assertRaises(TypeError, dropwhile, operator.pow)
312 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
313 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
314 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000315
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000316 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000317 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000318 def irange(n):
319 for i in xrange(n):
320 yield i
321
322 a, b = tee([]) # test empty iterator
323 self.assertEqual(list(a), [])
324 self.assertEqual(list(b), [])
325
326 a, b = tee(irange(n)) # test 100% interleaved
Guido van Rossum801f0d72006-08-24 19:48:10 +0000327 self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000328
329 a, b = tee(irange(n)) # test 0% interleaved
330 self.assertEqual(list(a), range(n))
331 self.assertEqual(list(b), range(n))
332
333 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000334 for i in xrange(100):
335 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000336 del a
337 self.assertEqual(list(b), range(n))
338
339 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000340 for i in xrange(100):
341 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000342 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000343 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000344
345 for j in xrange(5): # test randomly interleaved
346 order = [0]*n + [1]*n
347 random.shuffle(order)
348 lists = ([], [])
349 its = tee(irange(n))
350 for i in order:
351 value = its[i].next()
352 lists[i].append(value)
353 self.assertEqual(lists[0], range(n))
354 self.assertEqual(lists[1], range(n))
355
Raymond Hettingerad983e72003-11-12 14:32:26 +0000356 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000357 self.assertRaises(TypeError, tee)
358 self.assertRaises(TypeError, tee, 3)
359 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000360 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000361
Raymond Hettingerad983e72003-11-12 14:32:26 +0000362 # tee object should be instantiable
363 a, b = tee('abc')
364 c = type(a)('def')
365 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000366
Raymond Hettingerad983e72003-11-12 14:32:26 +0000367 # test long-lagged and multi-way split
368 a, b, c = tee(xrange(2000), 3)
369 for i in xrange(100):
370 self.assertEqual(a.next(), i)
371 self.assertEqual(list(b), range(2000))
372 self.assertEqual([c.next(), c.next()], range(2))
373 self.assertEqual(list(a), range(100,2000))
374 self.assertEqual(list(c), range(2,2000))
375
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000376 # test values of n
377 self.assertRaises(TypeError, tee, 'abc', 'invalid')
378 for n in xrange(5):
379 result = tee('abc', n)
380 self.assertEqual(type(result), tuple)
381 self.assertEqual(len(result), n)
382 self.assertEqual(map(list, result), [list('abc')]*n)
383
Raymond Hettingerad983e72003-11-12 14:32:26 +0000384 # tee pass-through to copyable iterator
385 a, b = tee('abc')
386 c, d = tee(a)
387 self.assert_(a is c)
388
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000389 # test tee_new
390 t1, t2 = tee('abc')
391 tnew = type(t1)
392 self.assertRaises(TypeError, tnew)
393 self.assertRaises(TypeError, tnew, 10)
394 t3 = tnew(t1)
395 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000396
Raymond Hettingera9f60922004-10-17 16:40:14 +0000397 # test that tee objects are weak referencable
398 a, b = tee(xrange(10))
399 p = proxy(a)
400 self.assertEqual(getattr(p, '__class__'), type(b))
401 del a
402 self.assertRaises(ReferenceError, getattr, p, '__class__')
403
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000404 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000405 self.assertRaises(StopIteration, izip().next)
406
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000407 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000408 self.assertRaises(StopIteration, f([]).next)
409 self.assertRaises(StopIteration, f(StopNow()).next)
410
411 self.assertRaises(StopIteration, islice([], None).next)
412 self.assertRaises(StopIteration, islice(StopNow(), None).next)
413
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000414 p, q = tee([])
415 self.assertRaises(StopIteration, p.next)
416 self.assertRaises(StopIteration, q.next)
417 p, q = tee(StopNow())
418 self.assertRaises(StopIteration, p.next)
419 self.assertRaises(StopIteration, q.next)
420
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000421 self.assertRaises(StopIteration, repeat(None, 0).next)
422
423 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
424 self.assertRaises(StopIteration, f(lambda x:x, []).next)
425 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
426
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000427class TestGC(unittest.TestCase):
428
429 def makecycle(self, iterator, container):
430 container.append(iterator)
431 iterator.next()
432 del container, iterator
433
434 def test_chain(self):
435 a = []
436 self.makecycle(chain(a), a)
437
438 def test_cycle(self):
439 a = []
440 self.makecycle(cycle([a]*2), a)
441
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000442 def test_dropwhile(self):
443 a = []
444 self.makecycle(dropwhile(bool, [0, a, a]), a)
445
446 def test_groupby(self):
447 a = []
448 self.makecycle(groupby([a]*2, lambda x:x), a)
449
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000450 def test_ifilter(self):
451 a = []
452 self.makecycle(ifilter(lambda x:True, [a]*2), a)
453
454 def test_ifilterfalse(self):
455 a = []
456 self.makecycle(ifilterfalse(lambda x:False, a), a)
457
458 def test_izip(self):
459 a = []
460 self.makecycle(izip([a]*2, [a]*3), a)
461
462 def test_imap(self):
463 a = []
464 self.makecycle(imap(lambda x:x, [a]*2), a)
465
466 def test_islice(self):
467 a = []
468 self.makecycle(islice([a]*2, None), a)
469
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000470 def test_repeat(self):
471 a = []
472 self.makecycle(repeat(a), a)
473
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000474 def test_starmap(self):
475 a = []
476 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
477
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000478 def test_takewhile(self):
479 a = []
480 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
481
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000482def R(seqn):
483 'Regular generator'
484 for i in seqn:
485 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000486
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000487class G:
488 'Sequence using __getitem__'
489 def __init__(self, seqn):
490 self.seqn = seqn
491 def __getitem__(self, i):
492 return self.seqn[i]
493
494class I:
495 'Sequence using iterator protocol'
496 def __init__(self, seqn):
497 self.seqn = seqn
498 self.i = 0
499 def __iter__(self):
500 return self
501 def next(self):
502 if self.i >= len(self.seqn): raise StopIteration
503 v = self.seqn[self.i]
504 self.i += 1
505 return v
506
507class Ig:
508 'Sequence using iterator protocol defined with a generator'
509 def __init__(self, seqn):
510 self.seqn = seqn
511 self.i = 0
512 def __iter__(self):
513 for val in self.seqn:
514 yield val
515
516class X:
517 'Missing __getitem__ and __iter__'
518 def __init__(self, seqn):
519 self.seqn = seqn
520 self.i = 0
521 def next(self):
522 if self.i >= len(self.seqn): raise StopIteration
523 v = self.seqn[self.i]
524 self.i += 1
525 return v
526
527class N:
528 'Iterator missing next()'
529 def __init__(self, seqn):
530 self.seqn = seqn
531 self.i = 0
532 def __iter__(self):
533 return self
534
535class E:
536 'Test propagation of exceptions'
537 def __init__(self, seqn):
538 self.seqn = seqn
539 self.i = 0
540 def __iter__(self):
541 return self
542 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +0000543 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000544
545class S:
546 'Test immediate stop'
547 def __init__(self, seqn):
548 pass
549 def __iter__(self):
550 return self
551 def next(self):
552 raise StopIteration
553
554def L(seqn):
555 'Test multiple tiers of iterators'
556 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
557
558
559class TestVariousIteratorArgs(unittest.TestCase):
560
561 def test_chain(self):
562 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
563 for g in (G, I, Ig, S, L, R):
564 self.assertEqual(list(chain(g(s))), list(g(s)))
565 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
566 self.assertRaises(TypeError, chain, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000567 self.assertRaises(TypeError, chain, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000568 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
569
570 def test_cycle(self):
571 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
572 for g in (G, I, Ig, S, L, R):
573 tgtlen = len(s) * 3
574 expected = list(g(s))*3
575 actual = list(islice(cycle(g(s)), tgtlen))
576 self.assertEqual(actual, expected)
577 self.assertRaises(TypeError, cycle, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000578 self.assertRaises(TypeError, cycle, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000579 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
580
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000581 def test_groupby(self):
582 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
583 for g in (G, I, Ig, S, L, R):
584 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
585 self.assertRaises(TypeError, groupby, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000586 self.assertRaises(TypeError, groupby, N(s))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000587 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
588
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000589 def test_ifilter(self):
590 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
591 for g in (G, I, Ig, S, L, R):
592 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
593 self.assertRaises(TypeError, ifilter, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000594 self.assertRaises(TypeError, ifilter, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000595 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
596
597 def test_ifilterfalse(self):
598 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
599 for g in (G, I, Ig, S, L, R):
600 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
601 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000602 self.assertRaises(TypeError, ifilterfalse, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000603 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
604
605 def test_izip(self):
606 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
607 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000608 self.assertEqual(list(izip(g(s))), lzip(g(s)))
609 self.assertEqual(list(izip(g(s), g(s))), lzip(g(s), g(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000610 self.assertRaises(TypeError, izip, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000611 self.assertRaises(TypeError, izip, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000612 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
613
614 def test_imap(self):
615 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
616 for g in (G, I, Ig, S, L, R):
617 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
618 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
619 self.assertRaises(TypeError, imap, onearg, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000620 self.assertRaises(TypeError, imap, onearg, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000621 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
622
623 def test_islice(self):
624 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
625 for g in (G, I, Ig, S, L, R):
626 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
627 self.assertRaises(TypeError, islice, X(s), 10)
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000628 self.assertRaises(TypeError, islice, N(s), 10)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000629 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
630
631 def test_starmap(self):
632 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
633 for g in (G, I, Ig, S, L, R):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000634 ss = lzip(s, s)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000635 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
636 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000637 self.assertRaises(TypeError, starmap, operator.pow, N(ss))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000638 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
639
640 def test_takewhile(self):
641 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
642 for g in (G, I, Ig, S, L, R):
643 tgt = []
644 for elem in g(s):
645 if not isEven(elem): break
646 tgt.append(elem)
647 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
648 self.assertRaises(TypeError, takewhile, isEven, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000649 self.assertRaises(TypeError, takewhile, isEven, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000650 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
651
652 def test_dropwhile(self):
653 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
654 for g in (G, I, Ig, S, L, R):
655 tgt = []
656 for elem in g(s):
657 if not tgt and isOdd(elem): continue
658 tgt.append(elem)
659 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
660 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000661 self.assertRaises(TypeError, dropwhile, isOdd, N(s))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000662 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
663
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000664 def test_tee(self):
665 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
666 for g in (G, I, Ig, S, L, R):
667 it1, it2 = tee(g(s))
668 self.assertEqual(list(it1), list(g(s)))
669 self.assertEqual(list(it2), list(g(s)))
670 self.assertRaises(TypeError, tee, X(s))
Thomas Wouters8690c4e2006-04-15 09:07:20 +0000671 self.assertRaises(TypeError, tee, N(s))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000672 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
673
Raymond Hettinger5cab2e32004-02-10 09:25:40 +0000674class LengthTransparency(unittest.TestCase):
675
676 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +0000677 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +0000678 self.assertEqual(len(repeat(None, 50)), 50)
679 self.assertRaises(TypeError, len, repeat(None))
680
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000681class RegressionTests(unittest.TestCase):
682
683 def test_sf_793826(self):
684 # Fix Armin Rigo's successful efforts to wreak havoc
685
686 def mutatingtuple(tuple1, f, tuple2):
687 # this builds a tuple t which is a copy of tuple1,
688 # then calls f(t), then mutates t to be equal to tuple2
689 # (needs len(tuple1) == len(tuple2)).
690 def g(value, first=[1]):
691 if first:
692 del first[:]
693 f(z.next())
694 return value
695 items = list(tuple2)
696 items[1:1] = list(tuple1)
697 gen = imap(g, items)
698 z = izip(*[gen]*len(tuple1))
699 z.next()
700
701 def f(t):
702 global T
703 T = t
704 first[:] = list(T)
705
706 first = []
707 mutatingtuple((1,2,3), f, (4,5,6))
708 second = list(T)
709 self.assertEqual(first, second)
710
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000711
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000712 def test_sf_950057(self):
713 # Make sure that chain() and cycle() catch exceptions immediately
714 # rather than when shifting between input sources
715
716 def gen1():
717 hist.append(0)
718 yield 1
719 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +0000720 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000721 hist.append(2)
722
723 def gen2(x):
724 hist.append(3)
725 yield 2
726 hist.append(4)
727 if x:
728 raise StopIteration
729
730 hist = []
731 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
732 self.assertEqual(hist, [0,1])
733
734 hist = []
735 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
736 self.assertEqual(hist, [0,1])
737
738 hist = []
739 self.assertRaises(AssertionError, list, cycle(gen1()))
740 self.assertEqual(hist, [0,1])
741
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000742libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000743
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000744
745>>> amounts = [120.15, 764.05, 823.14]
746>>> for checknum, amount in izip(count(1200), amounts):
747... print 'Check %d is for $%.2f' % (checknum, amount)
748...
749Check 1200 is for $120.15
750Check 1201 is for $764.05
751Check 1202 is for $823.14
752
753>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000754>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
755... print cube
756...
7571
7588
75927
760
761>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +0000762>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000763... print name.title()
764...
765Alex
766Laura
767Martin
768Walter
769Samuele
770
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000771>>> from operator import itemgetter
772>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Thomas Wouters4d70c3d2006-06-08 14:42:34 +0000773>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000774>>> for k, g in groupby(di, itemgetter(1)):
775... print k, map(itemgetter(0), g)
776...
7771 ['a', 'c', 'e']
7782 ['b', 'd', 'f']
7793 ['g']
780
Raymond Hettinger734fb572004-01-20 20:04:40 +0000781# Find runs of consecutive numbers using groupby. The key to the solution
782# is differencing with a range so that consecutive numbers all appear in
783# same group.
784>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
785>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
786... print map(operator.itemgetter(1), g)
787...
788[1]
789[4, 5, 6]
790[10]
791[15, 16, 17, 18]
792[22]
793[25, 26, 27, 28]
794
Raymond Hettingera098b332003-09-08 23:58:40 +0000795>>> def take(n, seq):
796... return list(islice(seq, n))
797
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000798>>> def enumerate(iterable):
799... return izip(count(), iterable)
800
801>>> def tabulate(function):
802... "Return function(0), function(1), ..."
803... return imap(function, count())
804
805>>> def iteritems(mapping):
806... return izip(mapping.iterkeys(), mapping.itervalues())
807
808>>> def nth(iterable, n):
809... "Returns the nth item"
Raymond Hettinger60eca932003-02-09 06:40:58 +0000810... return list(islice(iterable, n, n+1))
811
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000812>>> def all(seq, pred=None):
813... "Returns True if pred(x) is true for every element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000814... for elem in ifilterfalse(pred, seq):
815... return False
816... return True
Raymond Hettinger60eca932003-02-09 06:40:58 +0000817
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000818>>> def any(seq, pred=None):
819... "Returns True if pred(x) is true for at least one element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000820... for elem in ifilter(pred, seq):
821... return True
822... return False
Raymond Hettinger60eca932003-02-09 06:40:58 +0000823
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000824>>> def no(seq, pred=None):
825... "Returns True if pred(x) is false for every element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000826... for elem in ifilter(pred, seq):
827... return False
828... return True
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000829
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000830>>> def quantify(seq, pred=None):
831... "Count how many times the predicate is true in the sequence"
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000832... return sum(imap(pred, seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000833
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000834>>> def padnone(seq):
835... "Returns the sequence elements and then returns None indefinitely"
836... return chain(seq, repeat(None))
837
838>>> def ncycles(seq, n):
839... "Returns the sequence elements n times"
840... return chain(*repeat(seq, n))
841
842>>> def dotproduct(vec1, vec2):
843... return sum(imap(operator.mul, vec1, vec2))
844
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000845>>> def flatten(listOfLists):
846... return list(chain(*listOfLists))
847
848>>> def repeatfunc(func, times=None, *args):
849... "Repeat calls to func with specified arguments."
850... " Example: repeatfunc(random.random)"
851... if times is None:
852... return starmap(func, repeat(args))
853... else:
854... return starmap(func, repeat(args, times))
855
Raymond Hettingerd591f662003-10-26 15:34:50 +0000856>>> def pairwise(iterable):
857... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
858... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000859... try:
860... b.next()
861... except StopIteration:
862... pass
863... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000864
865This is not part of the examples but it tests to make sure the definitions
866perform as purported.
867
Raymond Hettingera098b332003-09-08 23:58:40 +0000868>>> take(10, count())
869[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
870
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000871>>> list(enumerate('abc'))
872[(0, 'a'), (1, 'b'), (2, 'c')]
873
874>>> list(islice(tabulate(lambda x: 2*x), 4))
875[0, 2, 4, 6]
876
877>>> nth('abcde', 3)
878['d']
879
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000880>>> all([2, 4, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000881True
882
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000883>>> all([2, 3, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000884False
885
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000886>>> any([2, 4, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000887True
888
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000889>>> any([1, 3, 5, 9], lambda x: x%2==0,)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000890False
891
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000892>>> no([1, 3, 5, 9], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000893True
894
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000895>>> no([1, 2, 5, 9], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000896False
897
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000898>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +000089950
900
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000901>>> a = [[1, 2, 3], [4, 5, 6]]
902>>> flatten(a)
903[1, 2, 3, 4, 5, 6]
904
905>>> list(repeatfunc(pow, 5, 2, 3))
906[8, 8, 8, 8, 8]
907
908>>> import random
909>>> take(5, imap(int, repeatfunc(random.random)))
910[0, 0, 0, 0, 0]
911
Raymond Hettingerd591f662003-10-26 15:34:50 +0000912>>> list(pairwise('abcd'))
913[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000914
Raymond Hettingerd591f662003-10-26 15:34:50 +0000915>>> list(pairwise([]))
916[]
917
918>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +0000919[]
920
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000921>>> list(islice(padnone('abc'), 0, 6))
922['a', 'b', 'c', None, None, None]
923
924>>> list(ncycles('abc', 3))
925['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
926
927>>> dotproduct([1,2,3], [4,5,6])
92832
929
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000930"""
931
932__test__ = {'libreftest' : libreftest}
933
934def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000935 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Raymond Hettinger5cab2e32004-02-10 09:25:40 +0000936 RegressionTests, LengthTransparency)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000937 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000938
939 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000940 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +0000941 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000942 counts = [None] * 5
943 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000944 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +0000945 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000946 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000947 print counts
948
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000949 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +0000950 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000951
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000952if __name__ == "__main__":
953 test_main(verbose=True)