blob: 8d8e0596579bd56d0da586670aa1b79fda71e8cc [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):
Raymond Hettingerad47fa12008-03-06 20:52:01 +000052
53 def chain2(*iterables):
54 'Pure python version in the docs'
55 for it in iterables:
56 for element in it:
57 yield element
58
59 for c in (chain, chain2):
60 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
61 self.assertEqual(list(c('abc')), list('abc'))
62 self.assertEqual(list(c('')), [])
63 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
64 self.assertRaises(TypeError, list,c(2, 3))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000065
Raymond Hettingerb4cbc982008-02-28 22:46:41 +000066 def test_chain_from_iterable(self):
67 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
68 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
69 self.assertEqual(list(chain.from_iterable([''])), [])
70 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
71 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
72
Raymond Hettinger93e804d2008-02-26 23:40:50 +000073 def test_combinations(self):
Raymond Hettinger5b913e32009-01-08 06:39:04 +000074 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Raymond Hettinger93e804d2008-02-26 23:40:50 +000075 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Raymond Hettingerd553d852008-03-04 04:17:08 +000076 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Raymond Hettinger93e804d2008-02-26 23:40:50 +000077 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Raymond Hettinger5b913e32009-01-08 06:39:04 +000078 self.assertEqual(list(combinations('abc', 32)), []) # r > n
Raymond Hettinger93e804d2008-02-26 23:40:50 +000079 self.assertEqual(list(combinations(range(4), 3)),
80 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
Raymond Hettingerd553d852008-03-04 04:17:08 +000081
82 def combinations1(iterable, r):
83 'Pure python version shown in the docs'
84 pool = tuple(iterable)
85 n = len(pool)
Raymond Hettinger5b913e32009-01-08 06:39:04 +000086 if r > n:
87 return
Raymond Hettingerd553d852008-03-04 04:17:08 +000088 indices = range(r)
89 yield tuple(pool[i] for i in indices)
90 while 1:
91 for i in reversed(range(r)):
92 if indices[i] != i + n - r:
93 break
94 else:
95 return
96 indices[i] += 1
97 for j in range(i+1, r):
98 indices[j] = indices[j-1] + 1
99 yield tuple(pool[i] for i in indices)
100
101 def combinations2(iterable, r):
102 'Pure python version shown in the docs'
103 pool = tuple(iterable)
104 n = len(pool)
105 for indices in permutations(range(n), r):
106 if sorted(indices) == list(indices):
107 yield tuple(pool[i] for i in indices)
108
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000109 def combinations3(iterable, r):
110 'Pure python version from cwr()'
111 pool = tuple(iterable)
112 n = len(pool)
113 for indices in combinations_with_replacement(range(n), r):
114 if len(set(indices)) == r:
115 yield tuple(pool[i] for i in indices)
116
Raymond Hettingerd553d852008-03-04 04:17:08 +0000117 for n in range(7):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000118 values = [5*x-12 for x in range(n)]
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000119 for r in range(n+2):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000120 result = list(combinations(values, r))
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000121 self.assertEqual(len(result), 0 if r>n else fact(n) / fact(r) / fact(n-r)) # right number of combs
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000122 self.assertEqual(len(result), len(set(result))) # no repeats
123 self.assertEqual(result, sorted(result)) # lexicographic order
124 for c in result:
125 self.assertEqual(len(c), r) # r-length combinations
126 self.assertEqual(len(set(c)), r) # no duplicate elements
127 self.assertEqual(list(c), sorted(c)) # keep original ordering
128 self.assert_(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000129 self.assertEqual(list(c),
130 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Raymond Hettingerd553d852008-03-04 04:17:08 +0000131 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000132 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000133 self.assertEqual(result, list(combinations3(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000134
135 # Test implementation detail: tuple re-use
136 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
137 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
138
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000139 def test_combinations_with_replacement(self):
140 cwr = combinations_with_replacement
141 self.assertRaises(TypeError, cwr, 'abc') # missing r argument
142 self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
143 self.assertRaises(TypeError, cwr, None) # pool is not iterable
144 self.assertRaises(ValueError, cwr, 'abc', -2) # r is negative
145 self.assertEqual(list(cwr('ABC', 2)),
146 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
147
148 def cwr1(iterable, r):
149 'Pure python version shown in the docs'
150 # number items returned: (n+r-1)! / r! / (n-1)! when n>0
151 pool = tuple(iterable)
152 n = len(pool)
153 if not n and r:
154 return
155 indices = [0] * r
156 yield tuple(pool[i] for i in indices)
157 while 1:
158 for i in reversed(range(r)):
159 if indices[i] != n - 1:
160 break
161 else:
162 return
163 indices[i:] = [indices[i] + 1] * (r - i)
164 yield tuple(pool[i] for i in indices)
165
166 def cwr2(iterable, r):
167 'Pure python version shown in the docs'
168 pool = tuple(iterable)
169 n = len(pool)
170 for indices in product(range(n), repeat=r):
171 if sorted(indices) == list(indices):
172 yield tuple(pool[i] for i in indices)
173
174 def numcombs(n, r):
175 if not n:
176 return 0 if r else 1
177 return fact(n+r-1) / fact(r)/ fact(n-1)
178
179 for n in range(7):
180 values = [5*x-12 for x in range(n)]
181 for r in range(n+2):
182 result = list(cwr(values, r))
183
184 self.assertEqual(len(result), numcombs(n, r)) # right number of combs
185 self.assertEqual(len(result), len(set(result))) # no repeats
186 self.assertEqual(result, sorted(result)) # lexicographic order
187
188 regular_combs = list(combinations(values, r)) # compare to combs without replacement
189 if n == 0 or r <= 1:
190 self.assertEquals(result, regular_combs) # cases that should be identical
191 else:
192 self.assert_(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
193
194 for c in result:
195 self.assertEqual(len(c), r) # r-length combinations
196 noruns = [k for k,v in groupby(c)] # combo without consecutive repeats
197 self.assertEqual(len(noruns), len(set(noruns))) # no repeats other than consecutive
198 self.assertEqual(list(c), sorted(c)) # keep original ordering
199 self.assert_(all(e in values for e in c)) # elements taken from input iterable
200 self.assertEqual(noruns,
201 [e for e in values if e in c]) # comb is a subsequence of the input iterable
202 self.assertEqual(result, list(cwr1(values, r))) # matches first pure python version
203 self.assertEqual(result, list(cwr2(values, r))) # matches second pure python version
204
205 # Test implementation detail: tuple re-use
206 self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
207 self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
208
Raymond Hettingerd553d852008-03-04 04:17:08 +0000209 def test_permutations(self):
210 self.assertRaises(TypeError, permutations) # too few arguments
211 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000212 self.assertRaises(TypeError, permutations, None) # pool is not iterable
213 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000214 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000215 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Raymond Hettingerd553d852008-03-04 04:17:08 +0000216 self.assertEqual(list(permutations(range(3), 2)),
217 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
218
219 def permutations1(iterable, r=None):
220 'Pure python version shown in the docs'
221 pool = tuple(iterable)
222 n = len(pool)
223 r = n if r is None else r
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000224 if r > n:
225 return
Raymond Hettingerd553d852008-03-04 04:17:08 +0000226 indices = range(n)
Raymond Hettingere70bb8d2008-03-23 00:55:46 +0000227 cycles = range(n, n-r, -1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000228 yield tuple(pool[i] for i in indices[:r])
229 while n:
230 for i in reversed(range(r)):
231 cycles[i] -= 1
232 if cycles[i] == 0:
233 indices[i:] = indices[i+1:] + indices[i:i+1]
234 cycles[i] = n - i
235 else:
236 j = cycles[i]
237 indices[i], indices[-j] = indices[-j], indices[i]
238 yield tuple(pool[i] for i in indices[:r])
239 break
240 else:
241 return
242
243 def permutations2(iterable, r=None):
244 'Pure python version shown in the docs'
245 pool = tuple(iterable)
246 n = len(pool)
247 r = n if r is None else r
248 for indices in product(range(n), repeat=r):
249 if len(set(indices)) == r:
250 yield tuple(pool[i] for i in indices)
251
252 for n in range(7):
253 values = [5*x-12 for x in range(n)]
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000254 for r in range(n+2):
Raymond Hettingerd553d852008-03-04 04:17:08 +0000255 result = list(permutations(values, r))
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000256 self.assertEqual(len(result), 0 if r>n else fact(n) / fact(n-r)) # right number of perms
Raymond Hettingerd553d852008-03-04 04:17:08 +0000257 self.assertEqual(len(result), len(set(result))) # no repeats
258 self.assertEqual(result, sorted(result)) # lexicographic order
259 for p in result:
260 self.assertEqual(len(p), r) # r-length permutations
261 self.assertEqual(len(set(p)), r) # no duplicate elements
262 self.assert_(all(e in values for e in p)) # elements taken from input iterable
263 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000264 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000265 if r == n:
266 self.assertEqual(result, list(permutations(values, None))) # test r as None
267 self.assertEqual(result, list(permutations(values))) # test default r
268
269 # Test implementation detail: tuple re-use
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000270 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000271 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000272
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000273 def test_combinatorics(self):
274 # Test relationships between product(), permutations(),
275 # combinations() and combinations_with_replacement().
276
277 s = 'ABCDE'
278 for r in range(8):
279 prod = list(product(s, repeat=r))
280 cwr = list(combinations_with_replacement(s, r))
281 perm = list(permutations(s, r))
282 comb = list(combinations(s, r))
283
284 self.assertEquals(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
285 self.assertEquals(perm, [t for t in prod if len(set(t))==r]) # perm: prods with no dups
286 self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
287 self.assertEqual(comb, [t for t in cwr if len(set(t))==r]) # comb: cwrs without dups
288 self.assertEqual(set(comb), set(cwr) & set(perm)) # comb: both a cwr and a perm
289
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000290 def test_compress(self):
291 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
292 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
293 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
294 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
295 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
296 n = 10000
297 data = chain.from_iterable(repeat(range(6), n))
298 selectors = chain.from_iterable(repeat((0, 1)))
299 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
300 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
301 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
302 self.assertRaises(TypeError, compress, range(6)) # too few args
303 self.assertRaises(TypeError, compress, range(6), None) # too many args
304
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000305 def test_count(self):
306 self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
307 self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000308 self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000309 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
310 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000311 self.assertRaises(TypeError, count, 2, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000312 self.assertRaises(TypeError, count, 'a')
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000313 self.assertEqual(list(islice(count(maxsize-5), 10)), range(maxsize-5, maxsize+5))
314 self.assertEqual(list(islice(count(-maxsize-5), 10)), range(-maxsize-5, -maxsize+5))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000315 c = count(3)
316 self.assertEqual(repr(c), 'count(3)')
317 c.next()
318 self.assertEqual(repr(c), 'count(4)')
Jack Diederich36234e82006-09-21 17:50:26 +0000319 c = count(-9)
320 self.assertEqual(repr(c), 'count(-9)')
321 c.next()
322 self.assertEqual(c.next(), -8)
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000323 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 +0000324 # Test repr (ignoring the L in longs)
325 r1 = repr(count(i)).replace('L', '')
326 r2 = 'count(%r)'.__mod__(i).replace('L', '')
327 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000328
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000329 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000330 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000331 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000332 self.assertRaises(TypeError, cycle)
333 self.assertRaises(TypeError, cycle, 5)
334 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000335
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000336 def test_groupby(self):
337 # Check whether it accepts arguments correctly
338 self.assertEqual([], list(groupby([])))
339 self.assertEqual([], list(groupby([], key=id)))
340 self.assertRaises(TypeError, list, groupby('abc', []))
341 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000342 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000343
344 # Check normal input
345 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
346 (2,15,22), (3,16,23), (3,17,23)]
347 dup = []
348 for k, g in groupby(s, lambda r:r[0]):
349 for elem in g:
350 self.assertEqual(k, elem[0])
351 dup.append(elem)
352 self.assertEqual(s, dup)
353
354 # Check nested case
355 dup = []
356 for k, g in groupby(s, lambda r:r[0]):
357 for ik, ig in groupby(g, lambda r:r[2]):
358 for elem in ig:
359 self.assertEqual(k, elem[0])
360 self.assertEqual(ik, elem[2])
361 dup.append(elem)
362 self.assertEqual(s, dup)
363
364 # Check case where inner iterator is not used
365 keys = [k for k, g in groupby(s, lambda r:r[0])]
366 expectedkeys = set([r[0] for r in s])
367 self.assertEqual(set(keys), expectedkeys)
368 self.assertEqual(len(keys), len(expectedkeys))
369
370 # Exercise pipes and filters style
371 s = 'abracadabra'
372 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000373 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000374 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
375 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000376 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000377 self.assertEqual(r, ['a', 'b', 'r'])
378 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000379 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000380 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
381 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000382 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000383 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
384
385 # iter.next failure
386 class ExpectedError(Exception):
387 pass
388 def delayed_raise(n=0):
389 for i in range(n):
390 yield 'yo'
391 raise ExpectedError
392 def gulp(iterable, keyp=None, func=list):
393 return [func(g) for k, g in groupby(iterable, keyp)]
394
395 # iter.next failure on outer object
396 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
397 # iter.next failure on inner object
398 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
399
400 # __cmp__ failure
401 class DummyCmp:
402 def __cmp__(self, dst):
403 raise ExpectedError
404 s = [DummyCmp(), DummyCmp(), None]
405
406 # __cmp__ failure on outer object
407 self.assertRaises(ExpectedError, gulp, s, func=id)
408 # __cmp__ failure on inner object
409 self.assertRaises(ExpectedError, gulp, s)
410
411 # keyfunc failure
412 def keyfunc(obj):
413 if keyfunc.skip > 0:
414 keyfunc.skip -= 1
415 return obj
416 else:
417 raise ExpectedError
418
419 # keyfunc failure on outer object
420 keyfunc.skip = 0
421 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
422 keyfunc.skip = 1
423 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
424
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000425 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000426 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000427 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000428 self.assertEqual(list(ifilter(bool, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000429 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000430 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000431 self.assertRaises(TypeError, ifilter, lambda x:x)
432 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000433 self.assertRaises(TypeError, ifilter, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000434 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000435
436 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000437 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
438 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000439 self.assertEqual(list(ifilterfalse(bool, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000440 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000441 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000442 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
443 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000444 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000445 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000446
447 def test_izip(self):
448 ans = [(x,y) for x, y in izip('abc',count())]
449 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000450 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
451 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
Raymond Hettinger02420702003-06-29 20:36:23 +0000452 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000453 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000454 self.assertEqual(list(izip()), zip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000455 self.assertRaises(TypeError, izip, 3)
456 self.assertRaises(TypeError, izip, range(3), 3)
457 # Check tuple re-use (implementation detail)
458 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
459 zip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000460 self.assertEqual([pair for pair in izip('abc', 'def')],
461 zip('abc', 'def'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000462 ids = map(id, izip('abc', 'def'))
463 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000464 ids = map(id, list(izip('abc', 'def')))
465 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000466
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000467 def test_iziplongest(self):
468 for args in [
469 ['abc', range(6)],
470 [range(6), 'abc'],
471 [range(1000), range(2000,2100), range(3000,3050)],
472 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
473 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
474 ]:
475 target = map(None, *args)
476 self.assertEqual(list(izip_longest(*args)), target)
477 self.assertEqual(list(izip_longest(*args, **{})), target)
478 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
479 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Tim Petersea5962f2007-03-12 18:07:52 +0000480
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000481 self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
482
483 self.assertEqual(list(izip_longest()), zip())
484 self.assertEqual(list(izip_longest([])), zip([]))
485 self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
Tim Petersea5962f2007-03-12 18:07:52 +0000486
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000487 self.assertEqual(list(izip_longest('abc', 'defg', **{})), map(None, 'abc', 'defg')) # empty keyword dict
488 self.assertRaises(TypeError, izip_longest, 3)
489 self.assertRaises(TypeError, izip_longest, range(3), 3)
490
491 for stmt in [
492 "izip_longest('abc', fv=1)",
Tim Petersea5962f2007-03-12 18:07:52 +0000493 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000494 ]:
495 try:
496 eval(stmt, globals(), locals())
497 except TypeError:
498 pass
499 else:
500 self.fail('Did not raise Type in: ' + stmt)
Tim Petersea5962f2007-03-12 18:07:52 +0000501
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000502 # Check tuple re-use (implementation detail)
503 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
504 zip('abc', 'def'))
505 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
506 zip('abc', 'def'))
507 ids = map(id, izip_longest('abc', 'def'))
508 self.assertEqual(min(ids), max(ids))
509 ids = map(id, list(izip_longest('abc', 'def')))
510 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
511
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000512 def test_product(self):
513 for args, result in [
Raymond Hettingerd553d852008-03-04 04:17:08 +0000514 ([], [()]), # zero iterables
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000515 (['ab'], [('a',), ('b',)]), # one iterable
516 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
517 ([range(0), range(2), range(3)], []), # first iterable with zero length
518 ([range(2), range(0), range(3)], []), # middle iterable with zero length
519 ([range(2), range(3), range(0)], []), # last iterable with zero length
520 ]:
521 self.assertEqual(list(product(*args)), result)
Raymond Hettinger08ff6822008-02-29 02:21:48 +0000522 for r in range(4):
523 self.assertEqual(list(product(*(args*r))),
524 list(product(*args, **dict(repeat=r))))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000525 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
526 self.assertRaises(TypeError, product, range(6), None)
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000527
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000528 def product1(*args, **kwds):
529 pools = map(tuple, args) * kwds.get('repeat', 1)
530 n = len(pools)
531 if n == 0:
532 yield ()
533 return
534 if any(len(pool) == 0 for pool in pools):
535 return
536 indices = [0] * n
537 yield tuple(pool[i] for pool, i in zip(pools, indices))
538 while 1:
539 for i in reversed(range(n)): # right to left
540 if indices[i] == len(pools[i]) - 1:
541 continue
542 indices[i] += 1
543 for j in range(i+1, n):
544 indices[j] = 0
545 yield tuple(pool[i] for pool, i in zip(pools, indices))
546 break
547 else:
548 return
549
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000550 def product2(*args, **kwds):
551 'Pure python version used in docs'
552 pools = map(tuple, args) * kwds.get('repeat', 1)
553 result = [[]]
554 for pool in pools:
555 result = [x+[y] for x in result for y in pool]
556 for prod in result:
557 yield tuple(prod)
558
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000559 argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
560 set('abcdefg'), range(11), tuple(range(13))]
561 for i in range(100):
562 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Raymond Hettingerd553d852008-03-04 04:17:08 +0000563 expected_len = prod(map(len, args))
564 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000565 self.assertEqual(list(product(*args)), list(product1(*args)))
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000566 self.assertEqual(list(product(*args)), list(product2(*args)))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000567 args = map(iter, args)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000568 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000569
Raymond Hettinger73d79632008-02-23 02:20:41 +0000570 # Test implementation detail: tuple re-use
571 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
572 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000573
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000574 def test_repeat(self):
575 self.assertEqual(zip(xrange(3),repeat('a')),
576 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000577 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000578 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000579 self.assertEqual(list(repeat('a', 0)), [])
580 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000581 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000582 self.assertRaises(TypeError, repeat, None, 3, 4)
583 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000584 r = repeat(1+0j)
585 self.assertEqual(repr(r), 'repeat((1+0j))')
586 r = repeat(1+0j, 5)
587 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
588 list(r)
589 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000590
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000591 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000592 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
593 [0**1, 1**2, 2**3])
594 self.assertEqual(list(imap(None, 'abc', range(5))),
595 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000596 self.assertEqual(list(imap(None, 'abc', count())),
597 [('a',0),('b',1),('c',2)])
598 self.assertEqual(take(2,imap(None, 'abc', count())),
599 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000600 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000601 self.assertRaises(TypeError, imap)
602 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000603 self.assertRaises(TypeError, imap(10, range(5)).next)
604 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
605 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000606
607 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000608 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
609 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000610 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
611 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000612 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger47317092008-01-17 03:02:14 +0000613 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
614 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000615 self.assertRaises(TypeError, starmap)
616 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
617 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
618 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
619 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000620
621 def test_islice(self):
622 for args in [ # islice(args) should agree with range(args)
623 (10, 20, 3),
624 (10, 3, 20),
625 (10, 20),
626 (10, 3),
627 (20,)
628 ]:
629 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
630
631 for args, tgtargs in [ # Stop when seqn is exhausted
632 ((10, 110, 3), ((10, 100, 3))),
633 ((10, 110), ((10, 100))),
634 ((110,), (100,))
635 ]:
636 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
637
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000638 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000639 self.assertEqual(list(islice(xrange(10), None)), range(10))
Raymond Hettingerb2594052004-12-05 09:25:51 +0000640 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
641 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000642 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
643 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
644
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000645 # Test number of items consumed SF #1171417
646 it = iter(range(10))
647 self.assertEqual(list(islice(it, 3)), range(3))
648 self.assertEqual(list(it), range(3, 10))
649
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000650 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000651 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000652 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
653 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
654 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
655 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
656 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000657 self.assertRaises(ValueError, islice, xrange(10), 'a')
658 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
659 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
660 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
661 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +0000662 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000663
664 def test_takewhile(self):
665 data = [1, 3, 5, 20, 2, 4, 6, 8]
666 underten = lambda x: x<10
667 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000668 self.assertEqual(list(takewhile(underten, [])), [])
669 self.assertRaises(TypeError, takewhile)
670 self.assertRaises(TypeError, takewhile, operator.pow)
671 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
672 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
673 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000674 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
675 self.assertEqual(list(t), [1, 1, 1])
676 self.assertRaises(StopIteration, t.next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000677
678 def test_dropwhile(self):
679 data = [1, 3, 5, 20, 2, 4, 6, 8]
680 underten = lambda x: x<10
681 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000682 self.assertEqual(list(dropwhile(underten, [])), [])
683 self.assertRaises(TypeError, dropwhile)
684 self.assertRaises(TypeError, dropwhile, operator.pow)
685 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
686 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
687 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000688
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000689 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000690 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000691 def irange(n):
692 for i in xrange(n):
693 yield i
694
695 a, b = tee([]) # test empty iterator
696 self.assertEqual(list(a), [])
697 self.assertEqual(list(b), [])
698
699 a, b = tee(irange(n)) # test 100% interleaved
700 self.assertEqual(zip(a,b), zip(range(n),range(n)))
701
702 a, b = tee(irange(n)) # test 0% interleaved
703 self.assertEqual(list(a), range(n))
704 self.assertEqual(list(b), range(n))
705
706 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000707 for i in xrange(100):
708 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000709 del a
710 self.assertEqual(list(b), range(n))
711
712 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000713 for i in xrange(100):
714 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000715 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000716 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000717
718 for j in xrange(5): # test randomly interleaved
719 order = [0]*n + [1]*n
720 random.shuffle(order)
721 lists = ([], [])
722 its = tee(irange(n))
723 for i in order:
724 value = its[i].next()
725 lists[i].append(value)
726 self.assertEqual(lists[0], range(n))
727 self.assertEqual(lists[1], range(n))
728
Raymond Hettingerad983e72003-11-12 14:32:26 +0000729 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000730 self.assertRaises(TypeError, tee)
731 self.assertRaises(TypeError, tee, 3)
732 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000733 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000734
Raymond Hettingerad983e72003-11-12 14:32:26 +0000735 # tee object should be instantiable
736 a, b = tee('abc')
737 c = type(a)('def')
738 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000739
Raymond Hettingerad983e72003-11-12 14:32:26 +0000740 # test long-lagged and multi-way split
741 a, b, c = tee(xrange(2000), 3)
742 for i in xrange(100):
743 self.assertEqual(a.next(), i)
744 self.assertEqual(list(b), range(2000))
745 self.assertEqual([c.next(), c.next()], range(2))
746 self.assertEqual(list(a), range(100,2000))
747 self.assertEqual(list(c), range(2,2000))
748
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000749 # test values of n
750 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Neal Norwitz69e88972006-09-02 02:58:13 +0000751 self.assertRaises(ValueError, tee, [], -1)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000752 for n in xrange(5):
753 result = tee('abc', n)
754 self.assertEqual(type(result), tuple)
755 self.assertEqual(len(result), n)
756 self.assertEqual(map(list, result), [list('abc')]*n)
757
Raymond Hettingerad983e72003-11-12 14:32:26 +0000758 # tee pass-through to copyable iterator
759 a, b = tee('abc')
760 c, d = tee(a)
761 self.assert_(a is c)
762
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000763 # test tee_new
764 t1, t2 = tee('abc')
765 tnew = type(t1)
766 self.assertRaises(TypeError, tnew)
767 self.assertRaises(TypeError, tnew, 10)
768 t3 = tnew(t1)
769 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000770
Raymond Hettingera9f60922004-10-17 16:40:14 +0000771 # test that tee objects are weak referencable
772 a, b = tee(xrange(10))
773 p = proxy(a)
774 self.assertEqual(getattr(p, '__class__'), type(b))
775 del a
776 self.assertRaises(ReferenceError, getattr, p, '__class__')
777
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000778 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000779 self.assertRaises(StopIteration, izip().next)
780
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000781 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000782 self.assertRaises(StopIteration, f([]).next)
783 self.assertRaises(StopIteration, f(StopNow()).next)
784
785 self.assertRaises(StopIteration, islice([], None).next)
786 self.assertRaises(StopIteration, islice(StopNow(), None).next)
787
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000788 p, q = tee([])
789 self.assertRaises(StopIteration, p.next)
790 self.assertRaises(StopIteration, q.next)
791 p, q = tee(StopNow())
792 self.assertRaises(StopIteration, p.next)
793 self.assertRaises(StopIteration, q.next)
794
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000795 self.assertRaises(StopIteration, repeat(None, 0).next)
796
797 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
798 self.assertRaises(StopIteration, f(lambda x:x, []).next)
799 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
800
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000801class TestExamples(unittest.TestCase):
802
803 def test_chain(self):
804 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
805
806 def test_chain_from_iterable(self):
807 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
808
809 def test_combinations(self):
810 self.assertEqual(list(combinations('ABCD', 2)),
811 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
812 self.assertEqual(list(combinations(range(4), 3)),
813 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
814
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000815 def test_combinations_with_replacement(self):
816 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
817 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
818
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000819 def test_compress(self):
820 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
821
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000822 def test_count(self):
823 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
824
825 def test_cycle(self):
826 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
827
828 def test_dropwhile(self):
829 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
830
831 def test_groupby(self):
832 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
833 list('ABCDAB'))
834 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
835 [list('AAAA'), list('BBB'), list('CC'), list('D')])
836
837 def test_ifilter(self):
838 self.assertEqual(list(ifilter(lambda x: x%2, range(10))), [1,3,5,7,9])
839
840 def test_ifilterfalse(self):
841 self.assertEqual(list(ifilterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
842
843 def test_imap(self):
844 self.assertEqual(list(imap(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
845
846 def test_islice(self):
847 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
848 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
849 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
850 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
851
852 def test_izip(self):
853 self.assertEqual(list(izip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
854
855 def test_izip_longest(self):
856 self.assertEqual(list(izip_longest('ABCD', 'xy', fillvalue='-')),
857 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
858
859 def test_permutations(self):
860 self.assertEqual(list(permutations('ABCD', 2)),
861 map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split()))
862 self.assertEqual(list(permutations(range(3))),
863 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
864
865 def test_product(self):
866 self.assertEqual(list(product('ABCD', 'xy')),
867 map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split()))
868 self.assertEqual(list(product(range(2), repeat=3)),
869 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
870 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
871
872 def test_repeat(self):
873 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
874
875 def test_stapmap(self):
876 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
877 [32, 9, 1000])
878
879 def test_takewhile(self):
880 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
881
882
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000883class TestGC(unittest.TestCase):
884
885 def makecycle(self, iterator, container):
886 container.append(iterator)
887 iterator.next()
888 del container, iterator
889
890 def test_chain(self):
891 a = []
892 self.makecycle(chain(a), a)
893
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000894 def test_chain_from_iterable(self):
895 a = []
896 self.makecycle(chain.from_iterable([a]), a)
897
898 def test_combinations(self):
899 a = []
900 self.makecycle(combinations([1,2,a,3], 3), a)
901
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000902 def test_combinations_with_replacement(self):
903 a = []
904 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
905
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000906 def test_compress(self):
907 a = []
908 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
909
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000910 def test_cycle(self):
911 a = []
912 self.makecycle(cycle([a]*2), a)
913
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000914 def test_dropwhile(self):
915 a = []
916 self.makecycle(dropwhile(bool, [0, a, a]), a)
917
918 def test_groupby(self):
919 a = []
920 self.makecycle(groupby([a]*2, lambda x:x), a)
921
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000922 def test_issue2246(self):
923 # Issue 2246 -- the _grouper iterator was not included in GC
924 n = 10
925 keyfunc = lambda x: x
926 for i, j in groupby(xrange(n), key=keyfunc):
927 keyfunc.__dict__.setdefault('x',[]).append(j)
928
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000929 def test_ifilter(self):
930 a = []
931 self.makecycle(ifilter(lambda x:True, [a]*2), a)
932
933 def test_ifilterfalse(self):
934 a = []
935 self.makecycle(ifilterfalse(lambda x:False, a), a)
936
937 def test_izip(self):
938 a = []
939 self.makecycle(izip([a]*2, [a]*3), a)
940
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000941 def test_izip_longest(self):
942 a = []
943 self.makecycle(izip_longest([a]*2, [a]*3), a)
944 b = [a, None]
945 self.makecycle(izip_longest([a]*2, [a]*3, fillvalue=b), a)
946
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000947 def test_imap(self):
948 a = []
949 self.makecycle(imap(lambda x:x, [a]*2), a)
950
951 def test_islice(self):
952 a = []
953 self.makecycle(islice([a]*2, None), a)
954
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000955 def test_permutations(self):
956 a = []
957 self.makecycle(permutations([1,2,a,3], 3), a)
958
959 def test_product(self):
960 a = []
961 self.makecycle(product([1,2,a,3], repeat=3), a)
962
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000963 def test_repeat(self):
964 a = []
965 self.makecycle(repeat(a), a)
966
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000967 def test_starmap(self):
968 a = []
969 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
970
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000971 def test_takewhile(self):
972 a = []
973 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
974
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000975def R(seqn):
976 'Regular generator'
977 for i in seqn:
978 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000979
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000980class G:
981 'Sequence using __getitem__'
982 def __init__(self, seqn):
983 self.seqn = seqn
984 def __getitem__(self, i):
985 return self.seqn[i]
986
987class I:
988 'Sequence using iterator protocol'
989 def __init__(self, seqn):
990 self.seqn = seqn
991 self.i = 0
992 def __iter__(self):
993 return self
994 def next(self):
995 if self.i >= len(self.seqn): raise StopIteration
996 v = self.seqn[self.i]
997 self.i += 1
998 return v
999
1000class Ig:
1001 'Sequence using iterator protocol defined with a generator'
1002 def __init__(self, seqn):
1003 self.seqn = seqn
1004 self.i = 0
1005 def __iter__(self):
1006 for val in self.seqn:
1007 yield val
1008
1009class X:
1010 'Missing __getitem__ and __iter__'
1011 def __init__(self, seqn):
1012 self.seqn = seqn
1013 self.i = 0
1014 def next(self):
1015 if self.i >= len(self.seqn): raise StopIteration
1016 v = self.seqn[self.i]
1017 self.i += 1
1018 return v
1019
1020class N:
1021 'Iterator missing next()'
1022 def __init__(self, seqn):
1023 self.seqn = seqn
1024 self.i = 0
1025 def __iter__(self):
1026 return self
1027
1028class E:
1029 'Test propagation of exceptions'
1030 def __init__(self, seqn):
1031 self.seqn = seqn
1032 self.i = 0
1033 def __iter__(self):
1034 return self
1035 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001036 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001037
1038class S:
1039 'Test immediate stop'
1040 def __init__(self, seqn):
1041 pass
1042 def __iter__(self):
1043 return self
1044 def next(self):
1045 raise StopIteration
1046
1047def L(seqn):
1048 'Test multiple tiers of iterators'
1049 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
1050
1051
1052class TestVariousIteratorArgs(unittest.TestCase):
1053
1054 def test_chain(self):
1055 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1056 for g in (G, I, Ig, S, L, R):
1057 self.assertEqual(list(chain(g(s))), list(g(s)))
1058 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001059 self.assertRaises(TypeError, list, chain(X(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001060 self.assertRaises(TypeError, list, chain(N(s)))
1061 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1062
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001063 def test_compress(self):
1064 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1065 n = len(s)
1066 for g in (G, I, Ig, S, L, R):
1067 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1068 self.assertRaises(TypeError, compress, X(s), repeat(1))
1069 self.assertRaises(TypeError, list, compress(N(s), repeat(1)))
1070 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1071
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001072 def test_product(self):
1073 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1074 self.assertRaises(TypeError, product, X(s))
1075 self.assertRaises(TypeError, product, N(s))
1076 self.assertRaises(ZeroDivisionError, product, E(s))
1077
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001078 def test_cycle(self):
1079 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1080 for g in (G, I, Ig, S, L, R):
1081 tgtlen = len(s) * 3
1082 expected = list(g(s))*3
1083 actual = list(islice(cycle(g(s)), tgtlen))
1084 self.assertEqual(actual, expected)
1085 self.assertRaises(TypeError, cycle, X(s))
1086 self.assertRaises(TypeError, list, cycle(N(s)))
1087 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1088
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001089 def test_groupby(self):
1090 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1091 for g in (G, I, Ig, S, L, R):
1092 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1093 self.assertRaises(TypeError, groupby, X(s))
1094 self.assertRaises(TypeError, list, groupby(N(s)))
1095 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1096
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001097 def test_ifilter(self):
1098 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1099 for g in (G, I, Ig, S, L, R):
1100 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
1101 self.assertRaises(TypeError, ifilter, isEven, X(s))
1102 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
1103 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
1104
1105 def test_ifilterfalse(self):
1106 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1107 for g in (G, I, Ig, S, L, R):
1108 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
1109 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
1110 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
1111 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
1112
1113 def test_izip(self):
1114 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1115 for g in (G, I, Ig, S, L, R):
1116 self.assertEqual(list(izip(g(s))), zip(g(s)))
1117 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
1118 self.assertRaises(TypeError, izip, X(s))
1119 self.assertRaises(TypeError, list, izip(N(s)))
1120 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
1121
Raymond Hettingerd36862c2007-02-21 05:20:38 +00001122 def test_iziplongest(self):
1123 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1124 for g in (G, I, Ig, S, L, R):
1125 self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
1126 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
1127 self.assertRaises(TypeError, izip_longest, X(s))
1128 self.assertRaises(TypeError, list, izip_longest(N(s)))
1129 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
1130
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001131 def test_imap(self):
1132 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1133 for g in (G, I, Ig, S, L, R):
1134 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
1135 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
1136 self.assertRaises(TypeError, imap, onearg, X(s))
1137 self.assertRaises(TypeError, list, imap(onearg, N(s)))
1138 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
1139
1140 def test_islice(self):
1141 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1142 for g in (G, I, Ig, S, L, R):
1143 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1144 self.assertRaises(TypeError, islice, X(s), 10)
1145 self.assertRaises(TypeError, list, islice(N(s), 10))
1146 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1147
1148 def test_starmap(self):
1149 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1150 for g in (G, I, Ig, S, L, R):
1151 ss = zip(s, s)
1152 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
1153 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
1154 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
1155 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1156
1157 def test_takewhile(self):
1158 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1159 for g in (G, I, Ig, S, L, R):
1160 tgt = []
1161 for elem in g(s):
1162 if not isEven(elem): break
1163 tgt.append(elem)
1164 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1165 self.assertRaises(TypeError, takewhile, isEven, X(s))
1166 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
1167 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1168
1169 def test_dropwhile(self):
1170 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1171 for g in (G, I, Ig, S, L, R):
1172 tgt = []
1173 for elem in g(s):
1174 if not tgt and isOdd(elem): continue
1175 tgt.append(elem)
1176 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1177 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
1178 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
1179 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1180
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001181 def test_tee(self):
1182 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1183 for g in (G, I, Ig, S, L, R):
1184 it1, it2 = tee(g(s))
1185 self.assertEqual(list(it1), list(g(s)))
1186 self.assertEqual(list(it2), list(g(s)))
1187 self.assertRaises(TypeError, tee, X(s))
1188 self.assertRaises(TypeError, list, tee(N(s))[0])
1189 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1190
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001191class LengthTransparency(unittest.TestCase):
1192
1193 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001194 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001195 self.assertEqual(len(repeat(None, 50)), 50)
1196 self.assertRaises(TypeError, len, repeat(None))
1197
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001198class RegressionTests(unittest.TestCase):
1199
1200 def test_sf_793826(self):
1201 # Fix Armin Rigo's successful efforts to wreak havoc
1202
1203 def mutatingtuple(tuple1, f, tuple2):
1204 # this builds a tuple t which is a copy of tuple1,
1205 # then calls f(t), then mutates t to be equal to tuple2
1206 # (needs len(tuple1) == len(tuple2)).
1207 def g(value, first=[1]):
1208 if first:
1209 del first[:]
1210 f(z.next())
1211 return value
1212 items = list(tuple2)
1213 items[1:1] = list(tuple1)
1214 gen = imap(g, items)
1215 z = izip(*[gen]*len(tuple1))
1216 z.next()
1217
1218 def f(t):
1219 global T
1220 T = t
1221 first[:] = list(T)
1222
1223 first = []
1224 mutatingtuple((1,2,3), f, (4,5,6))
1225 second = list(T)
1226 self.assertEqual(first, second)
1227
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001228
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001229 def test_sf_950057(self):
1230 # Make sure that chain() and cycle() catch exceptions immediately
1231 # rather than when shifting between input sources
1232
1233 def gen1():
1234 hist.append(0)
1235 yield 1
1236 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001237 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001238 hist.append(2)
1239
1240 def gen2(x):
1241 hist.append(3)
1242 yield 2
1243 hist.append(4)
1244 if x:
1245 raise StopIteration
1246
1247 hist = []
1248 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1249 self.assertEqual(hist, [0,1])
1250
1251 hist = []
1252 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1253 self.assertEqual(hist, [0,1])
1254
1255 hist = []
1256 self.assertRaises(AssertionError, list, cycle(gen1()))
1257 self.assertEqual(hist, [0,1])
1258
Georg Brandlb84c1372007-01-21 10:28:43 +00001259class SubclassWithKwargsTest(unittest.TestCase):
1260 def test_keywords_in_subclass(self):
1261 # count is not subclassable...
1262 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001263 starmap, islice, takewhile, dropwhile, cycle, compress):
Georg Brandlb84c1372007-01-21 10:28:43 +00001264 class Subclass(cls):
1265 def __init__(self, newarg=None, *args):
1266 cls.__init__(self, *args)
1267 try:
1268 Subclass(newarg=1)
1269 except TypeError, err:
1270 # we expect type errors because of wrong argument count
1271 self.failIf("does not take keyword arguments" in err.args[0])
1272
1273
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001274libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001275
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001276
1277>>> amounts = [120.15, 764.05, 823.14]
1278>>> for checknum, amount in izip(count(1200), amounts):
1279... print 'Check %d is for $%.2f' % (checknum, amount)
1280...
1281Check 1200 is for $120.15
1282Check 1201 is for $764.05
1283Check 1202 is for $823.14
1284
1285>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001286>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1287... print cube
1288...
12891
12908
129127
1292
1293>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001294>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001295... print name.title()
1296...
1297Alex
1298Laura
1299Martin
1300Walter
1301Samuele
1302
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001303>>> from operator import itemgetter
1304>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Armin Rigoa3f09272006-05-28 19:13:17 +00001305>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001306>>> for k, g in groupby(di, itemgetter(1)):
1307... print k, map(itemgetter(0), g)
1308...
13091 ['a', 'c', 'e']
13102 ['b', 'd', 'f']
13113 ['g']
1312
Raymond Hettinger734fb572004-01-20 20:04:40 +00001313# Find runs of consecutive numbers using groupby. The key to the solution
1314# is differencing with a range so that consecutive numbers all appear in
1315# same group.
1316>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
1317>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
1318... print map(operator.itemgetter(1), g)
1319...
1320[1]
1321[4, 5, 6]
1322[10]
1323[15, 16, 17, 18]
1324[22]
1325[25, 26, 27, 28]
1326
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001327>>> def take(n, iterable):
1328... "Return first n items of the iterable as a list"
1329... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001330
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001331>>> def enumerate(iterable, start=0):
1332... return izip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001333
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001334>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001335... "Return function(0), function(1), ..."
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001336... return imap(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001337
1338>>> def nth(iterable, n):
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001339... "Returns the nth item or empty list"
Raymond Hettinger60eca932003-02-09 06:40:58 +00001340... return list(islice(iterable, n, n+1))
1341
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001342>>> def quantify(iterable, pred=bool):
1343... "Count how many times the predicate is true"
1344... return sum(imap(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001345
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001346>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001347... "Returns the sequence elements and then returns None indefinitely"
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001348... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001349
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001350>>> def ncycles(iterable, n):
1351... "Returns the seqeuence elements n times"
1352... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001353
1354>>> def dotproduct(vec1, vec2):
1355... return sum(imap(operator.mul, vec1, vec2))
1356
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001357>>> def flatten(listOfLists):
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001358... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001359
1360>>> def repeatfunc(func, times=None, *args):
1361... "Repeat calls to func with specified arguments."
1362... " Example: repeatfunc(random.random)"
1363... if times is None:
1364... return starmap(func, repeat(args))
1365... else:
1366... return starmap(func, repeat(args, times))
1367
Raymond Hettingerd591f662003-10-26 15:34:50 +00001368>>> def pairwise(iterable):
1369... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1370... a, b = tee(iterable)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001371... for elem in b:
1372... break
Raymond Hettingerad983e72003-11-12 14:32:26 +00001373... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001374
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001375>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001376... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001377... args = [iter(iterable)] * n
Raymond Hettingerf080e6d2008-07-31 01:19:50 +00001378... return izip_longest(fillvalue=fillvalue, *args)
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001379
1380>>> def roundrobin(*iterables):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001381... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001382... # Recipe credited to George Sakkis
1383... pending = len(iterables)
1384... nexts = cycle(iter(it).next for it in iterables)
1385... while pending:
1386... try:
1387... for next in nexts:
1388... yield next()
1389... except StopIteration:
1390... pending -= 1
1391... nexts = cycle(islice(nexts, pending))
1392
1393>>> def powerset(iterable):
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001394... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1395... s = list(iterable)
1396... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001397
Raymond Hettinger44e15812009-01-02 21:26:45 +00001398>>> def unique_everseen(iterable, key=None):
1399... "List unique elements, preserving order. Remember all elements ever seen."
1400... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1401... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1402... seen = set()
1403... seen_add = seen.add
1404... if key is None:
1405... for element in iterable:
1406... if element not in seen:
1407... seen_add(element)
1408... yield element
1409... else:
1410... for element in iterable:
1411... k = key(element)
1412... if k not in seen:
1413... seen_add(k)
1414... yield element
1415
1416>>> def unique_justseen(iterable, key=None):
1417... "List unique elements, preserving order. Remember only the element just seen."
1418... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1419... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1420... return imap(next, imap(itemgetter(1), groupby(iterable, key)))
1421
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001422This is not part of the examples but it tests to make sure the definitions
1423perform as purported.
1424
Raymond Hettingera098b332003-09-08 23:58:40 +00001425>>> take(10, count())
1426[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1427
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001428>>> list(enumerate('abc'))
1429[(0, 'a'), (1, 'b'), (2, 'c')]
1430
1431>>> list(islice(tabulate(lambda x: 2*x), 4))
1432[0, 2, 4, 6]
1433
1434>>> nth('abcde', 3)
1435['d']
1436
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001437>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000143850
1439
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001440>>> a = [[1, 2, 3], [4, 5, 6]]
1441>>> flatten(a)
1442[1, 2, 3, 4, 5, 6]
1443
1444>>> list(repeatfunc(pow, 5, 2, 3))
1445[8, 8, 8, 8, 8]
1446
1447>>> import random
1448>>> take(5, imap(int, repeatfunc(random.random)))
1449[0, 0, 0, 0, 0]
1450
Raymond Hettingerd591f662003-10-26 15:34:50 +00001451>>> list(pairwise('abcd'))
1452[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001453
Raymond Hettingerd591f662003-10-26 15:34:50 +00001454>>> list(pairwise([]))
1455[]
1456
1457>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001458[]
1459
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001460>>> list(islice(padnone('abc'), 0, 6))
1461['a', 'b', 'c', None, None, None]
1462
1463>>> list(ncycles('abc', 3))
1464['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1465
1466>>> dotproduct([1,2,3], [4,5,6])
146732
1468
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001469>>> list(grouper(3, 'abcdefg', 'x'))
1470[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1471
1472>>> list(roundrobin('abc', 'd', 'ef'))
1473['a', 'd', 'e', 'b', 'f', 'c']
1474
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001475>>> list(powerset([1,2,3]))
1476[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001477
Raymond Hettinger44e15812009-01-02 21:26:45 +00001478>>> list(unique_everseen('AAAABBBCCDAABBB'))
1479['A', 'B', 'C', 'D']
1480
1481>>> list(unique_everseen('ABBCcAD', str.lower))
1482['A', 'B', 'C', 'D']
1483
1484>>> list(unique_justseen('AAAABBBCCDAABBB'))
1485['A', 'B', 'C', 'D', 'A', 'B']
1486
1487>>> list(unique_justseen('ABBCcAD', str.lower))
1488['A', 'B', 'C', 'A', 'D']
1489
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001490"""
1491
1492__test__ = {'libreftest' : libreftest}
1493
1494def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001495 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Georg Brandlb84c1372007-01-21 10:28:43 +00001496 RegressionTests, LengthTransparency,
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001497 SubclassWithKwargsTest, TestExamples)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001498 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001499
1500 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001501 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001502 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001503 counts = [None] * 5
1504 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001505 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001506 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001507 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001508 print counts
1509
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001510 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001511 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001512
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001513if __name__ == "__main__":
1514 test_main(verbose=True)