blob: 09908d699b2a6c2e793412968601dbe26f621469 [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
Raymond Hettingerf93f3032009-01-27 09:52:35 +0000284 self.assertEquals(len(prod), len(s)**r)
285 self.assertEquals(prod, sorted(set(prod))) # prod in lexicographic order without repeats
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000286 self.assertEquals(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
287 self.assertEquals(perm, [t for t in prod if len(set(t))==r]) # perm: prods with no dups
288 self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
289 self.assertEqual(comb, [t for t in cwr if len(set(t))==r]) # comb: cwrs without dups
Raymond Hettingerf93f3032009-01-27 09:52:35 +0000290 self.assertEqual(comb, filter(set(cwr).__contains__, perm)) # comb: perm that is a cwr
291 self.assertEqual(comb, filter(set(perm).__contains__, cwr)) # comb: cwr that is a perm
292 self.assertEqual(comb, sorted(set(cwr) & set(perm))) # comb: both a cwr and a perm
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000293
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000294 def test_compress(self):
295 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
296 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
297 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
298 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
299 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
300 n = 10000
301 data = chain.from_iterable(repeat(range(6), n))
302 selectors = chain.from_iterable(repeat((0, 1)))
303 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
304 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
305 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
306 self.assertRaises(TypeError, compress, range(6)) # too few args
307 self.assertRaises(TypeError, compress, range(6), None) # too many args
308
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000309 def test_count(self):
310 self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
311 self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000312 self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000313 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
314 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000315 self.assertRaises(TypeError, count, 2, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000316 self.assertRaises(TypeError, count, 'a')
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000317 self.assertEqual(list(islice(count(maxsize-5), 10)), range(maxsize-5, maxsize+5))
318 self.assertEqual(list(islice(count(-maxsize-5), 10)), range(-maxsize-5, -maxsize+5))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000319 c = count(3)
320 self.assertEqual(repr(c), 'count(3)')
321 c.next()
322 self.assertEqual(repr(c), 'count(4)')
Jack Diederich36234e82006-09-21 17:50:26 +0000323 c = count(-9)
324 self.assertEqual(repr(c), 'count(-9)')
325 c.next()
326 self.assertEqual(c.next(), -8)
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000327 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 +0000328 # Test repr (ignoring the L in longs)
329 r1 = repr(count(i)).replace('L', '')
330 r2 = 'count(%r)'.__mod__(i).replace('L', '')
331 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000332
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000333 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000334 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000335 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000336 self.assertRaises(TypeError, cycle)
337 self.assertRaises(TypeError, cycle, 5)
338 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000339
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000340 def test_groupby(self):
341 # Check whether it accepts arguments correctly
342 self.assertEqual([], list(groupby([])))
343 self.assertEqual([], list(groupby([], key=id)))
344 self.assertRaises(TypeError, list, groupby('abc', []))
345 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000346 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000347
348 # Check normal input
349 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
350 (2,15,22), (3,16,23), (3,17,23)]
351 dup = []
352 for k, g in groupby(s, lambda r:r[0]):
353 for elem in g:
354 self.assertEqual(k, elem[0])
355 dup.append(elem)
356 self.assertEqual(s, dup)
357
358 # Check nested case
359 dup = []
360 for k, g in groupby(s, lambda r:r[0]):
361 for ik, ig in groupby(g, lambda r:r[2]):
362 for elem in ig:
363 self.assertEqual(k, elem[0])
364 self.assertEqual(ik, elem[2])
365 dup.append(elem)
366 self.assertEqual(s, dup)
367
368 # Check case where inner iterator is not used
369 keys = [k for k, g in groupby(s, lambda r:r[0])]
370 expectedkeys = set([r[0] for r in s])
371 self.assertEqual(set(keys), expectedkeys)
372 self.assertEqual(len(keys), len(expectedkeys))
373
374 # Exercise pipes and filters style
375 s = 'abracadabra'
376 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000377 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000378 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
379 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000380 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000381 self.assertEqual(r, ['a', 'b', 'r'])
382 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000383 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000384 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
385 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000386 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000387 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
388
389 # iter.next failure
390 class ExpectedError(Exception):
391 pass
392 def delayed_raise(n=0):
393 for i in range(n):
394 yield 'yo'
395 raise ExpectedError
396 def gulp(iterable, keyp=None, func=list):
397 return [func(g) for k, g in groupby(iterable, keyp)]
398
399 # iter.next failure on outer object
400 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
401 # iter.next failure on inner object
402 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
403
404 # __cmp__ failure
405 class DummyCmp:
406 def __cmp__(self, dst):
407 raise ExpectedError
408 s = [DummyCmp(), DummyCmp(), None]
409
410 # __cmp__ failure on outer object
411 self.assertRaises(ExpectedError, gulp, s, func=id)
412 # __cmp__ failure on inner object
413 self.assertRaises(ExpectedError, gulp, s)
414
415 # keyfunc failure
416 def keyfunc(obj):
417 if keyfunc.skip > 0:
418 keyfunc.skip -= 1
419 return obj
420 else:
421 raise ExpectedError
422
423 # keyfunc failure on outer object
424 keyfunc.skip = 0
425 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
426 keyfunc.skip = 1
427 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
428
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000429 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000430 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000431 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000432 self.assertEqual(list(ifilter(bool, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000433 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000434 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000435 self.assertRaises(TypeError, ifilter, lambda x:x)
436 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000437 self.assertRaises(TypeError, ifilter, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000438 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000439
440 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000441 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
442 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000443 self.assertEqual(list(ifilterfalse(bool, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000444 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000445 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000446 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
447 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000448 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000449 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000450
451 def test_izip(self):
452 ans = [(x,y) for x, y in izip('abc',count())]
453 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000454 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
455 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
Raymond Hettinger02420702003-06-29 20:36:23 +0000456 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000457 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000458 self.assertEqual(list(izip()), zip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000459 self.assertRaises(TypeError, izip, 3)
460 self.assertRaises(TypeError, izip, range(3), 3)
461 # Check tuple re-use (implementation detail)
462 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
463 zip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000464 self.assertEqual([pair for pair in izip('abc', 'def')],
465 zip('abc', 'def'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000466 ids = map(id, izip('abc', 'def'))
467 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000468 ids = map(id, list(izip('abc', 'def')))
469 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000470
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000471 def test_iziplongest(self):
472 for args in [
473 ['abc', range(6)],
474 [range(6), 'abc'],
475 [range(1000), range(2000,2100), range(3000,3050)],
476 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
477 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
478 ]:
479 target = map(None, *args)
480 self.assertEqual(list(izip_longest(*args)), target)
481 self.assertEqual(list(izip_longest(*args, **{})), target)
482 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
483 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Tim Petersea5962f2007-03-12 18:07:52 +0000484
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000485 self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
486
487 self.assertEqual(list(izip_longest()), zip())
488 self.assertEqual(list(izip_longest([])), zip([]))
489 self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
Tim Petersea5962f2007-03-12 18:07:52 +0000490
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000491 self.assertEqual(list(izip_longest('abc', 'defg', **{})), map(None, 'abc', 'defg')) # empty keyword dict
492 self.assertRaises(TypeError, izip_longest, 3)
493 self.assertRaises(TypeError, izip_longest, range(3), 3)
494
495 for stmt in [
496 "izip_longest('abc', fv=1)",
Tim Petersea5962f2007-03-12 18:07:52 +0000497 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000498 ]:
499 try:
500 eval(stmt, globals(), locals())
501 except TypeError:
502 pass
503 else:
504 self.fail('Did not raise Type in: ' + stmt)
Tim Petersea5962f2007-03-12 18:07:52 +0000505
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000506 # Check tuple re-use (implementation detail)
507 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
508 zip('abc', 'def'))
509 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
510 zip('abc', 'def'))
511 ids = map(id, izip_longest('abc', 'def'))
512 self.assertEqual(min(ids), max(ids))
513 ids = map(id, list(izip_longest('abc', 'def')))
514 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
515
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000516 def test_product(self):
517 for args, result in [
Raymond Hettingerd553d852008-03-04 04:17:08 +0000518 ([], [()]), # zero iterables
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000519 (['ab'], [('a',), ('b',)]), # one iterable
520 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
521 ([range(0), range(2), range(3)], []), # first iterable with zero length
522 ([range(2), range(0), range(3)], []), # middle iterable with zero length
523 ([range(2), range(3), range(0)], []), # last iterable with zero length
524 ]:
525 self.assertEqual(list(product(*args)), result)
Raymond Hettinger08ff6822008-02-29 02:21:48 +0000526 for r in range(4):
527 self.assertEqual(list(product(*(args*r))),
528 list(product(*args, **dict(repeat=r))))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000529 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
530 self.assertRaises(TypeError, product, range(6), None)
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000531
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000532 def product1(*args, **kwds):
533 pools = map(tuple, args) * kwds.get('repeat', 1)
534 n = len(pools)
535 if n == 0:
536 yield ()
537 return
538 if any(len(pool) == 0 for pool in pools):
539 return
540 indices = [0] * n
541 yield tuple(pool[i] for pool, i in zip(pools, indices))
542 while 1:
543 for i in reversed(range(n)): # right to left
544 if indices[i] == len(pools[i]) - 1:
545 continue
546 indices[i] += 1
547 for j in range(i+1, n):
548 indices[j] = 0
549 yield tuple(pool[i] for pool, i in zip(pools, indices))
550 break
551 else:
552 return
553
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000554 def product2(*args, **kwds):
555 'Pure python version used in docs'
556 pools = map(tuple, args) * kwds.get('repeat', 1)
557 result = [[]]
558 for pool in pools:
559 result = [x+[y] for x in result for y in pool]
560 for prod in result:
561 yield tuple(prod)
562
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000563 argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
564 set('abcdefg'), range(11), tuple(range(13))]
565 for i in range(100):
566 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Raymond Hettingerd553d852008-03-04 04:17:08 +0000567 expected_len = prod(map(len, args))
568 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000569 self.assertEqual(list(product(*args)), list(product1(*args)))
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000570 self.assertEqual(list(product(*args)), list(product2(*args)))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000571 args = map(iter, args)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000572 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000573
Raymond Hettinger73d79632008-02-23 02:20:41 +0000574 # Test implementation detail: tuple re-use
575 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
576 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000577
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000578 def test_repeat(self):
579 self.assertEqual(zip(xrange(3),repeat('a')),
580 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000581 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000582 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000583 self.assertEqual(list(repeat('a', 0)), [])
584 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000585 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000586 self.assertRaises(TypeError, repeat, None, 3, 4)
587 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000588 r = repeat(1+0j)
589 self.assertEqual(repr(r), 'repeat((1+0j))')
590 r = repeat(1+0j, 5)
591 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
592 list(r)
593 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000594
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000595 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000596 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
597 [0**1, 1**2, 2**3])
598 self.assertEqual(list(imap(None, 'abc', range(5))),
599 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000600 self.assertEqual(list(imap(None, 'abc', count())),
601 [('a',0),('b',1),('c',2)])
602 self.assertEqual(take(2,imap(None, 'abc', count())),
603 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000604 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000605 self.assertRaises(TypeError, imap)
606 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000607 self.assertRaises(TypeError, imap(10, range(5)).next)
608 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
609 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000610
611 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000612 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
613 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000614 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
615 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000616 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger47317092008-01-17 03:02:14 +0000617 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
618 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000619 self.assertRaises(TypeError, starmap)
620 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
621 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
622 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
623 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000624
625 def test_islice(self):
626 for args in [ # islice(args) should agree with range(args)
627 (10, 20, 3),
628 (10, 3, 20),
629 (10, 20),
630 (10, 3),
631 (20,)
632 ]:
633 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
634
635 for args, tgtargs in [ # Stop when seqn is exhausted
636 ((10, 110, 3), ((10, 100, 3))),
637 ((10, 110), ((10, 100))),
638 ((110,), (100,))
639 ]:
640 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
641
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000642 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000643 self.assertEqual(list(islice(xrange(10), None)), range(10))
Raymond Hettingerb2594052004-12-05 09:25:51 +0000644 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
645 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000646 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
647 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
648
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000649 # Test number of items consumed SF #1171417
650 it = iter(range(10))
651 self.assertEqual(list(islice(it, 3)), range(3))
652 self.assertEqual(list(it), range(3, 10))
653
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000654 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000655 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000656 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
657 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
658 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
659 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
660 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000661 self.assertRaises(ValueError, islice, xrange(10), 'a')
662 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
663 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
664 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
665 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +0000666 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000667
668 def test_takewhile(self):
669 data = [1, 3, 5, 20, 2, 4, 6, 8]
670 underten = lambda x: x<10
671 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000672 self.assertEqual(list(takewhile(underten, [])), [])
673 self.assertRaises(TypeError, takewhile)
674 self.assertRaises(TypeError, takewhile, operator.pow)
675 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
676 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
677 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000678 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
679 self.assertEqual(list(t), [1, 1, 1])
680 self.assertRaises(StopIteration, t.next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000681
682 def test_dropwhile(self):
683 data = [1, 3, 5, 20, 2, 4, 6, 8]
684 underten = lambda x: x<10
685 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000686 self.assertEqual(list(dropwhile(underten, [])), [])
687 self.assertRaises(TypeError, dropwhile)
688 self.assertRaises(TypeError, dropwhile, operator.pow)
689 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
690 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
691 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000692
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000693 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000694 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000695 def irange(n):
696 for i in xrange(n):
697 yield i
698
699 a, b = tee([]) # test empty iterator
700 self.assertEqual(list(a), [])
701 self.assertEqual(list(b), [])
702
703 a, b = tee(irange(n)) # test 100% interleaved
704 self.assertEqual(zip(a,b), zip(range(n),range(n)))
705
706 a, b = tee(irange(n)) # test 0% interleaved
707 self.assertEqual(list(a), range(n))
708 self.assertEqual(list(b), range(n))
709
710 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000711 for i in xrange(100):
712 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000713 del a
714 self.assertEqual(list(b), range(n))
715
716 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000717 for i in xrange(100):
718 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000719 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000720 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000721
722 for j in xrange(5): # test randomly interleaved
723 order = [0]*n + [1]*n
724 random.shuffle(order)
725 lists = ([], [])
726 its = tee(irange(n))
727 for i in order:
728 value = its[i].next()
729 lists[i].append(value)
730 self.assertEqual(lists[0], range(n))
731 self.assertEqual(lists[1], range(n))
732
Raymond Hettingerad983e72003-11-12 14:32:26 +0000733 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000734 self.assertRaises(TypeError, tee)
735 self.assertRaises(TypeError, tee, 3)
736 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000737 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000738
Raymond Hettingerad983e72003-11-12 14:32:26 +0000739 # tee object should be instantiable
740 a, b = tee('abc')
741 c = type(a)('def')
742 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000743
Raymond Hettingerad983e72003-11-12 14:32:26 +0000744 # test long-lagged and multi-way split
745 a, b, c = tee(xrange(2000), 3)
746 for i in xrange(100):
747 self.assertEqual(a.next(), i)
748 self.assertEqual(list(b), range(2000))
749 self.assertEqual([c.next(), c.next()], range(2))
750 self.assertEqual(list(a), range(100,2000))
751 self.assertEqual(list(c), range(2,2000))
752
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000753 # test values of n
754 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Neal Norwitz69e88972006-09-02 02:58:13 +0000755 self.assertRaises(ValueError, tee, [], -1)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000756 for n in xrange(5):
757 result = tee('abc', n)
758 self.assertEqual(type(result), tuple)
759 self.assertEqual(len(result), n)
760 self.assertEqual(map(list, result), [list('abc')]*n)
761
Raymond Hettingerad983e72003-11-12 14:32:26 +0000762 # tee pass-through to copyable iterator
763 a, b = tee('abc')
764 c, d = tee(a)
765 self.assert_(a is c)
766
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000767 # test tee_new
768 t1, t2 = tee('abc')
769 tnew = type(t1)
770 self.assertRaises(TypeError, tnew)
771 self.assertRaises(TypeError, tnew, 10)
772 t3 = tnew(t1)
773 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000774
Raymond Hettingera9f60922004-10-17 16:40:14 +0000775 # test that tee objects are weak referencable
776 a, b = tee(xrange(10))
777 p = proxy(a)
778 self.assertEqual(getattr(p, '__class__'), type(b))
779 del a
780 self.assertRaises(ReferenceError, getattr, p, '__class__')
781
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000782 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000783 self.assertRaises(StopIteration, izip().next)
784
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000785 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000786 self.assertRaises(StopIteration, f([]).next)
787 self.assertRaises(StopIteration, f(StopNow()).next)
788
789 self.assertRaises(StopIteration, islice([], None).next)
790 self.assertRaises(StopIteration, islice(StopNow(), None).next)
791
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000792 p, q = tee([])
793 self.assertRaises(StopIteration, p.next)
794 self.assertRaises(StopIteration, q.next)
795 p, q = tee(StopNow())
796 self.assertRaises(StopIteration, p.next)
797 self.assertRaises(StopIteration, q.next)
798
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000799 self.assertRaises(StopIteration, repeat(None, 0).next)
800
801 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
802 self.assertRaises(StopIteration, f(lambda x:x, []).next)
803 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
804
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000805class TestExamples(unittest.TestCase):
806
807 def test_chain(self):
808 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
809
810 def test_chain_from_iterable(self):
811 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
812
813 def test_combinations(self):
814 self.assertEqual(list(combinations('ABCD', 2)),
815 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
816 self.assertEqual(list(combinations(range(4), 3)),
817 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
818
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000819 def test_combinations_with_replacement(self):
820 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
821 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
822
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000823 def test_compress(self):
824 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
825
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000826 def test_count(self):
827 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
828
829 def test_cycle(self):
830 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
831
832 def test_dropwhile(self):
833 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
834
835 def test_groupby(self):
836 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
837 list('ABCDAB'))
838 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
839 [list('AAAA'), list('BBB'), list('CC'), list('D')])
840
841 def test_ifilter(self):
842 self.assertEqual(list(ifilter(lambda x: x%2, range(10))), [1,3,5,7,9])
843
844 def test_ifilterfalse(self):
845 self.assertEqual(list(ifilterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
846
847 def test_imap(self):
848 self.assertEqual(list(imap(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
849
850 def test_islice(self):
851 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
852 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
853 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
854 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
855
856 def test_izip(self):
857 self.assertEqual(list(izip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
858
859 def test_izip_longest(self):
860 self.assertEqual(list(izip_longest('ABCD', 'xy', fillvalue='-')),
861 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
862
863 def test_permutations(self):
864 self.assertEqual(list(permutations('ABCD', 2)),
865 map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split()))
866 self.assertEqual(list(permutations(range(3))),
867 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
868
869 def test_product(self):
870 self.assertEqual(list(product('ABCD', 'xy')),
871 map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split()))
872 self.assertEqual(list(product(range(2), repeat=3)),
873 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
874 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
875
876 def test_repeat(self):
877 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
878
879 def test_stapmap(self):
880 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
881 [32, 9, 1000])
882
883 def test_takewhile(self):
884 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
885
886
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000887class TestGC(unittest.TestCase):
888
889 def makecycle(self, iterator, container):
890 container.append(iterator)
891 iterator.next()
892 del container, iterator
893
894 def test_chain(self):
895 a = []
896 self.makecycle(chain(a), a)
897
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000898 def test_chain_from_iterable(self):
899 a = []
900 self.makecycle(chain.from_iterable([a]), a)
901
902 def test_combinations(self):
903 a = []
904 self.makecycle(combinations([1,2,a,3], 3), a)
905
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000906 def test_combinations_with_replacement(self):
907 a = []
908 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
909
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000910 def test_compress(self):
911 a = []
912 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
913
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000914 def test_cycle(self):
915 a = []
916 self.makecycle(cycle([a]*2), a)
917
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000918 def test_dropwhile(self):
919 a = []
920 self.makecycle(dropwhile(bool, [0, a, a]), a)
921
922 def test_groupby(self):
923 a = []
924 self.makecycle(groupby([a]*2, lambda x:x), a)
925
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000926 def test_issue2246(self):
927 # Issue 2246 -- the _grouper iterator was not included in GC
928 n = 10
929 keyfunc = lambda x: x
930 for i, j in groupby(xrange(n), key=keyfunc):
931 keyfunc.__dict__.setdefault('x',[]).append(j)
932
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000933 def test_ifilter(self):
934 a = []
935 self.makecycle(ifilter(lambda x:True, [a]*2), a)
936
937 def test_ifilterfalse(self):
938 a = []
939 self.makecycle(ifilterfalse(lambda x:False, a), a)
940
941 def test_izip(self):
942 a = []
943 self.makecycle(izip([a]*2, [a]*3), a)
944
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000945 def test_izip_longest(self):
946 a = []
947 self.makecycle(izip_longest([a]*2, [a]*3), a)
948 b = [a, None]
949 self.makecycle(izip_longest([a]*2, [a]*3, fillvalue=b), a)
950
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000951 def test_imap(self):
952 a = []
953 self.makecycle(imap(lambda x:x, [a]*2), a)
954
955 def test_islice(self):
956 a = []
957 self.makecycle(islice([a]*2, None), a)
958
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000959 def test_permutations(self):
960 a = []
961 self.makecycle(permutations([1,2,a,3], 3), a)
962
963 def test_product(self):
964 a = []
965 self.makecycle(product([1,2,a,3], repeat=3), a)
966
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000967 def test_repeat(self):
968 a = []
969 self.makecycle(repeat(a), a)
970
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000971 def test_starmap(self):
972 a = []
973 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
974
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000975 def test_takewhile(self):
976 a = []
977 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
978
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000979def R(seqn):
980 'Regular generator'
981 for i in seqn:
982 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000983
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000984class G:
985 'Sequence using __getitem__'
986 def __init__(self, seqn):
987 self.seqn = seqn
988 def __getitem__(self, i):
989 return self.seqn[i]
990
991class I:
992 'Sequence using iterator protocol'
993 def __init__(self, seqn):
994 self.seqn = seqn
995 self.i = 0
996 def __iter__(self):
997 return self
998 def next(self):
999 if self.i >= len(self.seqn): raise StopIteration
1000 v = self.seqn[self.i]
1001 self.i += 1
1002 return v
1003
1004class Ig:
1005 'Sequence using iterator protocol defined with a generator'
1006 def __init__(self, seqn):
1007 self.seqn = seqn
1008 self.i = 0
1009 def __iter__(self):
1010 for val in self.seqn:
1011 yield val
1012
1013class X:
1014 'Missing __getitem__ and __iter__'
1015 def __init__(self, seqn):
1016 self.seqn = seqn
1017 self.i = 0
1018 def next(self):
1019 if self.i >= len(self.seqn): raise StopIteration
1020 v = self.seqn[self.i]
1021 self.i += 1
1022 return v
1023
1024class N:
1025 'Iterator missing next()'
1026 def __init__(self, seqn):
1027 self.seqn = seqn
1028 self.i = 0
1029 def __iter__(self):
1030 return self
1031
1032class E:
1033 'Test propagation of exceptions'
1034 def __init__(self, seqn):
1035 self.seqn = seqn
1036 self.i = 0
1037 def __iter__(self):
1038 return self
1039 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001040 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001041
1042class S:
1043 'Test immediate stop'
1044 def __init__(self, seqn):
1045 pass
1046 def __iter__(self):
1047 return self
1048 def next(self):
1049 raise StopIteration
1050
1051def L(seqn):
1052 'Test multiple tiers of iterators'
1053 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
1054
1055
1056class TestVariousIteratorArgs(unittest.TestCase):
1057
1058 def test_chain(self):
1059 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1060 for g in (G, I, Ig, S, L, R):
1061 self.assertEqual(list(chain(g(s))), list(g(s)))
1062 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001063 self.assertRaises(TypeError, list, chain(X(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001064 self.assertRaises(TypeError, list, chain(N(s)))
1065 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1066
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001067 def test_compress(self):
1068 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1069 n = len(s)
1070 for g in (G, I, Ig, S, L, R):
1071 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1072 self.assertRaises(TypeError, compress, X(s), repeat(1))
1073 self.assertRaises(TypeError, list, compress(N(s), repeat(1)))
1074 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1075
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001076 def test_product(self):
1077 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1078 self.assertRaises(TypeError, product, X(s))
1079 self.assertRaises(TypeError, product, N(s))
1080 self.assertRaises(ZeroDivisionError, product, E(s))
1081
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001082 def test_cycle(self):
1083 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1084 for g in (G, I, Ig, S, L, R):
1085 tgtlen = len(s) * 3
1086 expected = list(g(s))*3
1087 actual = list(islice(cycle(g(s)), tgtlen))
1088 self.assertEqual(actual, expected)
1089 self.assertRaises(TypeError, cycle, X(s))
1090 self.assertRaises(TypeError, list, cycle(N(s)))
1091 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1092
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001093 def test_groupby(self):
1094 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1095 for g in (G, I, Ig, S, L, R):
1096 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1097 self.assertRaises(TypeError, groupby, X(s))
1098 self.assertRaises(TypeError, list, groupby(N(s)))
1099 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1100
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001101 def test_ifilter(self):
1102 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1103 for g in (G, I, Ig, S, L, R):
1104 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
1105 self.assertRaises(TypeError, ifilter, isEven, X(s))
1106 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
1107 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
1108
1109 def test_ifilterfalse(self):
1110 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1111 for g in (G, I, Ig, S, L, R):
1112 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
1113 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
1114 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
1115 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
1116
1117 def test_izip(self):
1118 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1119 for g in (G, I, Ig, S, L, R):
1120 self.assertEqual(list(izip(g(s))), zip(g(s)))
1121 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
1122 self.assertRaises(TypeError, izip, X(s))
1123 self.assertRaises(TypeError, list, izip(N(s)))
1124 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
1125
Raymond Hettingerd36862c2007-02-21 05:20:38 +00001126 def test_iziplongest(self):
1127 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1128 for g in (G, I, Ig, S, L, R):
1129 self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
1130 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
1131 self.assertRaises(TypeError, izip_longest, X(s))
1132 self.assertRaises(TypeError, list, izip_longest(N(s)))
1133 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
1134
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001135 def test_imap(self):
1136 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1137 for g in (G, I, Ig, S, L, R):
1138 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
1139 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
1140 self.assertRaises(TypeError, imap, onearg, X(s))
1141 self.assertRaises(TypeError, list, imap(onearg, N(s)))
1142 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
1143
1144 def test_islice(self):
1145 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1146 for g in (G, I, Ig, S, L, R):
1147 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1148 self.assertRaises(TypeError, islice, X(s), 10)
1149 self.assertRaises(TypeError, list, islice(N(s), 10))
1150 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1151
1152 def test_starmap(self):
1153 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1154 for g in (G, I, Ig, S, L, R):
1155 ss = zip(s, s)
1156 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
1157 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
1158 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
1159 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1160
1161 def test_takewhile(self):
1162 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1163 for g in (G, I, Ig, S, L, R):
1164 tgt = []
1165 for elem in g(s):
1166 if not isEven(elem): break
1167 tgt.append(elem)
1168 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1169 self.assertRaises(TypeError, takewhile, isEven, X(s))
1170 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
1171 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1172
1173 def test_dropwhile(self):
1174 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1175 for g in (G, I, Ig, S, L, R):
1176 tgt = []
1177 for elem in g(s):
1178 if not tgt and isOdd(elem): continue
1179 tgt.append(elem)
1180 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1181 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
1182 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
1183 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1184
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001185 def test_tee(self):
1186 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1187 for g in (G, I, Ig, S, L, R):
1188 it1, it2 = tee(g(s))
1189 self.assertEqual(list(it1), list(g(s)))
1190 self.assertEqual(list(it2), list(g(s)))
1191 self.assertRaises(TypeError, tee, X(s))
1192 self.assertRaises(TypeError, list, tee(N(s))[0])
1193 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1194
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001195class LengthTransparency(unittest.TestCase):
1196
1197 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001198 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001199 self.assertEqual(len(repeat(None, 50)), 50)
1200 self.assertRaises(TypeError, len, repeat(None))
1201
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001202class RegressionTests(unittest.TestCase):
1203
1204 def test_sf_793826(self):
1205 # Fix Armin Rigo's successful efforts to wreak havoc
1206
1207 def mutatingtuple(tuple1, f, tuple2):
1208 # this builds a tuple t which is a copy of tuple1,
1209 # then calls f(t), then mutates t to be equal to tuple2
1210 # (needs len(tuple1) == len(tuple2)).
1211 def g(value, first=[1]):
1212 if first:
1213 del first[:]
1214 f(z.next())
1215 return value
1216 items = list(tuple2)
1217 items[1:1] = list(tuple1)
1218 gen = imap(g, items)
1219 z = izip(*[gen]*len(tuple1))
1220 z.next()
1221
1222 def f(t):
1223 global T
1224 T = t
1225 first[:] = list(T)
1226
1227 first = []
1228 mutatingtuple((1,2,3), f, (4,5,6))
1229 second = list(T)
1230 self.assertEqual(first, second)
1231
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001232
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001233 def test_sf_950057(self):
1234 # Make sure that chain() and cycle() catch exceptions immediately
1235 # rather than when shifting between input sources
1236
1237 def gen1():
1238 hist.append(0)
1239 yield 1
1240 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001241 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001242 hist.append(2)
1243
1244 def gen2(x):
1245 hist.append(3)
1246 yield 2
1247 hist.append(4)
1248 if x:
1249 raise StopIteration
1250
1251 hist = []
1252 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1253 self.assertEqual(hist, [0,1])
1254
1255 hist = []
1256 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1257 self.assertEqual(hist, [0,1])
1258
1259 hist = []
1260 self.assertRaises(AssertionError, list, cycle(gen1()))
1261 self.assertEqual(hist, [0,1])
1262
Georg Brandlb84c1372007-01-21 10:28:43 +00001263class SubclassWithKwargsTest(unittest.TestCase):
1264 def test_keywords_in_subclass(self):
1265 # count is not subclassable...
1266 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001267 starmap, islice, takewhile, dropwhile, cycle, compress):
Georg Brandlb84c1372007-01-21 10:28:43 +00001268 class Subclass(cls):
1269 def __init__(self, newarg=None, *args):
1270 cls.__init__(self, *args)
1271 try:
1272 Subclass(newarg=1)
1273 except TypeError, err:
1274 # we expect type errors because of wrong argument count
1275 self.failIf("does not take keyword arguments" in err.args[0])
1276
1277
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001278libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001279
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001280
1281>>> amounts = [120.15, 764.05, 823.14]
1282>>> for checknum, amount in izip(count(1200), amounts):
1283... print 'Check %d is for $%.2f' % (checknum, amount)
1284...
1285Check 1200 is for $120.15
1286Check 1201 is for $764.05
1287Check 1202 is for $823.14
1288
1289>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001290>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1291... print cube
1292...
12931
12948
129527
1296
1297>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001298>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001299... print name.title()
1300...
1301Alex
1302Laura
1303Martin
1304Walter
1305Samuele
1306
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001307>>> from operator import itemgetter
1308>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Armin Rigoa3f09272006-05-28 19:13:17 +00001309>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001310>>> for k, g in groupby(di, itemgetter(1)):
1311... print k, map(itemgetter(0), g)
1312...
13131 ['a', 'c', 'e']
13142 ['b', 'd', 'f']
13153 ['g']
1316
Raymond Hettinger734fb572004-01-20 20:04:40 +00001317# Find runs of consecutive numbers using groupby. The key to the solution
1318# is differencing with a range so that consecutive numbers all appear in
1319# same group.
1320>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
1321>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
1322... print map(operator.itemgetter(1), g)
1323...
1324[1]
1325[4, 5, 6]
1326[10]
1327[15, 16, 17, 18]
1328[22]
1329[25, 26, 27, 28]
1330
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001331>>> def take(n, iterable):
1332... "Return first n items of the iterable as a list"
1333... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001334
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001335>>> def enumerate(iterable, start=0):
1336... return izip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001337
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001338>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001339... "Return function(0), function(1), ..."
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001340... return imap(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001341
1342>>> def nth(iterable, n):
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001343... "Returns the nth item or empty list"
Raymond Hettinger60eca932003-02-09 06:40:58 +00001344... return list(islice(iterable, n, n+1))
1345
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001346>>> def quantify(iterable, pred=bool):
1347... "Count how many times the predicate is true"
1348... return sum(imap(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001349
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001350>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001351... "Returns the sequence elements and then returns None indefinitely"
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001352... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001353
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001354>>> def ncycles(iterable, n):
1355... "Returns the seqeuence elements n times"
1356... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001357
1358>>> def dotproduct(vec1, vec2):
1359... return sum(imap(operator.mul, vec1, vec2))
1360
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001361>>> def flatten(listOfLists):
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001362... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001363
1364>>> def repeatfunc(func, times=None, *args):
1365... "Repeat calls to func with specified arguments."
1366... " Example: repeatfunc(random.random)"
1367... if times is None:
1368... return starmap(func, repeat(args))
1369... else:
1370... return starmap(func, repeat(args, times))
1371
Raymond Hettingerd591f662003-10-26 15:34:50 +00001372>>> def pairwise(iterable):
1373... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1374... a, b = tee(iterable)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001375... for elem in b:
1376... break
Raymond Hettingerad983e72003-11-12 14:32:26 +00001377... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001378
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001379>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001380... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001381... args = [iter(iterable)] * n
Raymond Hettingerf080e6d2008-07-31 01:19:50 +00001382... return izip_longest(fillvalue=fillvalue, *args)
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001383
1384>>> def roundrobin(*iterables):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001385... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001386... # Recipe credited to George Sakkis
1387... pending = len(iterables)
1388... nexts = cycle(iter(it).next for it in iterables)
1389... while pending:
1390... try:
1391... for next in nexts:
1392... yield next()
1393... except StopIteration:
1394... pending -= 1
1395... nexts = cycle(islice(nexts, pending))
1396
1397>>> def powerset(iterable):
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001398... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1399... s = list(iterable)
1400... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001401
Raymond Hettinger44e15812009-01-02 21:26:45 +00001402>>> def unique_everseen(iterable, key=None):
1403... "List unique elements, preserving order. Remember all elements ever seen."
1404... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1405... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1406... seen = set()
1407... seen_add = seen.add
1408... if key is None:
1409... for element in iterable:
1410... if element not in seen:
1411... seen_add(element)
1412... yield element
1413... else:
1414... for element in iterable:
1415... k = key(element)
1416... if k not in seen:
1417... seen_add(k)
1418... yield element
1419
1420>>> def unique_justseen(iterable, key=None):
1421... "List unique elements, preserving order. Remember only the element just seen."
1422... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1423... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1424... return imap(next, imap(itemgetter(1), groupby(iterable, key)))
1425
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001426This is not part of the examples but it tests to make sure the definitions
1427perform as purported.
1428
Raymond Hettingera098b332003-09-08 23:58:40 +00001429>>> take(10, count())
1430[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1431
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001432>>> list(enumerate('abc'))
1433[(0, 'a'), (1, 'b'), (2, 'c')]
1434
1435>>> list(islice(tabulate(lambda x: 2*x), 4))
1436[0, 2, 4, 6]
1437
1438>>> nth('abcde', 3)
1439['d']
1440
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001441>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000144250
1443
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001444>>> a = [[1, 2, 3], [4, 5, 6]]
1445>>> flatten(a)
1446[1, 2, 3, 4, 5, 6]
1447
1448>>> list(repeatfunc(pow, 5, 2, 3))
1449[8, 8, 8, 8, 8]
1450
1451>>> import random
1452>>> take(5, imap(int, repeatfunc(random.random)))
1453[0, 0, 0, 0, 0]
1454
Raymond Hettingerd591f662003-10-26 15:34:50 +00001455>>> list(pairwise('abcd'))
1456[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001457
Raymond Hettingerd591f662003-10-26 15:34:50 +00001458>>> list(pairwise([]))
1459[]
1460
1461>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001462[]
1463
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001464>>> list(islice(padnone('abc'), 0, 6))
1465['a', 'b', 'c', None, None, None]
1466
1467>>> list(ncycles('abc', 3))
1468['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1469
1470>>> dotproduct([1,2,3], [4,5,6])
147132
1472
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001473>>> list(grouper(3, 'abcdefg', 'x'))
1474[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1475
1476>>> list(roundrobin('abc', 'd', 'ef'))
1477['a', 'd', 'e', 'b', 'f', 'c']
1478
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001479>>> list(powerset([1,2,3]))
1480[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001481
Raymond Hettinger44e15812009-01-02 21:26:45 +00001482>>> list(unique_everseen('AAAABBBCCDAABBB'))
1483['A', 'B', 'C', 'D']
1484
1485>>> list(unique_everseen('ABBCcAD', str.lower))
1486['A', 'B', 'C', 'D']
1487
1488>>> list(unique_justseen('AAAABBBCCDAABBB'))
1489['A', 'B', 'C', 'D', 'A', 'B']
1490
1491>>> list(unique_justseen('ABBCcAD', str.lower))
1492['A', 'B', 'C', 'A', 'D']
1493
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001494"""
1495
1496__test__ = {'libreftest' : libreftest}
1497
1498def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001499 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Georg Brandlb84c1372007-01-21 10:28:43 +00001500 RegressionTests, LengthTransparency,
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001501 SubclassWithKwargsTest, TestExamples)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001502 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001503
1504 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001505 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001506 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001507 counts = [None] * 5
1508 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001509 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001510 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001511 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001512 print counts
1513
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001514 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001515 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001516
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001517if __name__ == "__main__":
1518 test_main(verbose=True)