blob: a8f625c41e8da32e48e130d9c1379669619dfd04 [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
9def onearg(x):
10 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000011 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000012
13def errfunc(*args):
14 'Test function that raises an error'
15 raise ValueError
16
17def gen3():
18 'Non-restartable source sequence'
19 for i in (0, 1, 2):
20 yield i
21
22def isEven(x):
23 'Test predicate'
24 return x%2==0
25
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000026def isOdd(x):
27 'Test predicate'
28 return x%2==1
29
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000030class StopNow:
31 'Class emulating an empty iterable.'
32 def __iter__(self):
33 return self
34 def next(self):
35 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000036
Raymond Hettinger02420702003-06-29 20:36:23 +000037def take(n, seq):
38 'Convenience function for partially consuming a long of infinite iterable'
39 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000040
Raymond Hettinger96ef8112003-02-01 00:10:11 +000041class TestBasicOps(unittest.TestCase):
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000042 def test_chain(self):
43 self.assertEqual(list(chain('abc', 'def')), list('abcdef'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000044 self.assertEqual(list(chain('abc')), list('abc'))
45 self.assertEqual(list(chain('')), [])
Raymond Hettinger02420702003-06-29 20:36:23 +000046 self.assertEqual(take(4, chain('abc', 'def')), list('abcd'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000047 self.assertRaises(TypeError, chain, 2, 3)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000048
Raymond Hettinger96ef8112003-02-01 00:10:11 +000049 def test_count(self):
50 self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
51 self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
Raymond Hettinger02420702003-06-29 20:36:23 +000052 self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
Raymond Hettinger96ef8112003-02-01 00:10:11 +000053 self.assertRaises(TypeError, count, 2, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000054 self.assertRaises(TypeError, count, 'a')
Raymond Hettinger6d121f12007-02-08 00:07:32 +000055 self.assertRaises(OverflowError, list, islice(count(sys.maxint-5), 10))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +000056 c = count(3)
57 self.assertEqual(repr(c), 'count(3)')
58 c.next()
59 self.assertEqual(repr(c), 'count(4)')
Jack Diederich36234e82006-09-21 17:50:26 +000060 c = count(-9)
61 self.assertEqual(repr(c), 'count(-9)')
62 c.next()
63 self.assertEqual(c.next(), -8)
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:
138 def __cmp__(self, dst):
139 raise ExpectedError
140 s = [DummyCmp(), DummyCmp(), None]
141
142 # __cmp__ failure on outer object
143 self.assertRaises(ExpectedError, gulp, s, func=id)
144 # __cmp__ failure on inner object
145 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):
182 ans = [(x,y) for x, y in izip('abc',count())]
183 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000184 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
185 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
Raymond Hettinger02420702003-06-29 20:36:23 +0000186 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000187 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000188 self.assertEqual(list(izip()), zip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000189 self.assertRaises(TypeError, izip, 3)
190 self.assertRaises(TypeError, izip, range(3), 3)
191 # Check tuple re-use (implementation detail)
192 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
193 zip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000194 self.assertEqual([pair for pair in izip('abc', 'def')],
195 zip('abc', 'def'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000196 ids = map(id, izip('abc', 'def'))
197 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000198 ids = map(id, list(izip('abc', 'def')))
199 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000200
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000201 def test_iziplongest(self):
202 for args in [
203 ['abc', range(6)],
204 [range(6), 'abc'],
205 [range(1000), range(2000,2100), range(3000,3050)],
206 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
207 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
208 ]:
209 target = map(None, *args)
210 self.assertEqual(list(izip_longest(*args)), target)
211 self.assertEqual(list(izip_longest(*args, **{})), target)
212 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
213 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Tim Petersea5962f2007-03-12 18:07:52 +0000214
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000215 self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
216
217 self.assertEqual(list(izip_longest()), zip())
218 self.assertEqual(list(izip_longest([])), zip([]))
219 self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
Tim Petersea5962f2007-03-12 18:07:52 +0000220
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000221 self.assertEqual(list(izip_longest('abc', 'defg', **{})), map(None, 'abc', 'defg')) # empty keyword dict
222 self.assertRaises(TypeError, izip_longest, 3)
223 self.assertRaises(TypeError, izip_longest, range(3), 3)
224
225 for stmt in [
226 "izip_longest('abc', fv=1)",
Tim Petersea5962f2007-03-12 18:07:52 +0000227 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000228 ]:
229 try:
230 eval(stmt, globals(), locals())
231 except TypeError:
232 pass
233 else:
234 self.fail('Did not raise Type in: ' + stmt)
Tim Petersea5962f2007-03-12 18:07:52 +0000235
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000236 # Check tuple re-use (implementation detail)
237 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
238 zip('abc', 'def'))
239 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
240 zip('abc', 'def'))
241 ids = map(id, izip_longest('abc', 'def'))
242 self.assertEqual(min(ids), max(ids))
243 ids = map(id, list(izip_longest('abc', 'def')))
244 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
245
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000246 def test_repeat(self):
247 self.assertEqual(zip(xrange(3),repeat('a')),
248 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000249 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000250 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000251 self.assertEqual(list(repeat('a', 0)), [])
252 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000253 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000254 self.assertRaises(TypeError, repeat, None, 3, 4)
255 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000256 r = repeat(1+0j)
257 self.assertEqual(repr(r), 'repeat((1+0j))')
258 r = repeat(1+0j, 5)
259 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
260 list(r)
261 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000262
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000263 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000264 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
265 [0**1, 1**2, 2**3])
266 self.assertEqual(list(imap(None, 'abc', range(5))),
267 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000268 self.assertEqual(list(imap(None, 'abc', count())),
269 [('a',0),('b',1),('c',2)])
270 self.assertEqual(take(2,imap(None, 'abc', count())),
271 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000272 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000273 self.assertRaises(TypeError, imap)
274 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000275 self.assertRaises(TypeError, imap(10, range(5)).next)
276 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
277 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000278
279 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000280 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
281 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000282 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
283 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000284 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger2012f172003-02-07 05:32:58 +0000285 self.assertRaises(TypeError, list, starmap(operator.pow, [[4,5]]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000286 self.assertRaises(TypeError, starmap)
287 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
288 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
289 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
290 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000291
292 def test_islice(self):
293 for args in [ # islice(args) should agree with range(args)
294 (10, 20, 3),
295 (10, 3, 20),
296 (10, 20),
297 (10, 3),
298 (20,)
299 ]:
300 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
301
302 for args, tgtargs in [ # Stop when seqn is exhausted
303 ((10, 110, 3), ((10, 100, 3))),
304 ((10, 110), ((10, 100))),
305 ((110,), (100,))
306 ]:
307 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
308
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000309 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000310 self.assertEqual(list(islice(xrange(10), None)), range(10))
Raymond Hettingerb2594052004-12-05 09:25:51 +0000311 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
312 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000313 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
314 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
315
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000316 # Test number of items consumed SF #1171417
317 it = iter(range(10))
318 self.assertEqual(list(islice(it, 3)), range(3))
319 self.assertEqual(list(it), range(3, 10))
320
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000321 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000322 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000323 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
324 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
325 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
326 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
327 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000328 self.assertRaises(ValueError, islice, xrange(10), 'a')
329 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
330 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
331 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
332 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Raymond Hettinger2012f172003-02-07 05:32:58 +0000333 self.assertEqual(len(list(islice(count(), 1, 10, sys.maxint))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000334
335 def test_takewhile(self):
336 data = [1, 3, 5, 20, 2, 4, 6, 8]
337 underten = lambda x: x<10
338 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000339 self.assertEqual(list(takewhile(underten, [])), [])
340 self.assertRaises(TypeError, takewhile)
341 self.assertRaises(TypeError, takewhile, operator.pow)
342 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
343 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
344 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000345 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
346 self.assertEqual(list(t), [1, 1, 1])
347 self.assertRaises(StopIteration, t.next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000348
349 def test_dropwhile(self):
350 data = [1, 3, 5, 20, 2, 4, 6, 8]
351 underten = lambda x: x<10
352 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000353 self.assertEqual(list(dropwhile(underten, [])), [])
354 self.assertRaises(TypeError, dropwhile)
355 self.assertRaises(TypeError, dropwhile, operator.pow)
356 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
357 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
358 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000359
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000360 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000361 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000362 def irange(n):
363 for i in xrange(n):
364 yield i
365
366 a, b = tee([]) # test empty iterator
367 self.assertEqual(list(a), [])
368 self.assertEqual(list(b), [])
369
370 a, b = tee(irange(n)) # test 100% interleaved
371 self.assertEqual(zip(a,b), zip(range(n),range(n)))
372
373 a, b = tee(irange(n)) # test 0% interleaved
374 self.assertEqual(list(a), range(n))
375 self.assertEqual(list(b), range(n))
376
377 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000378 for i in xrange(100):
379 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000380 del a
381 self.assertEqual(list(b), range(n))
382
383 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000384 for i in xrange(100):
385 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000386 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000387 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000388
389 for j in xrange(5): # test randomly interleaved
390 order = [0]*n + [1]*n
391 random.shuffle(order)
392 lists = ([], [])
393 its = tee(irange(n))
394 for i in order:
395 value = its[i].next()
396 lists[i].append(value)
397 self.assertEqual(lists[0], range(n))
398 self.assertEqual(lists[1], range(n))
399
Raymond Hettingerad983e72003-11-12 14:32:26 +0000400 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000401 self.assertRaises(TypeError, tee)
402 self.assertRaises(TypeError, tee, 3)
403 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000404 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000405
Raymond Hettingerad983e72003-11-12 14:32:26 +0000406 # tee object should be instantiable
407 a, b = tee('abc')
408 c = type(a)('def')
409 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000410
Raymond Hettingerad983e72003-11-12 14:32:26 +0000411 # test long-lagged and multi-way split
412 a, b, c = tee(xrange(2000), 3)
413 for i in xrange(100):
414 self.assertEqual(a.next(), i)
415 self.assertEqual(list(b), range(2000))
416 self.assertEqual([c.next(), c.next()], range(2))
417 self.assertEqual(list(a), range(100,2000))
418 self.assertEqual(list(c), range(2,2000))
419
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000420 # test values of n
421 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Neal Norwitz69e88972006-09-02 02:58:13 +0000422 self.assertRaises(ValueError, tee, [], -1)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000423 for n in xrange(5):
424 result = tee('abc', n)
425 self.assertEqual(type(result), tuple)
426 self.assertEqual(len(result), n)
427 self.assertEqual(map(list, result), [list('abc')]*n)
428
Raymond Hettingerad983e72003-11-12 14:32:26 +0000429 # tee pass-through to copyable iterator
430 a, b = tee('abc')
431 c, d = tee(a)
432 self.assert_(a is c)
433
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000434 # test tee_new
435 t1, t2 = tee('abc')
436 tnew = type(t1)
437 self.assertRaises(TypeError, tnew)
438 self.assertRaises(TypeError, tnew, 10)
439 t3 = tnew(t1)
440 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000441
Raymond Hettingera9f60922004-10-17 16:40:14 +0000442 # test that tee objects are weak referencable
443 a, b = tee(xrange(10))
444 p = proxy(a)
445 self.assertEqual(getattr(p, '__class__'), type(b))
446 del a
447 self.assertRaises(ReferenceError, getattr, p, '__class__')
448
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000449 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000450 self.assertRaises(StopIteration, izip().next)
451
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000452 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000453 self.assertRaises(StopIteration, f([]).next)
454 self.assertRaises(StopIteration, f(StopNow()).next)
455
456 self.assertRaises(StopIteration, islice([], None).next)
457 self.assertRaises(StopIteration, islice(StopNow(), None).next)
458
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000459 p, q = tee([])
460 self.assertRaises(StopIteration, p.next)
461 self.assertRaises(StopIteration, q.next)
462 p, q = tee(StopNow())
463 self.assertRaises(StopIteration, p.next)
464 self.assertRaises(StopIteration, q.next)
465
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000466 self.assertRaises(StopIteration, repeat(None, 0).next)
467
468 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
469 self.assertRaises(StopIteration, f(lambda x:x, []).next)
470 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
471
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000472class TestGC(unittest.TestCase):
473
474 def makecycle(self, iterator, container):
475 container.append(iterator)
476 iterator.next()
477 del container, iterator
478
479 def test_chain(self):
480 a = []
481 self.makecycle(chain(a), a)
482
483 def test_cycle(self):
484 a = []
485 self.makecycle(cycle([a]*2), a)
486
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000487 def test_dropwhile(self):
488 a = []
489 self.makecycle(dropwhile(bool, [0, a, a]), a)
490
491 def test_groupby(self):
492 a = []
493 self.makecycle(groupby([a]*2, lambda x:x), a)
494
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000495 def test_ifilter(self):
496 a = []
497 self.makecycle(ifilter(lambda x:True, [a]*2), a)
498
499 def test_ifilterfalse(self):
500 a = []
501 self.makecycle(ifilterfalse(lambda x:False, a), a)
502
503 def test_izip(self):
504 a = []
505 self.makecycle(izip([a]*2, [a]*3), a)
506
507 def test_imap(self):
508 a = []
509 self.makecycle(imap(lambda x:x, [a]*2), a)
510
511 def test_islice(self):
512 a = []
513 self.makecycle(islice([a]*2, None), a)
514
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000515 def test_repeat(self):
516 a = []
517 self.makecycle(repeat(a), a)
518
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000519 def test_starmap(self):
520 a = []
521 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
522
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000523 def test_takewhile(self):
524 a = []
525 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
526
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000527def R(seqn):
528 'Regular generator'
529 for i in seqn:
530 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000531
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000532class G:
533 'Sequence using __getitem__'
534 def __init__(self, seqn):
535 self.seqn = seqn
536 def __getitem__(self, i):
537 return self.seqn[i]
538
539class I:
540 'Sequence using iterator protocol'
541 def __init__(self, seqn):
542 self.seqn = seqn
543 self.i = 0
544 def __iter__(self):
545 return self
546 def next(self):
547 if self.i >= len(self.seqn): raise StopIteration
548 v = self.seqn[self.i]
549 self.i += 1
550 return v
551
552class Ig:
553 'Sequence using iterator protocol defined with a generator'
554 def __init__(self, seqn):
555 self.seqn = seqn
556 self.i = 0
557 def __iter__(self):
558 for val in self.seqn:
559 yield val
560
561class X:
562 'Missing __getitem__ and __iter__'
563 def __init__(self, seqn):
564 self.seqn = seqn
565 self.i = 0
566 def next(self):
567 if self.i >= len(self.seqn): raise StopIteration
568 v = self.seqn[self.i]
569 self.i += 1
570 return v
571
572class N:
573 'Iterator missing next()'
574 def __init__(self, seqn):
575 self.seqn = seqn
576 self.i = 0
577 def __iter__(self):
578 return self
579
580class E:
581 'Test propagation of exceptions'
582 def __init__(self, seqn):
583 self.seqn = seqn
584 self.i = 0
585 def __iter__(self):
586 return self
587 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +0000588 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000589
590class S:
591 'Test immediate stop'
592 def __init__(self, seqn):
593 pass
594 def __iter__(self):
595 return self
596 def next(self):
597 raise StopIteration
598
599def L(seqn):
600 'Test multiple tiers of iterators'
601 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
602
603
604class TestVariousIteratorArgs(unittest.TestCase):
605
606 def test_chain(self):
607 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
608 for g in (G, I, Ig, S, L, R):
609 self.assertEqual(list(chain(g(s))), list(g(s)))
610 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
611 self.assertRaises(TypeError, chain, X(s))
612 self.assertRaises(TypeError, list, chain(N(s)))
613 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
614
615 def test_cycle(self):
616 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
617 for g in (G, I, Ig, S, L, R):
618 tgtlen = len(s) * 3
619 expected = list(g(s))*3
620 actual = list(islice(cycle(g(s)), tgtlen))
621 self.assertEqual(actual, expected)
622 self.assertRaises(TypeError, cycle, X(s))
623 self.assertRaises(TypeError, list, cycle(N(s)))
624 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
625
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000626 def test_groupby(self):
627 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
628 for g in (G, I, Ig, S, L, R):
629 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
630 self.assertRaises(TypeError, groupby, X(s))
631 self.assertRaises(TypeError, list, groupby(N(s)))
632 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
633
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000634 def test_ifilter(self):
635 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
636 for g in (G, I, Ig, S, L, R):
637 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
638 self.assertRaises(TypeError, ifilter, isEven, X(s))
639 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
640 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
641
642 def test_ifilterfalse(self):
643 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
644 for g in (G, I, Ig, S, L, R):
645 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
646 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
647 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
648 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
649
650 def test_izip(self):
651 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
652 for g in (G, I, Ig, S, L, R):
653 self.assertEqual(list(izip(g(s))), zip(g(s)))
654 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
655 self.assertRaises(TypeError, izip, X(s))
656 self.assertRaises(TypeError, list, izip(N(s)))
657 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
658
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000659 def test_iziplongest(self):
660 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
661 for g in (G, I, Ig, S, L, R):
662 self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
663 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
664 self.assertRaises(TypeError, izip_longest, X(s))
665 self.assertRaises(TypeError, list, izip_longest(N(s)))
666 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
667
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000668 def test_imap(self):
669 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
670 for g in (G, I, Ig, S, L, R):
671 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
672 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
673 self.assertRaises(TypeError, imap, onearg, X(s))
674 self.assertRaises(TypeError, list, imap(onearg, N(s)))
675 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
676
677 def test_islice(self):
678 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
679 for g in (G, I, Ig, S, L, R):
680 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
681 self.assertRaises(TypeError, islice, X(s), 10)
682 self.assertRaises(TypeError, list, islice(N(s), 10))
683 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
684
685 def test_starmap(self):
686 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
687 for g in (G, I, Ig, S, L, R):
688 ss = zip(s, s)
689 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
690 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
691 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
692 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
693
694 def test_takewhile(self):
695 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
696 for g in (G, I, Ig, S, L, R):
697 tgt = []
698 for elem in g(s):
699 if not isEven(elem): break
700 tgt.append(elem)
701 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
702 self.assertRaises(TypeError, takewhile, isEven, X(s))
703 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
704 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
705
706 def test_dropwhile(self):
707 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
708 for g in (G, I, Ig, S, L, R):
709 tgt = []
710 for elem in g(s):
711 if not tgt and isOdd(elem): continue
712 tgt.append(elem)
713 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
714 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
715 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
716 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
717
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000718 def test_tee(self):
719 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
720 for g in (G, I, Ig, S, L, R):
721 it1, it2 = tee(g(s))
722 self.assertEqual(list(it1), list(g(s)))
723 self.assertEqual(list(it2), list(g(s)))
724 self.assertRaises(TypeError, tee, X(s))
725 self.assertRaises(TypeError, list, tee(N(s))[0])
726 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
727
Raymond Hettinger5cab2e32004-02-10 09:25:40 +0000728class LengthTransparency(unittest.TestCase):
729
730 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +0000731 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +0000732 self.assertEqual(len(repeat(None, 50)), 50)
733 self.assertRaises(TypeError, len, repeat(None))
734
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000735class RegressionTests(unittest.TestCase):
736
737 def test_sf_793826(self):
738 # Fix Armin Rigo's successful efforts to wreak havoc
739
740 def mutatingtuple(tuple1, f, tuple2):
741 # this builds a tuple t which is a copy of tuple1,
742 # then calls f(t), then mutates t to be equal to tuple2
743 # (needs len(tuple1) == len(tuple2)).
744 def g(value, first=[1]):
745 if first:
746 del first[:]
747 f(z.next())
748 return value
749 items = list(tuple2)
750 items[1:1] = list(tuple1)
751 gen = imap(g, items)
752 z = izip(*[gen]*len(tuple1))
753 z.next()
754
755 def f(t):
756 global T
757 T = t
758 first[:] = list(T)
759
760 first = []
761 mutatingtuple((1,2,3), f, (4,5,6))
762 second = list(T)
763 self.assertEqual(first, second)
764
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000765
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000766 def test_sf_950057(self):
767 # Make sure that chain() and cycle() catch exceptions immediately
768 # rather than when shifting between input sources
769
770 def gen1():
771 hist.append(0)
772 yield 1
773 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +0000774 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000775 hist.append(2)
776
777 def gen2(x):
778 hist.append(3)
779 yield 2
780 hist.append(4)
781 if x:
782 raise StopIteration
783
784 hist = []
785 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
786 self.assertEqual(hist, [0,1])
787
788 hist = []
789 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
790 self.assertEqual(hist, [0,1])
791
792 hist = []
793 self.assertRaises(AssertionError, list, cycle(gen1()))
794 self.assertEqual(hist, [0,1])
795
Georg Brandlb84c1372007-01-21 10:28:43 +0000796class SubclassWithKwargsTest(unittest.TestCase):
797 def test_keywords_in_subclass(self):
798 # count is not subclassable...
799 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
800 starmap, islice, takewhile, dropwhile, cycle):
801 class Subclass(cls):
802 def __init__(self, newarg=None, *args):
803 cls.__init__(self, *args)
804 try:
805 Subclass(newarg=1)
806 except TypeError, err:
807 # we expect type errors because of wrong argument count
808 self.failIf("does not take keyword arguments" in err.args[0])
809
810
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000811libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000812
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000813
814>>> amounts = [120.15, 764.05, 823.14]
815>>> for checknum, amount in izip(count(1200), amounts):
816... print 'Check %d is for $%.2f' % (checknum, amount)
817...
818Check 1200 is for $120.15
819Check 1201 is for $764.05
820Check 1202 is for $823.14
821
822>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000823>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
824... print cube
825...
8261
8278
82827
829
830>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +0000831>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000832... print name.title()
833...
834Alex
835Laura
836Martin
837Walter
838Samuele
839
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000840>>> from operator import itemgetter
841>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Armin Rigoa3f09272006-05-28 19:13:17 +0000842>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000843>>> for k, g in groupby(di, itemgetter(1)):
844... print k, map(itemgetter(0), g)
845...
8461 ['a', 'c', 'e']
8472 ['b', 'd', 'f']
8483 ['g']
849
Raymond Hettinger734fb572004-01-20 20:04:40 +0000850# Find runs of consecutive numbers using groupby. The key to the solution
851# is differencing with a range so that consecutive numbers all appear in
852# same group.
853>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
854>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
855... print map(operator.itemgetter(1), g)
856...
857[1]
858[4, 5, 6]
859[10]
860[15, 16, 17, 18]
861[22]
862[25, 26, 27, 28]
863
Raymond Hettingera098b332003-09-08 23:58:40 +0000864>>> def take(n, seq):
865... return list(islice(seq, n))
866
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000867>>> def enumerate(iterable):
868... return izip(count(), iterable)
869
870>>> def tabulate(function):
871... "Return function(0), function(1), ..."
872... return imap(function, count())
873
874>>> def iteritems(mapping):
875... return izip(mapping.iterkeys(), mapping.itervalues())
876
877>>> def nth(iterable, n):
878... "Returns the nth item"
Raymond Hettinger60eca932003-02-09 06:40:58 +0000879... return list(islice(iterable, n, n+1))
880
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000881>>> def all(seq, pred=None):
882... "Returns True if pred(x) is true for every element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000883... for elem in ifilterfalse(pred, seq):
884... return False
885... return True
Raymond Hettinger60eca932003-02-09 06:40:58 +0000886
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000887>>> def any(seq, pred=None):
888... "Returns True if pred(x) is true for at least one element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000889... for elem in ifilter(pred, seq):
890... return True
891... return False
Raymond Hettinger60eca932003-02-09 06:40:58 +0000892
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000893>>> def no(seq, pred=None):
894... "Returns True if pred(x) is false for every element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000895... for elem in ifilter(pred, seq):
896... return False
897... return True
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000898
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000899>>> def quantify(seq, pred=None):
900... "Count how many times the predicate is true in the sequence"
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000901... return sum(imap(pred, seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000902
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000903>>> def padnone(seq):
904... "Returns the sequence elements and then returns None indefinitely"
905... return chain(seq, repeat(None))
906
907>>> def ncycles(seq, n):
908... "Returns the sequence elements n times"
909... return chain(*repeat(seq, n))
910
911>>> def dotproduct(vec1, vec2):
912... return sum(imap(operator.mul, vec1, vec2))
913
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000914>>> def flatten(listOfLists):
915... return list(chain(*listOfLists))
916
917>>> def repeatfunc(func, times=None, *args):
918... "Repeat calls to func with specified arguments."
919... " Example: repeatfunc(random.random)"
920... if times is None:
921... return starmap(func, repeat(args))
922... else:
923... return starmap(func, repeat(args, times))
924
Raymond Hettingerd591f662003-10-26 15:34:50 +0000925>>> def pairwise(iterable):
926... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
927... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000928... try:
929... b.next()
930... except StopIteration:
931... pass
932... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000933
934This is not part of the examples but it tests to make sure the definitions
935perform as purported.
936
Raymond Hettingera098b332003-09-08 23:58:40 +0000937>>> take(10, count())
938[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
939
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000940>>> list(enumerate('abc'))
941[(0, 'a'), (1, 'b'), (2, 'c')]
942
943>>> list(islice(tabulate(lambda x: 2*x), 4))
944[0, 2, 4, 6]
945
946>>> nth('abcde', 3)
947['d']
948
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000949>>> all([2, 4, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000950True
951
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000952>>> all([2, 3, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000953False
954
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000955>>> any([2, 4, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000956True
957
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000958>>> any([1, 3, 5, 9], lambda x: x%2==0,)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000959False
960
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000961>>> no([1, 3, 5, 9], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000962True
963
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000964>>> no([1, 2, 5, 9], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000965False
966
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000967>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +000096850
969
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000970>>> a = [[1, 2, 3], [4, 5, 6]]
971>>> flatten(a)
972[1, 2, 3, 4, 5, 6]
973
974>>> list(repeatfunc(pow, 5, 2, 3))
975[8, 8, 8, 8, 8]
976
977>>> import random
978>>> take(5, imap(int, repeatfunc(random.random)))
979[0, 0, 0, 0, 0]
980
Raymond Hettingerd591f662003-10-26 15:34:50 +0000981>>> list(pairwise('abcd'))
982[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000983
Raymond Hettingerd591f662003-10-26 15:34:50 +0000984>>> list(pairwise([]))
985[]
986
987>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +0000988[]
989
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000990>>> list(islice(padnone('abc'), 0, 6))
991['a', 'b', 'c', None, None, None]
992
993>>> list(ncycles('abc', 3))
994['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
995
996>>> dotproduct([1,2,3], [4,5,6])
99732
998
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000999"""
1000
1001__test__ = {'libreftest' : libreftest}
1002
1003def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001004 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Georg Brandlb84c1372007-01-21 10:28:43 +00001005 RegressionTests, LengthTransparency,
1006 SubclassWithKwargsTest)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001007 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001008
1009 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001010 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001011 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001012 counts = [None] * 5
1013 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001014 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001015 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001016 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001017 print counts
1018
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001019 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001020 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001021
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001022if __name__ == "__main__":
1023 test_main(verbose=True)