Commit | Line | Data |
---|---|---|
920dae64 AT |
1 | import unittest |
2 | from test import test_support | |
3 | from itertools import * | |
4 | from weakref import proxy | |
5 | import sys | |
6 | import operator | |
7 | import random | |
8 | ||
9 | def onearg(x): | |
10 | 'Test function of one argument' | |
11 | return 2*x | |
12 | ||
13 | def errfunc(*args): | |
14 | 'Test function that raises an error' | |
15 | raise ValueError | |
16 | ||
17 | def gen3(): | |
18 | 'Non-restartable source sequence' | |
19 | for i in (0, 1, 2): | |
20 | yield i | |
21 | ||
22 | def isEven(x): | |
23 | 'Test predicate' | |
24 | return x%2==0 | |
25 | ||
26 | def isOdd(x): | |
27 | 'Test predicate' | |
28 | return x%2==1 | |
29 | ||
30 | class StopNow: | |
31 | 'Class emulating an empty iterable.' | |
32 | def __iter__(self): | |
33 | return self | |
34 | def next(self): | |
35 | raise StopIteration | |
36 | ||
37 | def take(n, seq): | |
38 | 'Convenience function for partially consuming a long of infinite iterable' | |
39 | return list(islice(seq, n)) | |
40 | ||
41 | class TestBasicOps(unittest.TestCase): | |
42 | def test_chain(self): | |
43 | self.assertEqual(list(chain('abc', 'def')), list('abcdef')) | |
44 | self.assertEqual(list(chain('abc')), list('abc')) | |
45 | self.assertEqual(list(chain('')), []) | |
46 | self.assertEqual(take(4, chain('abc', 'def')), list('abcd')) | |
47 | self.assertRaises(TypeError, chain, 2, 3) | |
48 | ||
49 | 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)]) | |
52 | self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)]) | |
53 | self.assertRaises(TypeError, count, 2, 3) | |
54 | self.assertRaises(TypeError, count, 'a') | |
55 | c = count(sys.maxint-2) # verify that rollover doesn't crash | |
56 | c.next(); c.next(); c.next(); c.next(); c.next() | |
57 | c = count(3) | |
58 | self.assertEqual(repr(c), 'count(3)') | |
59 | c.next() | |
60 | self.assertEqual(repr(c), 'count(4)') | |
61 | ||
62 | def test_cycle(self): | |
63 | self.assertEqual(take(10, cycle('abc')), list('abcabcabca')) | |
64 | self.assertEqual(list(cycle('')), []) | |
65 | self.assertRaises(TypeError, cycle) | |
66 | self.assertRaises(TypeError, cycle, 5) | |
67 | self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0]) | |
68 | ||
69 | def test_groupby(self): | |
70 | # Check whether it accepts arguments correctly | |
71 | self.assertEqual([], list(groupby([]))) | |
72 | self.assertEqual([], list(groupby([], key=id))) | |
73 | self.assertRaises(TypeError, list, groupby('abc', [])) | |
74 | self.assertRaises(TypeError, groupby, None) | |
75 | self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10) | |
76 | ||
77 | # Check normal input | |
78 | s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22), | |
79 | (2,15,22), (3,16,23), (3,17,23)] | |
80 | dup = [] | |
81 | for k, g in groupby(s, lambda r:r[0]): | |
82 | for elem in g: | |
83 | self.assertEqual(k, elem[0]) | |
84 | dup.append(elem) | |
85 | self.assertEqual(s, dup) | |
86 | ||
87 | # Check nested case | |
88 | dup = [] | |
89 | for k, g in groupby(s, lambda r:r[0]): | |
90 | for ik, ig in groupby(g, lambda r:r[2]): | |
91 | for elem in ig: | |
92 | self.assertEqual(k, elem[0]) | |
93 | self.assertEqual(ik, elem[2]) | |
94 | dup.append(elem) | |
95 | self.assertEqual(s, dup) | |
96 | ||
97 | # Check case where inner iterator is not used | |
98 | keys = [k for k, g in groupby(s, lambda r:r[0])] | |
99 | expectedkeys = set([r[0] for r in s]) | |
100 | self.assertEqual(set(keys), expectedkeys) | |
101 | self.assertEqual(len(keys), len(expectedkeys)) | |
102 | ||
103 | # Exercise pipes and filters style | |
104 | s = 'abracadabra' | |
105 | # sort s | uniq | |
106 | r = [k for k, g in groupby(sorted(s))] | |
107 | self.assertEqual(r, ['a', 'b', 'c', 'd', 'r']) | |
108 | # sort s | uniq -d | |
109 | r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))] | |
110 | self.assertEqual(r, ['a', 'b', 'r']) | |
111 | # sort s | uniq -c | |
112 | r = [(len(list(g)), k) for k, g in groupby(sorted(s))] | |
113 | self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')]) | |
114 | # sort s | uniq -c | sort -rn | head -3 | |
115 | r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3] | |
116 | self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')]) | |
117 | ||
118 | # iter.next failure | |
119 | class ExpectedError(Exception): | |
120 | pass | |
121 | def delayed_raise(n=0): | |
122 | for i in range(n): | |
123 | yield 'yo' | |
124 | raise ExpectedError | |
125 | def gulp(iterable, keyp=None, func=list): | |
126 | return [func(g) for k, g in groupby(iterable, keyp)] | |
127 | ||
128 | # iter.next failure on outer object | |
129 | self.assertRaises(ExpectedError, gulp, delayed_raise(0)) | |
130 | # iter.next failure on inner object | |
131 | self.assertRaises(ExpectedError, gulp, delayed_raise(1)) | |
132 | ||
133 | # __cmp__ failure | |
134 | class DummyCmp: | |
135 | def __cmp__(self, dst): | |
136 | raise ExpectedError | |
137 | s = [DummyCmp(), DummyCmp(), None] | |
138 | ||
139 | # __cmp__ failure on outer object | |
140 | self.assertRaises(ExpectedError, gulp, s, func=id) | |
141 | # __cmp__ failure on inner object | |
142 | self.assertRaises(ExpectedError, gulp, s) | |
143 | ||
144 | # keyfunc failure | |
145 | def keyfunc(obj): | |
146 | if keyfunc.skip > 0: | |
147 | keyfunc.skip -= 1 | |
148 | return obj | |
149 | else: | |
150 | raise ExpectedError | |
151 | ||
152 | # keyfunc failure on outer object | |
153 | keyfunc.skip = 0 | |
154 | self.assertRaises(ExpectedError, gulp, [None], keyfunc) | |
155 | keyfunc.skip = 1 | |
156 | self.assertRaises(ExpectedError, gulp, [None, None], keyfunc) | |
157 | ||
158 | def test_ifilter(self): | |
159 | self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4]) | |
160 | self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2]) | |
161 | self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6]) | |
162 | self.assertRaises(TypeError, ifilter) | |
163 | self.assertRaises(TypeError, ifilter, lambda x:x) | |
164 | self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7) | |
165 | self.assertRaises(TypeError, ifilter, isEven, 3) | |
166 | self.assertRaises(TypeError, ifilter(range(6), range(6)).next) | |
167 | ||
168 | def test_ifilterfalse(self): | |
169 | self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5]) | |
170 | self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0]) | |
171 | self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7]) | |
172 | self.assertRaises(TypeError, ifilterfalse) | |
173 | self.assertRaises(TypeError, ifilterfalse, lambda x:x) | |
174 | self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7) | |
175 | self.assertRaises(TypeError, ifilterfalse, isEven, 3) | |
176 | self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next) | |
177 | ||
178 | def test_izip(self): | |
179 | ans = [(x,y) for x, y in izip('abc',count())] | |
180 | self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)]) | |
181 | self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6))) | |
182 | self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3))) | |
183 | self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3))) | |
184 | self.assertEqual(list(izip('abcdef')), zip('abcdef')) | |
185 | self.assertEqual(list(izip()), zip()) | |
186 | self.assertRaises(TypeError, izip, 3) | |
187 | self.assertRaises(TypeError, izip, range(3), 3) | |
188 | # Check tuple re-use (implementation detail) | |
189 | self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')], | |
190 | zip('abc', 'def')) | |
191 | self.assertEqual([pair for pair in izip('abc', 'def')], | |
192 | zip('abc', 'def')) | |
193 | ids = map(id, izip('abc', 'def')) | |
194 | self.assertEqual(min(ids), max(ids)) | |
195 | ids = map(id, list(izip('abc', 'def'))) | |
196 | self.assertEqual(len(dict.fromkeys(ids)), len(ids)) | |
197 | ||
198 | def test_repeat(self): | |
199 | self.assertEqual(zip(xrange(3),repeat('a')), | |
200 | [(0, 'a'), (1, 'a'), (2, 'a')]) | |
201 | self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a']) | |
202 | self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a']) | |
203 | self.assertEqual(list(repeat('a', 0)), []) | |
204 | self.assertEqual(list(repeat('a', -3)), []) | |
205 | self.assertRaises(TypeError, repeat) | |
206 | self.assertRaises(TypeError, repeat, None, 3, 4) | |
207 | self.assertRaises(TypeError, repeat, None, 'a') | |
208 | r = repeat(1+0j) | |
209 | self.assertEqual(repr(r), 'repeat((1+0j))') | |
210 | r = repeat(1+0j, 5) | |
211 | self.assertEqual(repr(r), 'repeat((1+0j), 5)') | |
212 | list(r) | |
213 | self.assertEqual(repr(r), 'repeat((1+0j), 0)') | |
214 | ||
215 | def test_imap(self): | |
216 | self.assertEqual(list(imap(operator.pow, range(3), range(1,7))), | |
217 | [0**1, 1**2, 2**3]) | |
218 | self.assertEqual(list(imap(None, 'abc', range(5))), | |
219 | [('a',0),('b',1),('c',2)]) | |
220 | self.assertEqual(list(imap(None, 'abc', count())), | |
221 | [('a',0),('b',1),('c',2)]) | |
222 | self.assertEqual(take(2,imap(None, 'abc', count())), | |
223 | [('a',0),('b',1)]) | |
224 | self.assertEqual(list(imap(operator.pow, [])), []) | |
225 | self.assertRaises(TypeError, imap) | |
226 | self.assertRaises(TypeError, imap, operator.neg) | |
227 | self.assertRaises(TypeError, imap(10, range(5)).next) | |
228 | self.assertRaises(ValueError, imap(errfunc, [4], [5]).next) | |
229 | self.assertRaises(TypeError, imap(onearg, [4], [5]).next) | |
230 | ||
231 | def test_starmap(self): | |
232 | self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))), | |
233 | [0**1, 1**2, 2**3]) | |
234 | self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))), | |
235 | [0**1, 1**2, 2**3]) | |
236 | self.assertEqual(list(starmap(operator.pow, [])), []) | |
237 | self.assertRaises(TypeError, list, starmap(operator.pow, [[4,5]])) | |
238 | self.assertRaises(TypeError, starmap) | |
239 | self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra') | |
240 | self.assertRaises(TypeError, starmap(10, [(4,5)]).next) | |
241 | self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next) | |
242 | self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next) | |
243 | ||
244 | def test_islice(self): | |
245 | for args in [ # islice(args) should agree with range(args) | |
246 | (10, 20, 3), | |
247 | (10, 3, 20), | |
248 | (10, 20), | |
249 | (10, 3), | |
250 | (20,) | |
251 | ]: | |
252 | self.assertEqual(list(islice(xrange(100), *args)), range(*args)) | |
253 | ||
254 | for args, tgtargs in [ # Stop when seqn is exhausted | |
255 | ((10, 110, 3), ((10, 100, 3))), | |
256 | ((10, 110), ((10, 100))), | |
257 | ((110,), (100,)) | |
258 | ]: | |
259 | self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs)) | |
260 | ||
261 | # Test stop=None | |
262 | self.assertEqual(list(islice(xrange(10), None)), range(10)) | |
263 | self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10)) | |
264 | self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2)) | |
265 | ||
266 | # Test invalid arguments | |
267 | self.assertRaises(TypeError, islice, xrange(10)) | |
268 | self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4) | |
269 | self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1) | |
270 | self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1) | |
271 | self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1) | |
272 | self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0) | |
273 | self.assertRaises(ValueError, islice, xrange(10), 'a') | |
274 | self.assertRaises(ValueError, islice, xrange(10), 'a', 1) | |
275 | self.assertRaises(ValueError, islice, xrange(10), 1, 'a') | |
276 | self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1) | |
277 | self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1) | |
278 | self.assertEqual(len(list(islice(count(), 1, 10, sys.maxint))), 1) | |
279 | ||
280 | def test_takewhile(self): | |
281 | data = [1, 3, 5, 20, 2, 4, 6, 8] | |
282 | underten = lambda x: x<10 | |
283 | self.assertEqual(list(takewhile(underten, data)), [1, 3, 5]) | |
284 | self.assertEqual(list(takewhile(underten, [])), []) | |
285 | self.assertRaises(TypeError, takewhile) | |
286 | self.assertRaises(TypeError, takewhile, operator.pow) | |
287 | self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra') | |
288 | self.assertRaises(TypeError, takewhile(10, [(4,5)]).next) | |
289 | self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next) | |
290 | t = takewhile(bool, [1, 1, 1, 0, 0, 0]) | |
291 | self.assertEqual(list(t), [1, 1, 1]) | |
292 | self.assertRaises(StopIteration, t.next) | |
293 | ||
294 | def test_dropwhile(self): | |
295 | data = [1, 3, 5, 20, 2, 4, 6, 8] | |
296 | underten = lambda x: x<10 | |
297 | self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8]) | |
298 | self.assertEqual(list(dropwhile(underten, [])), []) | |
299 | self.assertRaises(TypeError, dropwhile) | |
300 | self.assertRaises(TypeError, dropwhile, operator.pow) | |
301 | self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra') | |
302 | self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next) | |
303 | self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next) | |
304 | ||
305 | def test_tee(self): | |
306 | n = 200 | |
307 | def irange(n): | |
308 | for i in xrange(n): | |
309 | yield i | |
310 | ||
311 | a, b = tee([]) # test empty iterator | |
312 | self.assertEqual(list(a), []) | |
313 | self.assertEqual(list(b), []) | |
314 | ||
315 | a, b = tee(irange(n)) # test 100% interleaved | |
316 | self.assertEqual(zip(a,b), zip(range(n),range(n))) | |
317 | ||
318 | a, b = tee(irange(n)) # test 0% interleaved | |
319 | self.assertEqual(list(a), range(n)) | |
320 | self.assertEqual(list(b), range(n)) | |
321 | ||
322 | a, b = tee(irange(n)) # test dealloc of leading iterator | |
323 | for i in xrange(100): | |
324 | self.assertEqual(a.next(), i) | |
325 | del a | |
326 | self.assertEqual(list(b), range(n)) | |
327 | ||
328 | a, b = tee(irange(n)) # test dealloc of trailing iterator | |
329 | for i in xrange(100): | |
330 | self.assertEqual(a.next(), i) | |
331 | del b | |
332 | self.assertEqual(list(a), range(100, n)) | |
333 | ||
334 | for j in xrange(5): # test randomly interleaved | |
335 | order = [0]*n + [1]*n | |
336 | random.shuffle(order) | |
337 | lists = ([], []) | |
338 | its = tee(irange(n)) | |
339 | for i in order: | |
340 | value = its[i].next() | |
341 | lists[i].append(value) | |
342 | self.assertEqual(lists[0], range(n)) | |
343 | self.assertEqual(lists[1], range(n)) | |
344 | ||
345 | # test argument format checking | |
346 | self.assertRaises(TypeError, tee) | |
347 | self.assertRaises(TypeError, tee, 3) | |
348 | self.assertRaises(TypeError, tee, [1,2], 'x') | |
349 | self.assertRaises(TypeError, tee, [1,2], 3, 'x') | |
350 | ||
351 | # tee object should be instantiable | |
352 | a, b = tee('abc') | |
353 | c = type(a)('def') | |
354 | self.assertEqual(list(c), list('def')) | |
355 | ||
356 | # test long-lagged and multi-way split | |
357 | a, b, c = tee(xrange(2000), 3) | |
358 | for i in xrange(100): | |
359 | self.assertEqual(a.next(), i) | |
360 | self.assertEqual(list(b), range(2000)) | |
361 | self.assertEqual([c.next(), c.next()], range(2)) | |
362 | self.assertEqual(list(a), range(100,2000)) | |
363 | self.assertEqual(list(c), range(2,2000)) | |
364 | ||
365 | # test values of n | |
366 | self.assertRaises(TypeError, tee, 'abc', 'invalid') | |
367 | for n in xrange(5): | |
368 | result = tee('abc', n) | |
369 | self.assertEqual(type(result), tuple) | |
370 | self.assertEqual(len(result), n) | |
371 | self.assertEqual(map(list, result), [list('abc')]*n) | |
372 | ||
373 | # tee pass-through to copyable iterator | |
374 | a, b = tee('abc') | |
375 | c, d = tee(a) | |
376 | self.assert_(a is c) | |
377 | ||
378 | # test tee_new | |
379 | t1, t2 = tee('abc') | |
380 | tnew = type(t1) | |
381 | self.assertRaises(TypeError, tnew) | |
382 | self.assertRaises(TypeError, tnew, 10) | |
383 | t3 = tnew(t1) | |
384 | self.assert_(list(t1) == list(t2) == list(t3) == list('abc')) | |
385 | ||
386 | # test that tee objects are weak referencable | |
387 | a, b = tee(xrange(10)) | |
388 | p = proxy(a) | |
389 | self.assertEqual(getattr(p, '__class__'), type(b)) | |
390 | del a | |
391 | self.assertRaises(ReferenceError, getattr, p, '__class__') | |
392 | ||
393 | def test_StopIteration(self): | |
394 | self.assertRaises(StopIteration, izip().next) | |
395 | ||
396 | for f in (chain, cycle, izip, groupby): | |
397 | self.assertRaises(StopIteration, f([]).next) | |
398 | self.assertRaises(StopIteration, f(StopNow()).next) | |
399 | ||
400 | self.assertRaises(StopIteration, islice([], None).next) | |
401 | self.assertRaises(StopIteration, islice(StopNow(), None).next) | |
402 | ||
403 | p, q = tee([]) | |
404 | self.assertRaises(StopIteration, p.next) | |
405 | self.assertRaises(StopIteration, q.next) | |
406 | p, q = tee(StopNow()) | |
407 | self.assertRaises(StopIteration, p.next) | |
408 | self.assertRaises(StopIteration, q.next) | |
409 | ||
410 | self.assertRaises(StopIteration, repeat(None, 0).next) | |
411 | ||
412 | for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap): | |
413 | self.assertRaises(StopIteration, f(lambda x:x, []).next) | |
414 | self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next) | |
415 | ||
416 | class TestGC(unittest.TestCase): | |
417 | ||
418 | def makecycle(self, iterator, container): | |
419 | container.append(iterator) | |
420 | iterator.next() | |
421 | del container, iterator | |
422 | ||
423 | def test_chain(self): | |
424 | a = [] | |
425 | self.makecycle(chain(a), a) | |
426 | ||
427 | def test_cycle(self): | |
428 | a = [] | |
429 | self.makecycle(cycle([a]*2), a) | |
430 | ||
431 | def test_dropwhile(self): | |
432 | a = [] | |
433 | self.makecycle(dropwhile(bool, [0, a, a]), a) | |
434 | ||
435 | def test_groupby(self): | |
436 | a = [] | |
437 | self.makecycle(groupby([a]*2, lambda x:x), a) | |
438 | ||
439 | def test_ifilter(self): | |
440 | a = [] | |
441 | self.makecycle(ifilter(lambda x:True, [a]*2), a) | |
442 | ||
443 | def test_ifilterfalse(self): | |
444 | a = [] | |
445 | self.makecycle(ifilterfalse(lambda x:False, a), a) | |
446 | ||
447 | def test_izip(self): | |
448 | a = [] | |
449 | self.makecycle(izip([a]*2, [a]*3), a) | |
450 | ||
451 | def test_imap(self): | |
452 | a = [] | |
453 | self.makecycle(imap(lambda x:x, [a]*2), a) | |
454 | ||
455 | def test_islice(self): | |
456 | a = [] | |
457 | self.makecycle(islice([a]*2, None), a) | |
458 | ||
459 | def test_repeat(self): | |
460 | a = [] | |
461 | self.makecycle(repeat(a), a) | |
462 | ||
463 | def test_starmap(self): | |
464 | a = [] | |
465 | self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a) | |
466 | ||
467 | def test_takewhile(self): | |
468 | a = [] | |
469 | self.makecycle(takewhile(bool, [1, 0, a, a]), a) | |
470 | ||
471 | def R(seqn): | |
472 | 'Regular generator' | |
473 | for i in seqn: | |
474 | yield i | |
475 | ||
476 | class G: | |
477 | 'Sequence using __getitem__' | |
478 | def __init__(self, seqn): | |
479 | self.seqn = seqn | |
480 | def __getitem__(self, i): | |
481 | return self.seqn[i] | |
482 | ||
483 | class I: | |
484 | 'Sequence using iterator protocol' | |
485 | def __init__(self, seqn): | |
486 | self.seqn = seqn | |
487 | self.i = 0 | |
488 | def __iter__(self): | |
489 | return self | |
490 | def next(self): | |
491 | if self.i >= len(self.seqn): raise StopIteration | |
492 | v = self.seqn[self.i] | |
493 | self.i += 1 | |
494 | return v | |
495 | ||
496 | class Ig: | |
497 | 'Sequence using iterator protocol defined with a generator' | |
498 | def __init__(self, seqn): | |
499 | self.seqn = seqn | |
500 | self.i = 0 | |
501 | def __iter__(self): | |
502 | for val in self.seqn: | |
503 | yield val | |
504 | ||
505 | class X: | |
506 | 'Missing __getitem__ and __iter__' | |
507 | def __init__(self, seqn): | |
508 | self.seqn = seqn | |
509 | self.i = 0 | |
510 | def next(self): | |
511 | if self.i >= len(self.seqn): raise StopIteration | |
512 | v = self.seqn[self.i] | |
513 | self.i += 1 | |
514 | return v | |
515 | ||
516 | class N: | |
517 | 'Iterator missing next()' | |
518 | def __init__(self, seqn): | |
519 | self.seqn = seqn | |
520 | self.i = 0 | |
521 | def __iter__(self): | |
522 | return self | |
523 | ||
524 | class E: | |
525 | 'Test propagation of exceptions' | |
526 | def __init__(self, seqn): | |
527 | self.seqn = seqn | |
528 | self.i = 0 | |
529 | def __iter__(self): | |
530 | return self | |
531 | def next(self): | |
532 | 3 // 0 | |
533 | ||
534 | class S: | |
535 | 'Test immediate stop' | |
536 | def __init__(self, seqn): | |
537 | pass | |
538 | def __iter__(self): | |
539 | return self | |
540 | def next(self): | |
541 | raise StopIteration | |
542 | ||
543 | def L(seqn): | |
544 | 'Test multiple tiers of iterators' | |
545 | return chain(imap(lambda x:x, R(Ig(G(seqn))))) | |
546 | ||
547 | ||
548 | class TestVariousIteratorArgs(unittest.TestCase): | |
549 | ||
550 | def test_chain(self): | |
551 | for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)): | |
552 | for g in (G, I, Ig, S, L, R): | |
553 | self.assertEqual(list(chain(g(s))), list(g(s))) | |
554 | self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s))) | |
555 | self.assertRaises(TypeError, chain, X(s)) | |
556 | self.assertRaises(TypeError, list, chain(N(s))) | |
557 | self.assertRaises(ZeroDivisionError, list, chain(E(s))) | |
558 | ||
559 | def test_cycle(self): | |
560 | for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)): | |
561 | for g in (G, I, Ig, S, L, R): | |
562 | tgtlen = len(s) * 3 | |
563 | expected = list(g(s))*3 | |
564 | actual = list(islice(cycle(g(s)), tgtlen)) | |
565 | self.assertEqual(actual, expected) | |
566 | self.assertRaises(TypeError, cycle, X(s)) | |
567 | self.assertRaises(TypeError, list, cycle(N(s))) | |
568 | self.assertRaises(ZeroDivisionError, list, cycle(E(s))) | |
569 | ||
570 | def test_groupby(self): | |
571 | for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)): | |
572 | for g in (G, I, Ig, S, L, R): | |
573 | self.assertEqual([k for k, sb in groupby(g(s))], list(g(s))) | |
574 | self.assertRaises(TypeError, groupby, X(s)) | |
575 | self.assertRaises(TypeError, list, groupby(N(s))) | |
576 | self.assertRaises(ZeroDivisionError, list, groupby(E(s))) | |
577 | ||
578 | def test_ifilter(self): | |
579 | for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)): | |
580 | for g in (G, I, Ig, S, L, R): | |
581 | self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s))) | |
582 | self.assertRaises(TypeError, ifilter, isEven, X(s)) | |
583 | self.assertRaises(TypeError, list, ifilter(isEven, N(s))) | |
584 | self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s))) | |
585 | ||
586 | def test_ifilterfalse(self): | |
587 | for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)): | |
588 | for g in (G, I, Ig, S, L, R): | |
589 | self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s))) | |
590 | self.assertRaises(TypeError, ifilterfalse, isEven, X(s)) | |
591 | self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s))) | |
592 | self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s))) | |
593 | ||
594 | def test_izip(self): | |
595 | for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)): | |
596 | for g in (G, I, Ig, S, L, R): | |
597 | self.assertEqual(list(izip(g(s))), zip(g(s))) | |
598 | self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s))) | |
599 | self.assertRaises(TypeError, izip, X(s)) | |
600 | self.assertRaises(TypeError, list, izip(N(s))) | |
601 | self.assertRaises(ZeroDivisionError, list, izip(E(s))) | |
602 | ||
603 | def test_imap(self): | |
604 | for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)): | |
605 | for g in (G, I, Ig, S, L, R): | |
606 | self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s))) | |
607 | self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s))) | |
608 | self.assertRaises(TypeError, imap, onearg, X(s)) | |
609 | self.assertRaises(TypeError, list, imap(onearg, N(s))) | |
610 | self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s))) | |
611 | ||
612 | def test_islice(self): | |
613 | for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)): | |
614 | for g in (G, I, Ig, S, L, R): | |
615 | self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2]) | |
616 | self.assertRaises(TypeError, islice, X(s), 10) | |
617 | self.assertRaises(TypeError, list, islice(N(s), 10)) | |
618 | self.assertRaises(ZeroDivisionError, list, islice(E(s), 10)) | |
619 | ||
620 | def test_starmap(self): | |
621 | for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)): | |
622 | for g in (G, I, Ig, S, L, R): | |
623 | ss = zip(s, s) | |
624 | self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s))) | |
625 | self.assertRaises(TypeError, starmap, operator.pow, X(ss)) | |
626 | self.assertRaises(TypeError, list, starmap(operator.pow, N(ss))) | |
627 | self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss))) | |
628 | ||
629 | def test_takewhile(self): | |
630 | for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)): | |
631 | for g in (G, I, Ig, S, L, R): | |
632 | tgt = [] | |
633 | for elem in g(s): | |
634 | if not isEven(elem): break | |
635 | tgt.append(elem) | |
636 | self.assertEqual(list(takewhile(isEven, g(s))), tgt) | |
637 | self.assertRaises(TypeError, takewhile, isEven, X(s)) | |
638 | self.assertRaises(TypeError, list, takewhile(isEven, N(s))) | |
639 | self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s))) | |
640 | ||
641 | def test_dropwhile(self): | |
642 | for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)): | |
643 | for g in (G, I, Ig, S, L, R): | |
644 | tgt = [] | |
645 | for elem in g(s): | |
646 | if not tgt and isOdd(elem): continue | |
647 | tgt.append(elem) | |
648 | self.assertEqual(list(dropwhile(isOdd, g(s))), tgt) | |
649 | self.assertRaises(TypeError, dropwhile, isOdd, X(s)) | |
650 | self.assertRaises(TypeError, list, dropwhile(isOdd, N(s))) | |
651 | self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s))) | |
652 | ||
653 | def test_tee(self): | |
654 | for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)): | |
655 | for g in (G, I, Ig, S, L, R): | |
656 | it1, it2 = tee(g(s)) | |
657 | self.assertEqual(list(it1), list(g(s))) | |
658 | self.assertEqual(list(it2), list(g(s))) | |
659 | self.assertRaises(TypeError, tee, X(s)) | |
660 | self.assertRaises(TypeError, list, tee(N(s))[0]) | |
661 | self.assertRaises(ZeroDivisionError, list, tee(E(s))[0]) | |
662 | ||
663 | class LengthTransparency(unittest.TestCase): | |
664 | ||
665 | def test_repeat(self): | |
666 | self.assertEqual(len(repeat(None, 50)), 50) | |
667 | self.assertRaises(TypeError, len, repeat(None)) | |
668 | ||
669 | class RegressionTests(unittest.TestCase): | |
670 | ||
671 | def test_sf_793826(self): | |
672 | # Fix Armin Rigo's successful efforts to wreak havoc | |
673 | ||
674 | def mutatingtuple(tuple1, f, tuple2): | |
675 | # this builds a tuple t which is a copy of tuple1, | |
676 | # then calls f(t), then mutates t to be equal to tuple2 | |
677 | # (needs len(tuple1) == len(tuple2)). | |
678 | def g(value, first=[1]): | |
679 | if first: | |
680 | del first[:] | |
681 | f(z.next()) | |
682 | return value | |
683 | items = list(tuple2) | |
684 | items[1:1] = list(tuple1) | |
685 | gen = imap(g, items) | |
686 | z = izip(*[gen]*len(tuple1)) | |
687 | z.next() | |
688 | ||
689 | def f(t): | |
690 | global T | |
691 | T = t | |
692 | first[:] = list(T) | |
693 | ||
694 | first = [] | |
695 | mutatingtuple((1,2,3), f, (4,5,6)) | |
696 | second = list(T) | |
697 | self.assertEqual(first, second) | |
698 | ||
699 | ||
700 | def test_sf_950057(self): | |
701 | # Make sure that chain() and cycle() catch exceptions immediately | |
702 | # rather than when shifting between input sources | |
703 | ||
704 | def gen1(): | |
705 | hist.append(0) | |
706 | yield 1 | |
707 | hist.append(1) | |
708 | raise AssertionError | |
709 | hist.append(2) | |
710 | ||
711 | def gen2(x): | |
712 | hist.append(3) | |
713 | yield 2 | |
714 | hist.append(4) | |
715 | if x: | |
716 | raise StopIteration | |
717 | ||
718 | hist = [] | |
719 | self.assertRaises(AssertionError, list, chain(gen1(), gen2(False))) | |
720 | self.assertEqual(hist, [0,1]) | |
721 | ||
722 | hist = [] | |
723 | self.assertRaises(AssertionError, list, chain(gen1(), gen2(True))) | |
724 | self.assertEqual(hist, [0,1]) | |
725 | ||
726 | hist = [] | |
727 | self.assertRaises(AssertionError, list, cycle(gen1())) | |
728 | self.assertEqual(hist, [0,1]) | |
729 | ||
730 | libreftest = """ Doctest for examples in the library reference: libitertools.tex | |
731 | ||
732 | ||
733 | >>> amounts = [120.15, 764.05, 823.14] | |
734 | >>> for checknum, amount in izip(count(1200), amounts): | |
735 | ... print 'Check %d is for $%.2f' % (checknum, amount) | |
736 | ... | |
737 | Check 1200 is for $120.15 | |
738 | Check 1201 is for $764.05 | |
739 | Check 1202 is for $823.14 | |
740 | ||
741 | >>> import operator | |
742 | >>> for cube in imap(operator.pow, xrange(1,4), repeat(3)): | |
743 | ... print cube | |
744 | ... | |
745 | 1 | |
746 | 8 | |
747 | 27 | |
748 | ||
749 | >>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele'] | |
750 | >>> for name in islice(reportlines, 3, None, 2): | |
751 | ... print name.title() | |
752 | ... | |
753 | Alex | |
754 | Laura | |
755 | Martin | |
756 | Walter | |
757 | Samuele | |
758 | ||
759 | >>> from operator import itemgetter | |
760 | >>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3) | |
761 | >>> di = sorted(d.iteritems(), key=itemgetter(1)) | |
762 | >>> for k, g in groupby(di, itemgetter(1)): | |
763 | ... print k, map(itemgetter(0), g) | |
764 | ... | |
765 | 1 ['a', 'c', 'e'] | |
766 | 2 ['b', 'd', 'f'] | |
767 | 3 ['g'] | |
768 | ||
769 | # Find runs of consecutive numbers using groupby. The key to the solution | |
770 | # is differencing with a range so that consecutive numbers all appear in | |
771 | # same group. | |
772 | >>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28] | |
773 | >>> for k, g in groupby(enumerate(data), lambda (i,x):i-x): | |
774 | ... print map(operator.itemgetter(1), g) | |
775 | ... | |
776 | [1] | |
777 | [4, 5, 6] | |
778 | [10] | |
779 | [15, 16, 17, 18] | |
780 | [22] | |
781 | [25, 26, 27, 28] | |
782 | ||
783 | >>> def take(n, seq): | |
784 | ... return list(islice(seq, n)) | |
785 | ||
786 | >>> def enumerate(iterable): | |
787 | ... return izip(count(), iterable) | |
788 | ||
789 | >>> def tabulate(function): | |
790 | ... "Return function(0), function(1), ..." | |
791 | ... return imap(function, count()) | |
792 | ||
793 | >>> def iteritems(mapping): | |
794 | ... return izip(mapping.iterkeys(), mapping.itervalues()) | |
795 | ||
796 | >>> def nth(iterable, n): | |
797 | ... "Returns the nth item" | |
798 | ... return list(islice(iterable, n, n+1)) | |
799 | ||
800 | >>> def all(seq, pred=bool): | |
801 | ... "Returns True if pred(x) is True for every element in the iterable" | |
802 | ... for elem in ifilterfalse(pred, seq): | |
803 | ... return False | |
804 | ... return True | |
805 | ||
806 | >>> def any(seq, pred=bool): | |
807 | ... "Returns True if pred(x) is True for at least one element in the iterable" | |
808 | ... for elem in ifilter(pred, seq): | |
809 | ... return True | |
810 | ... return False | |
811 | ||
812 | >>> def no(seq, pred=bool): | |
813 | ... "Returns True if pred(x) is False for every element in the iterable" | |
814 | ... for elem in ifilter(pred, seq): | |
815 | ... return False | |
816 | ... return True | |
817 | ||
818 | >>> def quantify(seq, pred=bool): | |
819 | ... "Count how many times the predicate is True in the sequence" | |
820 | ... return sum(imap(pred, seq)) | |
821 | ||
822 | >>> def padnone(seq): | |
823 | ... "Returns the sequence elements and then returns None indefinitely" | |
824 | ... return chain(seq, repeat(None)) | |
825 | ||
826 | >>> def ncycles(seq, n): | |
827 | ... "Returns the sequence elements n times" | |
828 | ... return chain(*repeat(seq, n)) | |
829 | ||
830 | >>> def dotproduct(vec1, vec2): | |
831 | ... return sum(imap(operator.mul, vec1, vec2)) | |
832 | ||
833 | >>> def flatten(listOfLists): | |
834 | ... return list(chain(*listOfLists)) | |
835 | ||
836 | >>> def repeatfunc(func, times=None, *args): | |
837 | ... "Repeat calls to func with specified arguments." | |
838 | ... " Example: repeatfunc(random.random)" | |
839 | ... if times is None: | |
840 | ... return starmap(func, repeat(args)) | |
841 | ... else: | |
842 | ... return starmap(func, repeat(args, times)) | |
843 | ||
844 | >>> def pairwise(iterable): | |
845 | ... "s -> (s0,s1), (s1,s2), (s2, s3), ..." | |
846 | ... a, b = tee(iterable) | |
847 | ... try: | |
848 | ... b.next() | |
849 | ... except StopIteration: | |
850 | ... pass | |
851 | ... return izip(a, b) | |
852 | ||
853 | This is not part of the examples but it tests to make sure the definitions | |
854 | perform as purported. | |
855 | ||
856 | >>> take(10, count()) | |
857 | [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] | |
858 | ||
859 | >>> list(enumerate('abc')) | |
860 | [(0, 'a'), (1, 'b'), (2, 'c')] | |
861 | ||
862 | >>> list(islice(tabulate(lambda x: 2*x), 4)) | |
863 | [0, 2, 4, 6] | |
864 | ||
865 | >>> nth('abcde', 3) | |
866 | ['d'] | |
867 | ||
868 | >>> all([2, 4, 6, 8], lambda x: x%2==0) | |
869 | True | |
870 | ||
871 | >>> all([2, 3, 6, 8], lambda x: x%2==0) | |
872 | False | |
873 | ||
874 | >>> any([2, 4, 6, 8], lambda x: x%2==0) | |
875 | True | |
876 | ||
877 | >>> any([1, 3, 5, 9], lambda x: x%2==0,) | |
878 | False | |
879 | ||
880 | >>> no([1, 3, 5, 9], lambda x: x%2==0) | |
881 | True | |
882 | ||
883 | >>> no([1, 2, 5, 9], lambda x: x%2==0) | |
884 | False | |
885 | ||
886 | >>> quantify(xrange(99), lambda x: x%2==0) | |
887 | 50 | |
888 | ||
889 | >>> a = [[1, 2, 3], [4, 5, 6]] | |
890 | >>> flatten(a) | |
891 | [1, 2, 3, 4, 5, 6] | |
892 | ||
893 | >>> list(repeatfunc(pow, 5, 2, 3)) | |
894 | [8, 8, 8, 8, 8] | |
895 | ||
896 | >>> import random | |
897 | >>> take(5, imap(int, repeatfunc(random.random))) | |
898 | [0, 0, 0, 0, 0] | |
899 | ||
900 | >>> list(pairwise('abcd')) | |
901 | [('a', 'b'), ('b', 'c'), ('c', 'd')] | |
902 | ||
903 | >>> list(pairwise([])) | |
904 | [] | |
905 | ||
906 | >>> list(pairwise('a')) | |
907 | [] | |
908 | ||
909 | >>> list(islice(padnone('abc'), 0, 6)) | |
910 | ['a', 'b', 'c', None, None, None] | |
911 | ||
912 | >>> list(ncycles('abc', 3)) | |
913 | ['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c'] | |
914 | ||
915 | >>> dotproduct([1,2,3], [4,5,6]) | |
916 | 32 | |
917 | ||
918 | """ | |
919 | ||
920 | __test__ = {'libreftest' : libreftest} | |
921 | ||
922 | def test_main(verbose=None): | |
923 | test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC, | |
924 | RegressionTests, LengthTransparency) | |
925 | test_support.run_unittest(*test_classes) | |
926 | ||
927 | # verify reference counting | |
928 | if verbose and hasattr(sys, "gettotalrefcount"): | |
929 | import gc | |
930 | counts = [None] * 5 | |
931 | for i in xrange(len(counts)): | |
932 | test_support.run_unittest(*test_classes) | |
933 | gc.collect() | |
934 | counts[i] = sys.gettotalrefcount() | |
935 | print counts | |
936 | ||
937 | # doctest the examples in the library reference | |
938 | test_support.run_doctest(sys.modules[__name__], verbose) | |
939 | ||
940 | if __name__ == "__main__": | |
941 | test_main(verbose=True) |