blob: 069274745180f894a7ff79554ac969d798582024 [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
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +00008maxsize = test_support.MAX_Py_ssize_t
9minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000010
11def onearg(x):
12 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000013 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000014
15def errfunc(*args):
16 'Test function that raises an error'
17 raise ValueError
18
19def gen3():
20 'Non-restartable source sequence'
21 for i in (0, 1, 2):
22 yield i
23
24def isEven(x):
25 'Test predicate'
26 return x%2==0
27
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000028def isOdd(x):
29 'Test predicate'
30 return x%2==1
31
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000032class StopNow:
33 'Class emulating an empty iterable.'
34 def __iter__(self):
35 return self
36 def next(self):
37 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000038
Raymond Hettinger02420702003-06-29 20:36:23 +000039def take(n, seq):
40 'Convenience function for partially consuming a long of infinite iterable'
41 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000042
Raymond Hettingerd553d852008-03-04 04:17:08 +000043def prod(iterable):
44 return reduce(operator.mul, iterable, 1)
45
Raymond Hettinger93e804d2008-02-26 23:40:50 +000046def fact(n):
47 'Factorial'
Raymond Hettingerd553d852008-03-04 04:17:08 +000048 return prod(range(1, n+1))
49
Raymond Hettinger96ef8112003-02-01 00:10:11 +000050class TestBasicOps(unittest.TestCase):
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000051 def test_chain(self):
52 self.assertEqual(list(chain('abc', 'def')), list('abcdef'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000053 self.assertEqual(list(chain('abc')), list('abc'))
54 self.assertEqual(list(chain('')), [])
Raymond Hettinger02420702003-06-29 20:36:23 +000055 self.assertEqual(take(4, chain('abc', 'def')), list('abcd'))
Raymond Hettinger05bf6332008-02-28 22:30:42 +000056 self.assertRaises(TypeError, list,chain(2, 3))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000057
Raymond Hettingerb4cbc982008-02-28 22:46:41 +000058 def test_chain_from_iterable(self):
59 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
60 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
61 self.assertEqual(list(chain.from_iterable([''])), [])
62 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
63 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
64
Raymond Hettinger93e804d2008-02-26 23:40:50 +000065 def test_combinations(self):
66 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
67 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Raymond Hettingerd553d852008-03-04 04:17:08 +000068 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Raymond Hettinger93e804d2008-02-26 23:40:50 +000069 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
70 self.assertRaises(ValueError, combinations, 'abc', 32) # r is too big
71 self.assertEqual(list(combinations(range(4), 3)),
72 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
Raymond Hettingerd553d852008-03-04 04:17:08 +000073
74 def combinations1(iterable, r):
75 'Pure python version shown in the docs'
76 pool = tuple(iterable)
77 n = len(pool)
78 indices = range(r)
79 yield tuple(pool[i] for i in indices)
80 while 1:
81 for i in reversed(range(r)):
82 if indices[i] != i + n - r:
83 break
84 else:
85 return
86 indices[i] += 1
87 for j in range(i+1, r):
88 indices[j] = indices[j-1] + 1
89 yield tuple(pool[i] for i in indices)
90
91 def combinations2(iterable, r):
92 'Pure python version shown in the docs'
93 pool = tuple(iterable)
94 n = len(pool)
95 for indices in permutations(range(n), r):
96 if sorted(indices) == list(indices):
97 yield tuple(pool[i] for i in indices)
98
99 for n in range(7):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000100 values = [5*x-12 for x in range(n)]
101 for r in range(n+1):
102 result = list(combinations(values, r))
103 self.assertEqual(len(result), fact(n) / fact(r) / fact(n-r)) # right number of combs
104 self.assertEqual(len(result), len(set(result))) # no repeats
105 self.assertEqual(result, sorted(result)) # lexicographic order
106 for c in result:
107 self.assertEqual(len(c), r) # r-length combinations
108 self.assertEqual(len(set(c)), r) # no duplicate elements
109 self.assertEqual(list(c), sorted(c)) # keep original ordering
110 self.assert_(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000111 self.assertEqual(list(c),
112 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Raymond Hettingerd553d852008-03-04 04:17:08 +0000113 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
114 self.assertEqual(result, list(combinations2(values, r))) # matches first pure python version
115
116 # Test implementation detail: tuple re-use
117 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
118 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
119
120 def test_permutations(self):
121 self.assertRaises(TypeError, permutations) # too few arguments
122 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000123 self.assertRaises(TypeError, permutations, None) # pool is not iterable
124 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
125 self.assertRaises(ValueError, permutations, 'abc', 32) # r is too big
126 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Raymond Hettingerd553d852008-03-04 04:17:08 +0000127 self.assertEqual(list(permutations(range(3), 2)),
128 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
129
130 def permutations1(iterable, r=None):
131 'Pure python version shown in the docs'
132 pool = tuple(iterable)
133 n = len(pool)
134 r = n if r is None else r
135 indices = range(n)
136 cycles = range(n-r+1, n+1)[::-1]
137 yield tuple(pool[i] for i in indices[:r])
138 while n:
139 for i in reversed(range(r)):
140 cycles[i] -= 1
141 if cycles[i] == 0:
142 indices[i:] = indices[i+1:] + indices[i:i+1]
143 cycles[i] = n - i
144 else:
145 j = cycles[i]
146 indices[i], indices[-j] = indices[-j], indices[i]
147 yield tuple(pool[i] for i in indices[:r])
148 break
149 else:
150 return
151
152 def permutations2(iterable, r=None):
153 'Pure python version shown in the docs'
154 pool = tuple(iterable)
155 n = len(pool)
156 r = n if r is None else r
157 for indices in product(range(n), repeat=r):
158 if len(set(indices)) == r:
159 yield tuple(pool[i] for i in indices)
160
161 for n in range(7):
162 values = [5*x-12 for x in range(n)]
163 for r in range(n+1):
164 result = list(permutations(values, r))
165 self.assertEqual(len(result), fact(n) / fact(n-r)) # right number of perms
166 self.assertEqual(len(result), len(set(result))) # no repeats
167 self.assertEqual(result, sorted(result)) # lexicographic order
168 for p in result:
169 self.assertEqual(len(p), r) # r-length permutations
170 self.assertEqual(len(set(p)), r) # no duplicate elements
171 self.assert_(all(e in values for e in p)) # elements taken from input iterable
172 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
173 self.assertEqual(result, list(permutations2(values, r))) # matches first pure python version
174 if r == n:
175 self.assertEqual(result, list(permutations(values, None))) # test r as None
176 self.assertEqual(result, list(permutations(values))) # test default r
177
178 # Test implementation detail: tuple re-use
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000179 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000180 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000181
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000182 def test_count(self):
183 self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
184 self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000185 self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000186 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
187 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000188 self.assertRaises(TypeError, count, 2, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000189 self.assertRaises(TypeError, count, 'a')
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000190 self.assertEqual(list(islice(count(maxsize-5), 10)), range(maxsize-5, maxsize+5))
191 self.assertEqual(list(islice(count(-maxsize-5), 10)), range(-maxsize-5, -maxsize+5))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000192 c = count(3)
193 self.assertEqual(repr(c), 'count(3)')
194 c.next()
195 self.assertEqual(repr(c), 'count(4)')
Jack Diederich36234e82006-09-21 17:50:26 +0000196 c = count(-9)
197 self.assertEqual(repr(c), 'count(-9)')
198 c.next()
199 self.assertEqual(c.next(), -8)
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000200 for i in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 10, sys.maxint-5, sys.maxint+5):
Raymond Hettinger3a0de082007-10-12 17:53:11 +0000201 # Test repr (ignoring the L in longs)
202 r1 = repr(count(i)).replace('L', '')
203 r2 = 'count(%r)'.__mod__(i).replace('L', '')
204 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000205
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000206 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000207 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000208 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000209 self.assertRaises(TypeError, cycle)
210 self.assertRaises(TypeError, cycle, 5)
211 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000212
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000213 def test_groupby(self):
214 # Check whether it accepts arguments correctly
215 self.assertEqual([], list(groupby([])))
216 self.assertEqual([], list(groupby([], key=id)))
217 self.assertRaises(TypeError, list, groupby('abc', []))
218 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000219 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000220
221 # Check normal input
222 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
223 (2,15,22), (3,16,23), (3,17,23)]
224 dup = []
225 for k, g in groupby(s, lambda r:r[0]):
226 for elem in g:
227 self.assertEqual(k, elem[0])
228 dup.append(elem)
229 self.assertEqual(s, dup)
230
231 # Check nested case
232 dup = []
233 for k, g in groupby(s, lambda r:r[0]):
234 for ik, ig in groupby(g, lambda r:r[2]):
235 for elem in ig:
236 self.assertEqual(k, elem[0])
237 self.assertEqual(ik, elem[2])
238 dup.append(elem)
239 self.assertEqual(s, dup)
240
241 # Check case where inner iterator is not used
242 keys = [k for k, g in groupby(s, lambda r:r[0])]
243 expectedkeys = set([r[0] for r in s])
244 self.assertEqual(set(keys), expectedkeys)
245 self.assertEqual(len(keys), len(expectedkeys))
246
247 # Exercise pipes and filters style
248 s = 'abracadabra'
249 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000250 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000251 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
252 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000253 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000254 self.assertEqual(r, ['a', 'b', 'r'])
255 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000256 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000257 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
258 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000259 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000260 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
261
262 # iter.next failure
263 class ExpectedError(Exception):
264 pass
265 def delayed_raise(n=0):
266 for i in range(n):
267 yield 'yo'
268 raise ExpectedError
269 def gulp(iterable, keyp=None, func=list):
270 return [func(g) for k, g in groupby(iterable, keyp)]
271
272 # iter.next failure on outer object
273 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
274 # iter.next failure on inner object
275 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
276
277 # __cmp__ failure
278 class DummyCmp:
279 def __cmp__(self, dst):
280 raise ExpectedError
281 s = [DummyCmp(), DummyCmp(), None]
282
283 # __cmp__ failure on outer object
284 self.assertRaises(ExpectedError, gulp, s, func=id)
285 # __cmp__ failure on inner object
286 self.assertRaises(ExpectedError, gulp, s)
287
288 # keyfunc failure
289 def keyfunc(obj):
290 if keyfunc.skip > 0:
291 keyfunc.skip -= 1
292 return obj
293 else:
294 raise ExpectedError
295
296 # keyfunc failure on outer object
297 keyfunc.skip = 0
298 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
299 keyfunc.skip = 1
300 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
301
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000302 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000303 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000304 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000305 self.assertEqual(list(ifilter(bool, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000306 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000307 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000308 self.assertRaises(TypeError, ifilter, lambda x:x)
309 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000310 self.assertRaises(TypeError, ifilter, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000311 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000312
313 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000314 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
315 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000316 self.assertEqual(list(ifilterfalse(bool, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000317 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000318 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000319 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
320 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000321 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000322 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000323
324 def test_izip(self):
325 ans = [(x,y) for x, y in izip('abc',count())]
326 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000327 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
328 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
Raymond Hettinger02420702003-06-29 20:36:23 +0000329 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000330 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000331 self.assertEqual(list(izip()), zip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000332 self.assertRaises(TypeError, izip, 3)
333 self.assertRaises(TypeError, izip, range(3), 3)
334 # Check tuple re-use (implementation detail)
335 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
336 zip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000337 self.assertEqual([pair for pair in izip('abc', 'def')],
338 zip('abc', 'def'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000339 ids = map(id, izip('abc', 'def'))
340 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000341 ids = map(id, list(izip('abc', 'def')))
342 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000343
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000344 def test_iziplongest(self):
345 for args in [
346 ['abc', range(6)],
347 [range(6), 'abc'],
348 [range(1000), range(2000,2100), range(3000,3050)],
349 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
350 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
351 ]:
352 target = map(None, *args)
353 self.assertEqual(list(izip_longest(*args)), target)
354 self.assertEqual(list(izip_longest(*args, **{})), target)
355 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
356 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Tim Petersea5962f2007-03-12 18:07:52 +0000357
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000358 self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
359
360 self.assertEqual(list(izip_longest()), zip())
361 self.assertEqual(list(izip_longest([])), zip([]))
362 self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
Tim Petersea5962f2007-03-12 18:07:52 +0000363
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000364 self.assertEqual(list(izip_longest('abc', 'defg', **{})), map(None, 'abc', 'defg')) # empty keyword dict
365 self.assertRaises(TypeError, izip_longest, 3)
366 self.assertRaises(TypeError, izip_longest, range(3), 3)
367
368 for stmt in [
369 "izip_longest('abc', fv=1)",
Tim Petersea5962f2007-03-12 18:07:52 +0000370 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000371 ]:
372 try:
373 eval(stmt, globals(), locals())
374 except TypeError:
375 pass
376 else:
377 self.fail('Did not raise Type in: ' + stmt)
Tim Petersea5962f2007-03-12 18:07:52 +0000378
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000379 # Check tuple re-use (implementation detail)
380 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
381 zip('abc', 'def'))
382 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
383 zip('abc', 'def'))
384 ids = map(id, izip_longest('abc', 'def'))
385 self.assertEqual(min(ids), max(ids))
386 ids = map(id, list(izip_longest('abc', 'def')))
387 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
388
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000389 def test_product(self):
390 for args, result in [
Raymond Hettingerd553d852008-03-04 04:17:08 +0000391 ([], [()]), # zero iterables
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000392 (['ab'], [('a',), ('b',)]), # one iterable
393 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
394 ([range(0), range(2), range(3)], []), # first iterable with zero length
395 ([range(2), range(0), range(3)], []), # middle iterable with zero length
396 ([range(2), range(3), range(0)], []), # last iterable with zero length
397 ]:
398 self.assertEqual(list(product(*args)), result)
Raymond Hettinger08ff6822008-02-29 02:21:48 +0000399 for r in range(4):
400 self.assertEqual(list(product(*(args*r))),
401 list(product(*args, **dict(repeat=r))))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000402 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
403 self.assertRaises(TypeError, product, range(6), None)
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000404
405 def product2(*args, **kwds):
406 'Pure python version used in docs'
407 pools = map(tuple, args) * kwds.get('repeat', 1)
408 result = [[]]
409 for pool in pools:
410 result = [x+[y] for x in result for y in pool]
411 for prod in result:
412 yield tuple(prod)
413
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000414 argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
415 set('abcdefg'), range(11), tuple(range(13))]
416 for i in range(100):
417 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Raymond Hettingerd553d852008-03-04 04:17:08 +0000418 expected_len = prod(map(len, args))
419 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000420 self.assertEqual(list(product(*args)), list(product2(*args)))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000421 args = map(iter, args)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000422 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000423
Raymond Hettinger73d79632008-02-23 02:20:41 +0000424 # Test implementation detail: tuple re-use
425 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
426 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000427
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000428 def test_repeat(self):
429 self.assertEqual(zip(xrange(3),repeat('a')),
430 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000431 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000432 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000433 self.assertEqual(list(repeat('a', 0)), [])
434 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000435 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000436 self.assertRaises(TypeError, repeat, None, 3, 4)
437 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000438 r = repeat(1+0j)
439 self.assertEqual(repr(r), 'repeat((1+0j))')
440 r = repeat(1+0j, 5)
441 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
442 list(r)
443 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000444
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000445 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000446 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
447 [0**1, 1**2, 2**3])
448 self.assertEqual(list(imap(None, 'abc', range(5))),
449 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000450 self.assertEqual(list(imap(None, 'abc', count())),
451 [('a',0),('b',1),('c',2)])
452 self.assertEqual(take(2,imap(None, 'abc', count())),
453 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000454 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000455 self.assertRaises(TypeError, imap)
456 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000457 self.assertRaises(TypeError, imap(10, range(5)).next)
458 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
459 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000460
461 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000462 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
463 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000464 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
465 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000466 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger47317092008-01-17 03:02:14 +0000467 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
468 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000469 self.assertRaises(TypeError, starmap)
470 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
471 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
472 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
473 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000474
475 def test_islice(self):
476 for args in [ # islice(args) should agree with range(args)
477 (10, 20, 3),
478 (10, 3, 20),
479 (10, 20),
480 (10, 3),
481 (20,)
482 ]:
483 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
484
485 for args, tgtargs in [ # Stop when seqn is exhausted
486 ((10, 110, 3), ((10, 100, 3))),
487 ((10, 110), ((10, 100))),
488 ((110,), (100,))
489 ]:
490 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
491
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000492 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000493 self.assertEqual(list(islice(xrange(10), None)), range(10))
Raymond Hettingerb2594052004-12-05 09:25:51 +0000494 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
495 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000496 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
497 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
498
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000499 # Test number of items consumed SF #1171417
500 it = iter(range(10))
501 self.assertEqual(list(islice(it, 3)), range(3))
502 self.assertEqual(list(it), range(3, 10))
503
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000504 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000505 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000506 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
507 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
508 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
509 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
510 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000511 self.assertRaises(ValueError, islice, xrange(10), 'a')
512 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
513 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
514 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
515 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +0000516 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000517
518 def test_takewhile(self):
519 data = [1, 3, 5, 20, 2, 4, 6, 8]
520 underten = lambda x: x<10
521 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000522 self.assertEqual(list(takewhile(underten, [])), [])
523 self.assertRaises(TypeError, takewhile)
524 self.assertRaises(TypeError, takewhile, operator.pow)
525 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
526 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
527 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000528 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
529 self.assertEqual(list(t), [1, 1, 1])
530 self.assertRaises(StopIteration, t.next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000531
532 def test_dropwhile(self):
533 data = [1, 3, 5, 20, 2, 4, 6, 8]
534 underten = lambda x: x<10
535 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000536 self.assertEqual(list(dropwhile(underten, [])), [])
537 self.assertRaises(TypeError, dropwhile)
538 self.assertRaises(TypeError, dropwhile, operator.pow)
539 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
540 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
541 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000542
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000543 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000544 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000545 def irange(n):
546 for i in xrange(n):
547 yield i
548
549 a, b = tee([]) # test empty iterator
550 self.assertEqual(list(a), [])
551 self.assertEqual(list(b), [])
552
553 a, b = tee(irange(n)) # test 100% interleaved
554 self.assertEqual(zip(a,b), zip(range(n),range(n)))
555
556 a, b = tee(irange(n)) # test 0% interleaved
557 self.assertEqual(list(a), range(n))
558 self.assertEqual(list(b), range(n))
559
560 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000561 for i in xrange(100):
562 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000563 del a
564 self.assertEqual(list(b), range(n))
565
566 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000567 for i in xrange(100):
568 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000569 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000570 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000571
572 for j in xrange(5): # test randomly interleaved
573 order = [0]*n + [1]*n
574 random.shuffle(order)
575 lists = ([], [])
576 its = tee(irange(n))
577 for i in order:
578 value = its[i].next()
579 lists[i].append(value)
580 self.assertEqual(lists[0], range(n))
581 self.assertEqual(lists[1], range(n))
582
Raymond Hettingerad983e72003-11-12 14:32:26 +0000583 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000584 self.assertRaises(TypeError, tee)
585 self.assertRaises(TypeError, tee, 3)
586 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000587 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000588
Raymond Hettingerad983e72003-11-12 14:32:26 +0000589 # tee object should be instantiable
590 a, b = tee('abc')
591 c = type(a)('def')
592 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000593
Raymond Hettingerad983e72003-11-12 14:32:26 +0000594 # test long-lagged and multi-way split
595 a, b, c = tee(xrange(2000), 3)
596 for i in xrange(100):
597 self.assertEqual(a.next(), i)
598 self.assertEqual(list(b), range(2000))
599 self.assertEqual([c.next(), c.next()], range(2))
600 self.assertEqual(list(a), range(100,2000))
601 self.assertEqual(list(c), range(2,2000))
602
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000603 # test values of n
604 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Neal Norwitz69e88972006-09-02 02:58:13 +0000605 self.assertRaises(ValueError, tee, [], -1)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000606 for n in xrange(5):
607 result = tee('abc', n)
608 self.assertEqual(type(result), tuple)
609 self.assertEqual(len(result), n)
610 self.assertEqual(map(list, result), [list('abc')]*n)
611
Raymond Hettingerad983e72003-11-12 14:32:26 +0000612 # tee pass-through to copyable iterator
613 a, b = tee('abc')
614 c, d = tee(a)
615 self.assert_(a is c)
616
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000617 # test tee_new
618 t1, t2 = tee('abc')
619 tnew = type(t1)
620 self.assertRaises(TypeError, tnew)
621 self.assertRaises(TypeError, tnew, 10)
622 t3 = tnew(t1)
623 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000624
Raymond Hettingera9f60922004-10-17 16:40:14 +0000625 # test that tee objects are weak referencable
626 a, b = tee(xrange(10))
627 p = proxy(a)
628 self.assertEqual(getattr(p, '__class__'), type(b))
629 del a
630 self.assertRaises(ReferenceError, getattr, p, '__class__')
631
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000632 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000633 self.assertRaises(StopIteration, izip().next)
634
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000635 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000636 self.assertRaises(StopIteration, f([]).next)
637 self.assertRaises(StopIteration, f(StopNow()).next)
638
639 self.assertRaises(StopIteration, islice([], None).next)
640 self.assertRaises(StopIteration, islice(StopNow(), None).next)
641
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000642 p, q = tee([])
643 self.assertRaises(StopIteration, p.next)
644 self.assertRaises(StopIteration, q.next)
645 p, q = tee(StopNow())
646 self.assertRaises(StopIteration, p.next)
647 self.assertRaises(StopIteration, q.next)
648
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000649 self.assertRaises(StopIteration, repeat(None, 0).next)
650
651 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
652 self.assertRaises(StopIteration, f(lambda x:x, []).next)
653 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
654
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000655class TestGC(unittest.TestCase):
656
657 def makecycle(self, iterator, container):
658 container.append(iterator)
659 iterator.next()
660 del container, iterator
661
662 def test_chain(self):
663 a = []
664 self.makecycle(chain(a), a)
665
666 def test_cycle(self):
667 a = []
668 self.makecycle(cycle([a]*2), a)
669
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000670 def test_dropwhile(self):
671 a = []
672 self.makecycle(dropwhile(bool, [0, a, a]), a)
673
674 def test_groupby(self):
675 a = []
676 self.makecycle(groupby([a]*2, lambda x:x), a)
677
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000678 def test_ifilter(self):
679 a = []
680 self.makecycle(ifilter(lambda x:True, [a]*2), a)
681
682 def test_ifilterfalse(self):
683 a = []
684 self.makecycle(ifilterfalse(lambda x:False, a), a)
685
686 def test_izip(self):
687 a = []
688 self.makecycle(izip([a]*2, [a]*3), a)
689
690 def test_imap(self):
691 a = []
692 self.makecycle(imap(lambda x:x, [a]*2), a)
693
694 def test_islice(self):
695 a = []
696 self.makecycle(islice([a]*2, None), a)
697
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000698 def test_repeat(self):
699 a = []
700 self.makecycle(repeat(a), a)
701
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000702 def test_starmap(self):
703 a = []
704 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
705
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000706 def test_takewhile(self):
707 a = []
708 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
709
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000710def R(seqn):
711 'Regular generator'
712 for i in seqn:
713 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000714
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000715class G:
716 'Sequence using __getitem__'
717 def __init__(self, seqn):
718 self.seqn = seqn
719 def __getitem__(self, i):
720 return self.seqn[i]
721
722class I:
723 'Sequence using iterator protocol'
724 def __init__(self, seqn):
725 self.seqn = seqn
726 self.i = 0
727 def __iter__(self):
728 return self
729 def next(self):
730 if self.i >= len(self.seqn): raise StopIteration
731 v = self.seqn[self.i]
732 self.i += 1
733 return v
734
735class Ig:
736 'Sequence using iterator protocol defined with a generator'
737 def __init__(self, seqn):
738 self.seqn = seqn
739 self.i = 0
740 def __iter__(self):
741 for val in self.seqn:
742 yield val
743
744class X:
745 'Missing __getitem__ and __iter__'
746 def __init__(self, seqn):
747 self.seqn = seqn
748 self.i = 0
749 def next(self):
750 if self.i >= len(self.seqn): raise StopIteration
751 v = self.seqn[self.i]
752 self.i += 1
753 return v
754
755class N:
756 'Iterator missing next()'
757 def __init__(self, seqn):
758 self.seqn = seqn
759 self.i = 0
760 def __iter__(self):
761 return self
762
763class E:
764 'Test propagation of exceptions'
765 def __init__(self, seqn):
766 self.seqn = seqn
767 self.i = 0
768 def __iter__(self):
769 return self
770 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +0000771 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000772
773class S:
774 'Test immediate stop'
775 def __init__(self, seqn):
776 pass
777 def __iter__(self):
778 return self
779 def next(self):
780 raise StopIteration
781
782def L(seqn):
783 'Test multiple tiers of iterators'
784 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
785
786
787class TestVariousIteratorArgs(unittest.TestCase):
788
789 def test_chain(self):
790 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
791 for g in (G, I, Ig, S, L, R):
792 self.assertEqual(list(chain(g(s))), list(g(s)))
793 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Raymond Hettinger05bf6332008-02-28 22:30:42 +0000794 self.assertRaises(TypeError, list, chain(X(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000795 self.assertRaises(TypeError, list, chain(N(s)))
796 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
797
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000798 def test_product(self):
799 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
800 self.assertRaises(TypeError, product, X(s))
801 self.assertRaises(TypeError, product, N(s))
802 self.assertRaises(ZeroDivisionError, product, E(s))
803
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000804 def test_cycle(self):
805 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
806 for g in (G, I, Ig, S, L, R):
807 tgtlen = len(s) * 3
808 expected = list(g(s))*3
809 actual = list(islice(cycle(g(s)), tgtlen))
810 self.assertEqual(actual, expected)
811 self.assertRaises(TypeError, cycle, X(s))
812 self.assertRaises(TypeError, list, cycle(N(s)))
813 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
814
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000815 def test_groupby(self):
816 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
817 for g in (G, I, Ig, S, L, R):
818 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
819 self.assertRaises(TypeError, groupby, X(s))
820 self.assertRaises(TypeError, list, groupby(N(s)))
821 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
822
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000823 def test_ifilter(self):
824 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
825 for g in (G, I, Ig, S, L, R):
826 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
827 self.assertRaises(TypeError, ifilter, isEven, X(s))
828 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
829 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
830
831 def test_ifilterfalse(self):
832 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
833 for g in (G, I, Ig, S, L, R):
834 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
835 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
836 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
837 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
838
839 def test_izip(self):
840 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
841 for g in (G, I, Ig, S, L, R):
842 self.assertEqual(list(izip(g(s))), zip(g(s)))
843 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
844 self.assertRaises(TypeError, izip, X(s))
845 self.assertRaises(TypeError, list, izip(N(s)))
846 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
847
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000848 def test_iziplongest(self):
849 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
850 for g in (G, I, Ig, S, L, R):
851 self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
852 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
853 self.assertRaises(TypeError, izip_longest, X(s))
854 self.assertRaises(TypeError, list, izip_longest(N(s)))
855 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
856
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000857 def test_imap(self):
858 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
859 for g in (G, I, Ig, S, L, R):
860 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
861 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
862 self.assertRaises(TypeError, imap, onearg, X(s))
863 self.assertRaises(TypeError, list, imap(onearg, N(s)))
864 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
865
866 def test_islice(self):
867 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
868 for g in (G, I, Ig, S, L, R):
869 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
870 self.assertRaises(TypeError, islice, X(s), 10)
871 self.assertRaises(TypeError, list, islice(N(s), 10))
872 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
873
874 def test_starmap(self):
875 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
876 for g in (G, I, Ig, S, L, R):
877 ss = zip(s, s)
878 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
879 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
880 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
881 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
882
883 def test_takewhile(self):
884 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
885 for g in (G, I, Ig, S, L, R):
886 tgt = []
887 for elem in g(s):
888 if not isEven(elem): break
889 tgt.append(elem)
890 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
891 self.assertRaises(TypeError, takewhile, isEven, X(s))
892 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
893 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
894
895 def test_dropwhile(self):
896 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
897 for g in (G, I, Ig, S, L, R):
898 tgt = []
899 for elem in g(s):
900 if not tgt and isOdd(elem): continue
901 tgt.append(elem)
902 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
903 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
904 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
905 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
906
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000907 def test_tee(self):
908 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
909 for g in (G, I, Ig, S, L, R):
910 it1, it2 = tee(g(s))
911 self.assertEqual(list(it1), list(g(s)))
912 self.assertEqual(list(it2), list(g(s)))
913 self.assertRaises(TypeError, tee, X(s))
914 self.assertRaises(TypeError, list, tee(N(s))[0])
915 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
916
Raymond Hettinger5cab2e32004-02-10 09:25:40 +0000917class LengthTransparency(unittest.TestCase):
918
919 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +0000920 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +0000921 self.assertEqual(len(repeat(None, 50)), 50)
922 self.assertRaises(TypeError, len, repeat(None))
923
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000924class RegressionTests(unittest.TestCase):
925
926 def test_sf_793826(self):
927 # Fix Armin Rigo's successful efforts to wreak havoc
928
929 def mutatingtuple(tuple1, f, tuple2):
930 # this builds a tuple t which is a copy of tuple1,
931 # then calls f(t), then mutates t to be equal to tuple2
932 # (needs len(tuple1) == len(tuple2)).
933 def g(value, first=[1]):
934 if first:
935 del first[:]
936 f(z.next())
937 return value
938 items = list(tuple2)
939 items[1:1] = list(tuple1)
940 gen = imap(g, items)
941 z = izip(*[gen]*len(tuple1))
942 z.next()
943
944 def f(t):
945 global T
946 T = t
947 first[:] = list(T)
948
949 first = []
950 mutatingtuple((1,2,3), f, (4,5,6))
951 second = list(T)
952 self.assertEqual(first, second)
953
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000954
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000955 def test_sf_950057(self):
956 # Make sure that chain() and cycle() catch exceptions immediately
957 # rather than when shifting between input sources
958
959 def gen1():
960 hist.append(0)
961 yield 1
962 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +0000963 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000964 hist.append(2)
965
966 def gen2(x):
967 hist.append(3)
968 yield 2
969 hist.append(4)
970 if x:
971 raise StopIteration
972
973 hist = []
974 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
975 self.assertEqual(hist, [0,1])
976
977 hist = []
978 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
979 self.assertEqual(hist, [0,1])
980
981 hist = []
982 self.assertRaises(AssertionError, list, cycle(gen1()))
983 self.assertEqual(hist, [0,1])
984
Georg Brandlb84c1372007-01-21 10:28:43 +0000985class SubclassWithKwargsTest(unittest.TestCase):
986 def test_keywords_in_subclass(self):
987 # count is not subclassable...
988 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
989 starmap, islice, takewhile, dropwhile, cycle):
990 class Subclass(cls):
991 def __init__(self, newarg=None, *args):
992 cls.__init__(self, *args)
993 try:
994 Subclass(newarg=1)
995 except TypeError, err:
996 # we expect type errors because of wrong argument count
997 self.failIf("does not take keyword arguments" in err.args[0])
998
999
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001000libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001001
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001002
1003>>> amounts = [120.15, 764.05, 823.14]
1004>>> for checknum, amount in izip(count(1200), amounts):
1005... print 'Check %d is for $%.2f' % (checknum, amount)
1006...
1007Check 1200 is for $120.15
1008Check 1201 is for $764.05
1009Check 1202 is for $823.14
1010
1011>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001012>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1013... print cube
1014...
10151
10168
101727
1018
1019>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001020>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001021... print name.title()
1022...
1023Alex
1024Laura
1025Martin
1026Walter
1027Samuele
1028
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001029>>> from operator import itemgetter
1030>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Armin Rigoa3f09272006-05-28 19:13:17 +00001031>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001032>>> for k, g in groupby(di, itemgetter(1)):
1033... print k, map(itemgetter(0), g)
1034...
10351 ['a', 'c', 'e']
10362 ['b', 'd', 'f']
10373 ['g']
1038
Raymond Hettinger734fb572004-01-20 20:04:40 +00001039# Find runs of consecutive numbers using groupby. The key to the solution
1040# is differencing with a range so that consecutive numbers all appear in
1041# same group.
1042>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
1043>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
1044... print map(operator.itemgetter(1), g)
1045...
1046[1]
1047[4, 5, 6]
1048[10]
1049[15, 16, 17, 18]
1050[22]
1051[25, 26, 27, 28]
1052
Raymond Hettingera098b332003-09-08 23:58:40 +00001053>>> def take(n, seq):
1054... return list(islice(seq, n))
1055
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001056>>> def enumerate(iterable):
1057... return izip(count(), iterable)
1058
1059>>> def tabulate(function):
1060... "Return function(0), function(1), ..."
1061... return imap(function, count())
1062
1063>>> def iteritems(mapping):
1064... return izip(mapping.iterkeys(), mapping.itervalues())
1065
1066>>> def nth(iterable, n):
1067... "Returns the nth item"
Raymond Hettinger60eca932003-02-09 06:40:58 +00001068... return list(islice(iterable, n, n+1))
1069
Raymond Hettingerf77d0332005-03-11 22:17:30 +00001070>>> def all(seq, pred=None):
1071... "Returns True if pred(x) is true for every element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +00001072... for elem in ifilterfalse(pred, seq):
1073... return False
1074... return True
Raymond Hettinger60eca932003-02-09 06:40:58 +00001075
Raymond Hettingerf77d0332005-03-11 22:17:30 +00001076>>> def any(seq, pred=None):
1077... "Returns True if pred(x) is true for at least one element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +00001078... for elem in ifilter(pred, seq):
1079... return True
1080... return False
Raymond Hettinger60eca932003-02-09 06:40:58 +00001081
Raymond Hettingerf77d0332005-03-11 22:17:30 +00001082>>> def no(seq, pred=None):
1083... "Returns True if pred(x) is false for every element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +00001084... for elem in ifilter(pred, seq):
1085... return False
1086... return True
Raymond Hettinger77fe69b2003-08-08 04:33:19 +00001087
Raymond Hettingerf77d0332005-03-11 22:17:30 +00001088>>> def quantify(seq, pred=None):
1089... "Count how many times the predicate is true in the sequence"
Raymond Hettinger77fe69b2003-08-08 04:33:19 +00001090... return sum(imap(pred, seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001091
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001092>>> def padnone(seq):
1093... "Returns the sequence elements and then returns None indefinitely"
1094... return chain(seq, repeat(None))
1095
1096>>> def ncycles(seq, n):
1097... "Returns the sequence elements n times"
1098... return chain(*repeat(seq, n))
1099
1100>>> def dotproduct(vec1, vec2):
1101... return sum(imap(operator.mul, vec1, vec2))
1102
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001103>>> def flatten(listOfLists):
1104... return list(chain(*listOfLists))
1105
1106>>> def repeatfunc(func, times=None, *args):
1107... "Repeat calls to func with specified arguments."
1108... " Example: repeatfunc(random.random)"
1109... if times is None:
1110... return starmap(func, repeat(args))
1111... else:
1112... return starmap(func, repeat(args, times))
1113
Raymond Hettingerd591f662003-10-26 15:34:50 +00001114>>> def pairwise(iterable):
1115... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1116... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001117... try:
1118... b.next()
1119... except StopIteration:
1120... pass
1121... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001122
1123This is not part of the examples but it tests to make sure the definitions
1124perform as purported.
1125
Raymond Hettingera098b332003-09-08 23:58:40 +00001126>>> take(10, count())
1127[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1128
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001129>>> list(enumerate('abc'))
1130[(0, 'a'), (1, 'b'), (2, 'c')]
1131
1132>>> list(islice(tabulate(lambda x: 2*x), 4))
1133[0, 2, 4, 6]
1134
1135>>> nth('abcde', 3)
1136['d']
1137
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001138>>> all([2, 4, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001139True
1140
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001141>>> all([2, 3, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001142False
1143
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001144>>> any([2, 4, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001145True
1146
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001147>>> any([1, 3, 5, 9], lambda x: x%2==0,)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001148False
1149
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001150>>> no([1, 3, 5, 9], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001151True
1152
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001153>>> no([1, 2, 5, 9], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001154False
1155
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001156>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000115750
1158
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001159>>> a = [[1, 2, 3], [4, 5, 6]]
1160>>> flatten(a)
1161[1, 2, 3, 4, 5, 6]
1162
1163>>> list(repeatfunc(pow, 5, 2, 3))
1164[8, 8, 8, 8, 8]
1165
1166>>> import random
1167>>> take(5, imap(int, repeatfunc(random.random)))
1168[0, 0, 0, 0, 0]
1169
Raymond Hettingerd591f662003-10-26 15:34:50 +00001170>>> list(pairwise('abcd'))
1171[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001172
Raymond Hettingerd591f662003-10-26 15:34:50 +00001173>>> list(pairwise([]))
1174[]
1175
1176>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001177[]
1178
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001179>>> list(islice(padnone('abc'), 0, 6))
1180['a', 'b', 'c', None, None, None]
1181
1182>>> list(ncycles('abc', 3))
1183['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1184
1185>>> dotproduct([1,2,3], [4,5,6])
118632
1187
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001188"""
1189
1190__test__ = {'libreftest' : libreftest}
1191
1192def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001193 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Georg Brandlb84c1372007-01-21 10:28:43 +00001194 RegressionTests, LengthTransparency,
1195 SubclassWithKwargsTest)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001196 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001197
1198 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001199 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001200 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001201 counts = [None] * 5
1202 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001203 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001204 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001205 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001206 print counts
1207
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001208 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001209 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001210
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001211if __name__ == "__main__":
1212 test_main(verbose=True)