blob: 03f1d85111ef27a1bbdb00396774824b4b9416d5 [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 Hettingeraa044612009-02-12 12:04:26 +00005from decimal import Decimal
Raymond Hettingerdbe3bfb2009-02-12 12:43:01 +00006from fractions import Fraction
Raymond Hettinger2012f172003-02-07 05:32:58 +00007import sys
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00008import operator
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00009import random
Raymond Hettingere09f45a2009-11-30 19:44:40 +000010import copy
11import pickle
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +000012maxsize = test_support.MAX_Py_ssize_t
13minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000014
15def onearg(x):
16 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000017 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000018
19def errfunc(*args):
20 'Test function that raises an error'
21 raise ValueError
22
23def gen3():
24 'Non-restartable source sequence'
25 for i in (0, 1, 2):
26 yield i
27
28def isEven(x):
29 'Test predicate'
30 return x%2==0
31
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000032def isOdd(x):
33 'Test predicate'
34 return x%2==1
35
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000036class StopNow:
37 'Class emulating an empty iterable.'
38 def __iter__(self):
39 return self
40 def next(self):
41 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000042
Raymond Hettinger02420702003-06-29 20:36:23 +000043def take(n, seq):
44 'Convenience function for partially consuming a long of infinite iterable'
45 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000046
Raymond Hettingerd553d852008-03-04 04:17:08 +000047def prod(iterable):
48 return reduce(operator.mul, iterable, 1)
49
Raymond Hettinger93e804d2008-02-26 23:40:50 +000050def fact(n):
51 'Factorial'
Raymond Hettingerd553d852008-03-04 04:17:08 +000052 return prod(range(1, n+1))
53
Raymond Hettinger96ef8112003-02-01 00:10:11 +000054class TestBasicOps(unittest.TestCase):
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000055 def test_chain(self):
Raymond Hettingerad47fa12008-03-06 20:52:01 +000056
57 def chain2(*iterables):
58 'Pure python version in the docs'
59 for it in iterables:
60 for element in it:
61 yield element
62
63 for c in (chain, chain2):
64 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
65 self.assertEqual(list(c('abc')), list('abc'))
66 self.assertEqual(list(c('')), [])
67 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
68 self.assertRaises(TypeError, list,c(2, 3))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000069
Raymond Hettingerb4cbc982008-02-28 22:46:41 +000070 def test_chain_from_iterable(self):
71 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
72 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
73 self.assertEqual(list(chain.from_iterable([''])), [])
74 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
75 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
76
Raymond Hettinger93e804d2008-02-26 23:40:50 +000077 def test_combinations(self):
Raymond Hettinger5b913e32009-01-08 06:39:04 +000078 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Raymond Hettinger93e804d2008-02-26 23:40:50 +000079 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Raymond Hettingerd553d852008-03-04 04:17:08 +000080 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Raymond Hettinger93e804d2008-02-26 23:40:50 +000081 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Raymond Hettinger5b913e32009-01-08 06:39:04 +000082 self.assertEqual(list(combinations('abc', 32)), []) # r > n
Raymond Hettinger93e804d2008-02-26 23:40:50 +000083 self.assertEqual(list(combinations(range(4), 3)),
84 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
Raymond Hettingerd553d852008-03-04 04:17:08 +000085
86 def combinations1(iterable, r):
87 'Pure python version shown in the docs'
88 pool = tuple(iterable)
89 n = len(pool)
Raymond Hettinger5b913e32009-01-08 06:39:04 +000090 if r > n:
91 return
Raymond Hettingerd553d852008-03-04 04:17:08 +000092 indices = range(r)
93 yield tuple(pool[i] for i in indices)
94 while 1:
95 for i in reversed(range(r)):
96 if indices[i] != i + n - r:
97 break
98 else:
99 return
100 indices[i] += 1
101 for j in range(i+1, r):
102 indices[j] = indices[j-1] + 1
103 yield tuple(pool[i] for i in indices)
104
105 def combinations2(iterable, r):
106 'Pure python version shown in the docs'
107 pool = tuple(iterable)
108 n = len(pool)
109 for indices in permutations(range(n), r):
110 if sorted(indices) == list(indices):
111 yield tuple(pool[i] for i in indices)
112
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000113 def combinations3(iterable, r):
114 'Pure python version from cwr()'
115 pool = tuple(iterable)
116 n = len(pool)
117 for indices in combinations_with_replacement(range(n), r):
118 if len(set(indices)) == r:
119 yield tuple(pool[i] for i in indices)
120
Raymond Hettingerd553d852008-03-04 04:17:08 +0000121 for n in range(7):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000122 values = [5*x-12 for x in range(n)]
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000123 for r in range(n+2):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000124 result = list(combinations(values, r))
Senthil Kumarance8e33a2010-01-08 19:04:16 +0000125 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 +0000126 self.assertEqual(len(result), len(set(result))) # no repeats
127 self.assertEqual(result, sorted(result)) # lexicographic order
128 for c in result:
129 self.assertEqual(len(c), r) # r-length combinations
130 self.assertEqual(len(set(c)), r) # no duplicate elements
131 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Peterson5c8da862009-06-30 22:57:08 +0000132 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000133 self.assertEqual(list(c),
134 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Raymond Hettingerd553d852008-03-04 04:17:08 +0000135 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000136 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000137 self.assertEqual(result, list(combinations3(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000138
139 # Test implementation detail: tuple re-use
140 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
141 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
142
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000143 def test_combinations_with_replacement(self):
144 cwr = combinations_with_replacement
145 self.assertRaises(TypeError, cwr, 'abc') # missing r argument
146 self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
147 self.assertRaises(TypeError, cwr, None) # pool is not iterable
148 self.assertRaises(ValueError, cwr, 'abc', -2) # r is negative
149 self.assertEqual(list(cwr('ABC', 2)),
150 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
151
152 def cwr1(iterable, r):
153 'Pure python version shown in the docs'
154 # number items returned: (n+r-1)! / r! / (n-1)! when n>0
155 pool = tuple(iterable)
156 n = len(pool)
157 if not n and r:
158 return
159 indices = [0] * r
160 yield tuple(pool[i] for i in indices)
161 while 1:
162 for i in reversed(range(r)):
163 if indices[i] != n - 1:
164 break
165 else:
166 return
167 indices[i:] = [indices[i] + 1] * (r - i)
168 yield tuple(pool[i] for i in indices)
169
170 def cwr2(iterable, r):
171 'Pure python version shown in the docs'
172 pool = tuple(iterable)
173 n = len(pool)
174 for indices in product(range(n), repeat=r):
175 if sorted(indices) == list(indices):
176 yield tuple(pool[i] for i in indices)
177
178 def numcombs(n, r):
179 if not n:
180 return 0 if r else 1
Senthil Kumarance8e33a2010-01-08 19:04:16 +0000181 return fact(n+r-1) / fact(r)/ fact(n-1)
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000182
183 for n in range(7):
184 values = [5*x-12 for x in range(n)]
185 for r in range(n+2):
186 result = list(cwr(values, r))
187
188 self.assertEqual(len(result), numcombs(n, r)) # right number of combs
189 self.assertEqual(len(result), len(set(result))) # no repeats
190 self.assertEqual(result, sorted(result)) # lexicographic order
191
192 regular_combs = list(combinations(values, r)) # compare to combs without replacement
193 if n == 0 or r <= 1:
194 self.assertEquals(result, regular_combs) # cases that should be identical
195 else:
Benjamin Peterson5c8da862009-06-30 22:57:08 +0000196 self.assertTrue(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000197
198 for c in result:
199 self.assertEqual(len(c), r) # r-length combinations
200 noruns = [k for k,v in groupby(c)] # combo without consecutive repeats
201 self.assertEqual(len(noruns), len(set(noruns))) # no repeats other than consecutive
202 self.assertEqual(list(c), sorted(c)) # keep original ordering
Benjamin Peterson5c8da862009-06-30 22:57:08 +0000203 self.assertTrue(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000204 self.assertEqual(noruns,
205 [e for e in values if e in c]) # comb is a subsequence of the input iterable
206 self.assertEqual(result, list(cwr1(values, r))) # matches first pure python version
207 self.assertEqual(result, list(cwr2(values, r))) # matches second pure python version
208
209 # Test implementation detail: tuple re-use
210 self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
211 self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
212
Raymond Hettingerd553d852008-03-04 04:17:08 +0000213 def test_permutations(self):
214 self.assertRaises(TypeError, permutations) # too few arguments
215 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000216 self.assertRaises(TypeError, permutations, None) # pool is not iterable
217 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000218 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000219 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Raymond Hettingerd553d852008-03-04 04:17:08 +0000220 self.assertEqual(list(permutations(range(3), 2)),
221 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
222
223 def permutations1(iterable, r=None):
224 'Pure python version shown in the docs'
225 pool = tuple(iterable)
226 n = len(pool)
227 r = n if r is None else r
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000228 if r > n:
229 return
Raymond Hettingerd553d852008-03-04 04:17:08 +0000230 indices = range(n)
Raymond Hettingere70bb8d2008-03-23 00:55:46 +0000231 cycles = range(n, n-r, -1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000232 yield tuple(pool[i] for i in indices[:r])
233 while n:
234 for i in reversed(range(r)):
235 cycles[i] -= 1
236 if cycles[i] == 0:
237 indices[i:] = indices[i+1:] + indices[i:i+1]
238 cycles[i] = n - i
239 else:
240 j = cycles[i]
241 indices[i], indices[-j] = indices[-j], indices[i]
242 yield tuple(pool[i] for i in indices[:r])
243 break
244 else:
245 return
246
247 def permutations2(iterable, r=None):
248 'Pure python version shown in the docs'
249 pool = tuple(iterable)
250 n = len(pool)
251 r = n if r is None else r
252 for indices in product(range(n), repeat=r):
253 if len(set(indices)) == r:
254 yield tuple(pool[i] for i in indices)
255
256 for n in range(7):
257 values = [5*x-12 for x in range(n)]
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000258 for r in range(n+2):
Raymond Hettingerd553d852008-03-04 04:17:08 +0000259 result = list(permutations(values, r))
Senthil Kumarance8e33a2010-01-08 19:04:16 +0000260 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 +0000261 self.assertEqual(len(result), len(set(result))) # no repeats
262 self.assertEqual(result, sorted(result)) # lexicographic order
263 for p in result:
264 self.assertEqual(len(p), r) # r-length permutations
265 self.assertEqual(len(set(p)), r) # no duplicate elements
Benjamin Peterson5c8da862009-06-30 22:57:08 +0000266 self.assertTrue(all(e in values for e in p)) # elements taken from input iterable
Raymond Hettingerd553d852008-03-04 04:17:08 +0000267 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000268 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000269 if r == n:
270 self.assertEqual(result, list(permutations(values, None))) # test r as None
271 self.assertEqual(result, list(permutations(values))) # test default r
272
273 # Test implementation detail: tuple re-use
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000274 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000275 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000276
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000277 def test_combinatorics(self):
278 # Test relationships between product(), permutations(),
279 # combinations() and combinations_with_replacement().
280
Raymond Hettinger2f6c2e02009-01-27 10:36:14 +0000281 for n in range(6):
282 s = 'ABCDEFG'[:n]
283 for r in range(8):
284 prod = list(product(s, repeat=r))
285 cwr = list(combinations_with_replacement(s, r))
286 perm = list(permutations(s, r))
287 comb = list(combinations(s, r))
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000288
Raymond Hettinger2f6c2e02009-01-27 10:36:14 +0000289 # Check size
290 self.assertEquals(len(prod), n**r)
Senthil Kumarance8e33a2010-01-08 19:04:16 +0000291 self.assertEquals(len(cwr), (fact(n+r-1) / fact(r)/ fact(n-1)) if n else (not r))
292 self.assertEquals(len(perm), 0 if r>n else fact(n) / fact(n-r))
293 self.assertEquals(len(comb), 0 if r>n else fact(n) / fact(r) / fact(n-r))
Raymond Hettinger2f6c2e02009-01-27 10:36:14 +0000294
295 # Check lexicographic order without repeated tuples
296 self.assertEquals(prod, sorted(set(prod)))
297 self.assertEquals(cwr, sorted(set(cwr)))
298 self.assertEquals(perm, sorted(set(perm)))
299 self.assertEquals(comb, sorted(set(comb)))
300
301 # Check interrelationships
302 self.assertEquals(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
303 self.assertEquals(perm, [t for t in prod if len(set(t))==r]) # perm: prods with no dups
304 self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
305 self.assertEqual(comb, [t for t in cwr if len(set(t))==r]) # comb: cwrs without dups
306 self.assertEqual(comb, filter(set(cwr).__contains__, perm)) # comb: perm that is a cwr
307 self.assertEqual(comb, filter(set(perm).__contains__, cwr)) # comb: cwr that is a perm
308 self.assertEqual(comb, sorted(set(cwr) & set(perm))) # comb: both a cwr and a perm
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000309
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000310 def test_compress(self):
Raymond Hettinger2e2909f2009-02-19 02:15:14 +0000311 self.assertEqual(list(compress(data='ABCDEF', selectors=[1,0,1,0,1,1])), list('ACEF'))
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000312 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
313 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
314 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
315 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
316 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
317 n = 10000
318 data = chain.from_iterable(repeat(range(6), n))
319 selectors = chain.from_iterable(repeat((0, 1)))
320 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
321 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
322 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
323 self.assertRaises(TypeError, compress, range(6)) # too few args
324 self.assertRaises(TypeError, compress, range(6), None) # too many args
325
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000326 def test_count(self):
327 self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
328 self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000329 self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000330 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
331 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000332 self.assertRaises(TypeError, count, 2, 3, 4)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000333 self.assertRaises(TypeError, count, 'a')
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000334 self.assertEqual(list(islice(count(maxsize-5), 10)), range(maxsize-5, maxsize+5))
335 self.assertEqual(list(islice(count(-maxsize-5), 10)), range(-maxsize-5, -maxsize+5))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000336 c = count(3)
337 self.assertEqual(repr(c), 'count(3)')
338 c.next()
339 self.assertEqual(repr(c), 'count(4)')
Jack Diederich36234e82006-09-21 17:50:26 +0000340 c = count(-9)
341 self.assertEqual(repr(c), 'count(-9)')
342 c.next()
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000343 self.assertEqual(repr(count(10.25)), 'count(10.25)')
Jack Diederich36234e82006-09-21 17:50:26 +0000344 self.assertEqual(c.next(), -8)
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000345 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 +0000346 # Test repr (ignoring the L in longs)
347 r1 = repr(count(i)).replace('L', '')
348 r2 = 'count(%r)'.__mod__(i).replace('L', '')
349 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000350
Raymond Hettingere09f45a2009-11-30 19:44:40 +0000351 # check copy, deepcopy, pickle
352 for value in -3, 3, sys.maxint-5, sys.maxint+5:
353 c = count(value)
354 self.assertEqual(next(copy.copy(c)), value)
355 self.assertEqual(next(copy.deepcopy(c)), value)
356 self.assertEqual(next(pickle.loads(pickle.dumps(c))), value)
357
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000358 def test_count_with_stride(self):
359 self.assertEqual(zip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingera4038032009-02-14 00:25:51 +0000360 self.assertEqual(zip('abc',count(start=2,step=3)),
361 [('a', 2), ('b', 5), ('c', 8)])
Raymond Hettingeraa681c72009-02-21 07:17:22 +0000362 self.assertEqual(zip('abc',count(step=-1)),
363 [('a', 0), ('b', -1), ('c', -2)])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000364 self.assertEqual(zip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
365 self.assertEqual(zip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
366 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
367 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
368 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettingeraa044612009-02-12 12:04:26 +0000369 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
370 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettingerdbe3bfb2009-02-12 12:43:01 +0000371 self.assertEqual(take(3, count(Fraction(2,3), Fraction(1,7))),
372 [Fraction(2,3), Fraction(17,21), Fraction(20,21)])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000373 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
374 c = count(3, 5)
375 self.assertEqual(repr(c), 'count(3, 5)')
376 c.next()
377 self.assertEqual(repr(c), 'count(8, 5)')
378 c = count(-9, 0)
379 self.assertEqual(repr(c), 'count(-9, 0)')
380 c.next()
381 self.assertEqual(repr(c), 'count(-9, 0)')
382 c = count(-9, -3)
383 self.assertEqual(repr(c), 'count(-9, -3)')
384 c.next()
385 self.assertEqual(repr(c), 'count(-12, -3)')
386 self.assertEqual(repr(c), 'count(-12, -3)')
387 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
388 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
389 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
390 for i in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 10, sys.maxint-5, sys.maxint+5):
391 for j in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 1, 10, sys.maxint-5, sys.maxint+5):
392 # Test repr (ignoring the L in longs)
393 r1 = repr(count(i, j)).replace('L', '')
394 if j == 1:
395 r2 = ('count(%r)' % i).replace('L', '')
396 else:
397 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
398 self.assertEqual(r1, r2)
399
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000400 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000401 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000402 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000403 self.assertRaises(TypeError, cycle)
404 self.assertRaises(TypeError, cycle, 5)
405 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000406
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000407 def test_groupby(self):
408 # Check whether it accepts arguments correctly
409 self.assertEqual([], list(groupby([])))
410 self.assertEqual([], list(groupby([], key=id)))
411 self.assertRaises(TypeError, list, groupby('abc', []))
412 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000413 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000414
415 # Check normal input
416 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
417 (2,15,22), (3,16,23), (3,17,23)]
418 dup = []
419 for k, g in groupby(s, lambda r:r[0]):
420 for elem in g:
421 self.assertEqual(k, elem[0])
422 dup.append(elem)
423 self.assertEqual(s, dup)
424
425 # Check nested case
426 dup = []
427 for k, g in groupby(s, lambda r:r[0]):
428 for ik, ig in groupby(g, lambda r:r[2]):
429 for elem in ig:
430 self.assertEqual(k, elem[0])
431 self.assertEqual(ik, elem[2])
432 dup.append(elem)
433 self.assertEqual(s, dup)
434
435 # Check case where inner iterator is not used
436 keys = [k for k, g in groupby(s, lambda r:r[0])]
437 expectedkeys = set([r[0] for r in s])
438 self.assertEqual(set(keys), expectedkeys)
439 self.assertEqual(len(keys), len(expectedkeys))
440
441 # Exercise pipes and filters style
442 s = 'abracadabra'
443 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000444 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000445 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
446 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000447 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000448 self.assertEqual(r, ['a', 'b', 'r'])
449 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000450 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000451 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
452 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000453 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000454 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
455
456 # iter.next failure
457 class ExpectedError(Exception):
458 pass
459 def delayed_raise(n=0):
460 for i in range(n):
461 yield 'yo'
462 raise ExpectedError
463 def gulp(iterable, keyp=None, func=list):
464 return [func(g) for k, g in groupby(iterable, keyp)]
465
466 # iter.next failure on outer object
467 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
468 # iter.next failure on inner object
469 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
470
471 # __cmp__ failure
472 class DummyCmp:
473 def __cmp__(self, dst):
474 raise ExpectedError
475 s = [DummyCmp(), DummyCmp(), None]
476
477 # __cmp__ failure on outer object
478 self.assertRaises(ExpectedError, gulp, s, func=id)
479 # __cmp__ failure on inner object
480 self.assertRaises(ExpectedError, gulp, s)
481
482 # keyfunc failure
483 def keyfunc(obj):
484 if keyfunc.skip > 0:
485 keyfunc.skip -= 1
486 return obj
487 else:
488 raise ExpectedError
489
490 # keyfunc failure on outer object
491 keyfunc.skip = 0
492 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
493 keyfunc.skip = 1
494 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
495
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000496 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000497 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000498 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000499 self.assertEqual(list(ifilter(bool, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000500 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000501 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000502 self.assertRaises(TypeError, ifilter, lambda x:x)
503 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000504 self.assertRaises(TypeError, ifilter, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000505 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000506
507 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000508 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
509 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000510 self.assertEqual(list(ifilterfalse(bool, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000511 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000512 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000513 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
514 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000515 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000516 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000517
518 def test_izip(self):
519 ans = [(x,y) for x, y in izip('abc',count())]
520 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000521 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
522 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
Raymond Hettinger02420702003-06-29 20:36:23 +0000523 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000524 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000525 self.assertEqual(list(izip()), zip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000526 self.assertRaises(TypeError, izip, 3)
527 self.assertRaises(TypeError, izip, range(3), 3)
528 # Check tuple re-use (implementation detail)
529 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
530 zip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000531 self.assertEqual([pair for pair in izip('abc', 'def')],
532 zip('abc', 'def'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000533 ids = map(id, izip('abc', 'def'))
534 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000535 ids = map(id, list(izip('abc', 'def')))
536 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000537
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000538 def test_iziplongest(self):
539 for args in [
540 ['abc', range(6)],
541 [range(6), 'abc'],
542 [range(1000), range(2000,2100), range(3000,3050)],
543 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
544 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
545 ]:
Senthil Kumarance8e33a2010-01-08 19:04:16 +0000546 target = map(None, *args)
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000547 self.assertEqual(list(izip_longest(*args)), target)
548 self.assertEqual(list(izip_longest(*args, **{})), target)
549 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
550 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Tim Petersea5962f2007-03-12 18:07:52 +0000551
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000552 self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
553
554 self.assertEqual(list(izip_longest()), zip())
555 self.assertEqual(list(izip_longest([])), zip([]))
556 self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
Tim Petersea5962f2007-03-12 18:07:52 +0000557
Senthil Kumarance8e33a2010-01-08 19:04:16 +0000558 self.assertEqual(list(izip_longest('abc', 'defg', **{})), map(None, 'abc', 'defg')) # empty keyword dict
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000559 self.assertRaises(TypeError, izip_longest, 3)
560 self.assertRaises(TypeError, izip_longest, range(3), 3)
561
562 for stmt in [
563 "izip_longest('abc', fv=1)",
Tim Petersea5962f2007-03-12 18:07:52 +0000564 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000565 ]:
566 try:
567 eval(stmt, globals(), locals())
568 except TypeError:
569 pass
570 else:
571 self.fail('Did not raise Type in: ' + stmt)
Tim Petersea5962f2007-03-12 18:07:52 +0000572
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000573 # Check tuple re-use (implementation detail)
574 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
575 zip('abc', 'def'))
576 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
577 zip('abc', 'def'))
578 ids = map(id, izip_longest('abc', 'def'))
579 self.assertEqual(min(ids), max(ids))
580 ids = map(id, list(izip_longest('abc', 'def')))
581 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
582
Raymond Hettingerfa7dadd2009-11-01 20:45:16 +0000583 def test_bug_7244(self):
584
585 class Repeater(object):
586 # this class is similar to itertools.repeat
587 def __init__(self, o, t, e):
588 self.o = o
589 self.t = int(t)
590 self.e = e
591 def __iter__(self): # its iterator is itself
592 return self
593 def next(self):
594 if self.t > 0:
595 self.t -= 1
596 return self.o
597 else:
598 raise self.e
599
600 # Formerly this code in would fail in debug mode
601 # with Undetected Error and Stop Iteration
602 r1 = Repeater(1, 3, StopIteration)
603 r2 = Repeater(2, 4, StopIteration)
604 def run(r1, r2):
605 result = []
606 for i, j in izip_longest(r1, r2, fillvalue=0):
607 with test_support.captured_output('stdout'):
608 print (i, j)
609 result.append((i, j))
610 return result
611 self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
612
613 # Formerly, the RuntimeError would be lost
614 # and StopIteration would stop as expected
615 r1 = Repeater(1, 3, RuntimeError)
616 r2 = Repeater(2, 4, StopIteration)
617 it = izip_longest(r1, r2, fillvalue=0)
618 self.assertEqual(next(it), (1, 2))
619 self.assertEqual(next(it), (1, 2))
620 self.assertEqual(next(it), (1, 2))
621 self.assertRaises(RuntimeError, next, it)
622
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000623 def test_product(self):
624 for args, result in [
Raymond Hettingerd553d852008-03-04 04:17:08 +0000625 ([], [()]), # zero iterables
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000626 (['ab'], [('a',), ('b',)]), # one iterable
627 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
628 ([range(0), range(2), range(3)], []), # first iterable with zero length
629 ([range(2), range(0), range(3)], []), # middle iterable with zero length
630 ([range(2), range(3), range(0)], []), # last iterable with zero length
631 ]:
632 self.assertEqual(list(product(*args)), result)
Raymond Hettinger08ff6822008-02-29 02:21:48 +0000633 for r in range(4):
634 self.assertEqual(list(product(*(args*r))),
635 list(product(*args, **dict(repeat=r))))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000636 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
637 self.assertRaises(TypeError, product, range(6), None)
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000638
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000639 def product1(*args, **kwds):
640 pools = map(tuple, args) * kwds.get('repeat', 1)
641 n = len(pools)
642 if n == 0:
643 yield ()
644 return
645 if any(len(pool) == 0 for pool in pools):
646 return
647 indices = [0] * n
648 yield tuple(pool[i] for pool, i in zip(pools, indices))
649 while 1:
650 for i in reversed(range(n)): # right to left
651 if indices[i] == len(pools[i]) - 1:
652 continue
653 indices[i] += 1
654 for j in range(i+1, n):
655 indices[j] = 0
656 yield tuple(pool[i] for pool, i in zip(pools, indices))
657 break
658 else:
659 return
660
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000661 def product2(*args, **kwds):
662 'Pure python version used in docs'
663 pools = map(tuple, args) * kwds.get('repeat', 1)
664 result = [[]]
665 for pool in pools:
666 result = [x+[y] for x in result for y in pool]
667 for prod in result:
668 yield tuple(prod)
669
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000670 argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
671 set('abcdefg'), range(11), tuple(range(13))]
672 for i in range(100):
673 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Raymond Hettingerd553d852008-03-04 04:17:08 +0000674 expected_len = prod(map(len, args))
675 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000676 self.assertEqual(list(product(*args)), list(product1(*args)))
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000677 self.assertEqual(list(product(*args)), list(product2(*args)))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000678 args = map(iter, args)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000679 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000680
Raymond Hettinger73d79632008-02-23 02:20:41 +0000681 # Test implementation detail: tuple re-use
682 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
683 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000684
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000685 def test_repeat(self):
Raymond Hettinger182edae2009-02-19 02:38:25 +0000686 self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a'])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000687 self.assertEqual(zip(xrange(3),repeat('a')),
688 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000689 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000690 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000691 self.assertEqual(list(repeat('a', 0)), [])
692 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000693 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000694 self.assertRaises(TypeError, repeat, None, 3, 4)
695 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000696 r = repeat(1+0j)
697 self.assertEqual(repr(r), 'repeat((1+0j))')
698 r = repeat(1+0j, 5)
699 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
700 list(r)
701 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000702
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000703 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000704 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
705 [0**1, 1**2, 2**3])
706 self.assertEqual(list(imap(None, 'abc', range(5))),
707 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000708 self.assertEqual(list(imap(None, 'abc', count())),
709 [('a',0),('b',1),('c',2)])
710 self.assertEqual(take(2,imap(None, 'abc', count())),
711 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000712 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000713 self.assertRaises(TypeError, imap)
714 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000715 self.assertRaises(TypeError, imap(10, range(5)).next)
716 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
717 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000718
719 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000720 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
721 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000722 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
723 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000724 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger47317092008-01-17 03:02:14 +0000725 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
726 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000727 self.assertRaises(TypeError, starmap)
728 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
729 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
730 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
731 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000732
733 def test_islice(self):
734 for args in [ # islice(args) should agree with range(args)
735 (10, 20, 3),
736 (10, 3, 20),
737 (10, 20),
738 (10, 3),
739 (20,)
740 ]:
741 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
742
743 for args, tgtargs in [ # Stop when seqn is exhausted
744 ((10, 110, 3), ((10, 100, 3))),
745 ((10, 110), ((10, 100))),
746 ((110,), (100,))
747 ]:
748 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
749
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000750 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000751 self.assertEqual(list(islice(xrange(10), None)), range(10))
Raymond Hettingerb2594052004-12-05 09:25:51 +0000752 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
753 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000754 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
755 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
756
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000757 # Test number of items consumed SF #1171417
758 it = iter(range(10))
759 self.assertEqual(list(islice(it, 3)), range(3))
760 self.assertEqual(list(it), range(3, 10))
761
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000762 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000763 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000764 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
765 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
766 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
767 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
768 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000769 self.assertRaises(ValueError, islice, xrange(10), 'a')
770 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
771 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
772 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
773 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +0000774 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000775
776 def test_takewhile(self):
777 data = [1, 3, 5, 20, 2, 4, 6, 8]
778 underten = lambda x: x<10
779 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000780 self.assertEqual(list(takewhile(underten, [])), [])
781 self.assertRaises(TypeError, takewhile)
782 self.assertRaises(TypeError, takewhile, operator.pow)
783 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
784 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
785 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000786 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
787 self.assertEqual(list(t), [1, 1, 1])
788 self.assertRaises(StopIteration, t.next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000789
790 def test_dropwhile(self):
791 data = [1, 3, 5, 20, 2, 4, 6, 8]
792 underten = lambda x: x<10
793 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000794 self.assertEqual(list(dropwhile(underten, [])), [])
795 self.assertRaises(TypeError, dropwhile)
796 self.assertRaises(TypeError, dropwhile, operator.pow)
797 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
798 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
799 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000800
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000801 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000802 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000803 def irange(n):
804 for i in xrange(n):
805 yield i
806
807 a, b = tee([]) # test empty iterator
808 self.assertEqual(list(a), [])
809 self.assertEqual(list(b), [])
810
811 a, b = tee(irange(n)) # test 100% interleaved
812 self.assertEqual(zip(a,b), zip(range(n),range(n)))
813
814 a, b = tee(irange(n)) # test 0% interleaved
815 self.assertEqual(list(a), range(n))
816 self.assertEqual(list(b), range(n))
817
818 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000819 for i in xrange(100):
820 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000821 del a
822 self.assertEqual(list(b), range(n))
823
824 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000825 for i in xrange(100):
826 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000827 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000828 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000829
830 for j in xrange(5): # test randomly interleaved
831 order = [0]*n + [1]*n
832 random.shuffle(order)
833 lists = ([], [])
834 its = tee(irange(n))
835 for i in order:
836 value = its[i].next()
837 lists[i].append(value)
838 self.assertEqual(lists[0], range(n))
839 self.assertEqual(lists[1], range(n))
840
Raymond Hettingerad983e72003-11-12 14:32:26 +0000841 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000842 self.assertRaises(TypeError, tee)
843 self.assertRaises(TypeError, tee, 3)
844 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000845 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000846
Raymond Hettingerad983e72003-11-12 14:32:26 +0000847 # tee object should be instantiable
848 a, b = tee('abc')
849 c = type(a)('def')
850 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000851
Raymond Hettingerad983e72003-11-12 14:32:26 +0000852 # test long-lagged and multi-way split
853 a, b, c = tee(xrange(2000), 3)
854 for i in xrange(100):
855 self.assertEqual(a.next(), i)
856 self.assertEqual(list(b), range(2000))
857 self.assertEqual([c.next(), c.next()], range(2))
858 self.assertEqual(list(a), range(100,2000))
859 self.assertEqual(list(c), range(2,2000))
860
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000861 # test values of n
862 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Neal Norwitz69e88972006-09-02 02:58:13 +0000863 self.assertRaises(ValueError, tee, [], -1)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000864 for n in xrange(5):
865 result = tee('abc', n)
866 self.assertEqual(type(result), tuple)
867 self.assertEqual(len(result), n)
868 self.assertEqual(map(list, result), [list('abc')]*n)
869
Raymond Hettingerad983e72003-11-12 14:32:26 +0000870 # tee pass-through to copyable iterator
871 a, b = tee('abc')
872 c, d = tee(a)
Benjamin Peterson5c8da862009-06-30 22:57:08 +0000873 self.assertTrue(a is c)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000874
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000875 # test tee_new
876 t1, t2 = tee('abc')
877 tnew = type(t1)
878 self.assertRaises(TypeError, tnew)
879 self.assertRaises(TypeError, tnew, 10)
880 t3 = tnew(t1)
Benjamin Peterson5c8da862009-06-30 22:57:08 +0000881 self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000882
Raymond Hettingera9f60922004-10-17 16:40:14 +0000883 # test that tee objects are weak referencable
884 a, b = tee(xrange(10))
885 p = proxy(a)
886 self.assertEqual(getattr(p, '__class__'), type(b))
887 del a
888 self.assertRaises(ReferenceError, getattr, p, '__class__')
889
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000890 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000891 self.assertRaises(StopIteration, izip().next)
892
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000893 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000894 self.assertRaises(StopIteration, f([]).next)
895 self.assertRaises(StopIteration, f(StopNow()).next)
896
897 self.assertRaises(StopIteration, islice([], None).next)
898 self.assertRaises(StopIteration, islice(StopNow(), None).next)
899
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000900 p, q = tee([])
901 self.assertRaises(StopIteration, p.next)
902 self.assertRaises(StopIteration, q.next)
903 p, q = tee(StopNow())
904 self.assertRaises(StopIteration, p.next)
905 self.assertRaises(StopIteration, q.next)
906
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000907 self.assertRaises(StopIteration, repeat(None, 0).next)
908
909 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
910 self.assertRaises(StopIteration, f(lambda x:x, []).next)
911 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
912
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000913class TestExamples(unittest.TestCase):
914
915 def test_chain(self):
916 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
917
918 def test_chain_from_iterable(self):
919 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
920
921 def test_combinations(self):
922 self.assertEqual(list(combinations('ABCD', 2)),
923 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
924 self.assertEqual(list(combinations(range(4), 3)),
925 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
926
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000927 def test_combinations_with_replacement(self):
928 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
929 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
930
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000931 def test_compress(self):
932 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
933
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000934 def test_count(self):
935 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
936
937 def test_cycle(self):
938 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
939
940 def test_dropwhile(self):
941 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
942
943 def test_groupby(self):
944 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
945 list('ABCDAB'))
946 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
947 [list('AAAA'), list('BBB'), list('CC'), list('D')])
948
949 def test_ifilter(self):
950 self.assertEqual(list(ifilter(lambda x: x%2, range(10))), [1,3,5,7,9])
951
952 def test_ifilterfalse(self):
953 self.assertEqual(list(ifilterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
954
955 def test_imap(self):
956 self.assertEqual(list(imap(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
957
958 def test_islice(self):
959 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
960 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
961 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
962 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
963
964 def test_izip(self):
965 self.assertEqual(list(izip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
966
967 def test_izip_longest(self):
968 self.assertEqual(list(izip_longest('ABCD', 'xy', fillvalue='-')),
969 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
970
971 def test_permutations(self):
972 self.assertEqual(list(permutations('ABCD', 2)),
973 map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split()))
974 self.assertEqual(list(permutations(range(3))),
975 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
976
977 def test_product(self):
978 self.assertEqual(list(product('ABCD', 'xy')),
979 map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split()))
980 self.assertEqual(list(product(range(2), repeat=3)),
981 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
982 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
983
984 def test_repeat(self):
985 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
986
987 def test_stapmap(self):
988 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
989 [32, 9, 1000])
990
991 def test_takewhile(self):
992 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
993
994
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000995class TestGC(unittest.TestCase):
996
997 def makecycle(self, iterator, container):
998 container.append(iterator)
999 iterator.next()
1000 del container, iterator
1001
1002 def test_chain(self):
1003 a = []
1004 self.makecycle(chain(a), a)
1005
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001006 def test_chain_from_iterable(self):
1007 a = []
1008 self.makecycle(chain.from_iterable([a]), a)
1009
1010 def test_combinations(self):
1011 a = []
1012 self.makecycle(combinations([1,2,a,3], 3), a)
1013
Raymond Hettingerd081abc2009-01-27 02:58:49 +00001014 def test_combinations_with_replacement(self):
1015 a = []
1016 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
1017
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001018 def test_compress(self):
1019 a = []
1020 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
1021
Raymond Hettingerb21d8102009-02-16 20:39:12 +00001022 def test_count(self):
1023 a = []
1024 Int = type('Int', (int,), dict(x=a))
1025 self.makecycle(count(Int(0), Int(1)), a)
1026
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001027 def test_cycle(self):
1028 a = []
1029 self.makecycle(cycle([a]*2), a)
1030
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001031 def test_dropwhile(self):
1032 a = []
1033 self.makecycle(dropwhile(bool, [0, a, a]), a)
1034
1035 def test_groupby(self):
1036 a = []
1037 self.makecycle(groupby([a]*2, lambda x:x), a)
1038
Raymond Hettingera1ca94a2008-03-06 22:51:36 +00001039 def test_issue2246(self):
1040 # Issue 2246 -- the _grouper iterator was not included in GC
1041 n = 10
1042 keyfunc = lambda x: x
1043 for i, j in groupby(xrange(n), key=keyfunc):
1044 keyfunc.__dict__.setdefault('x',[]).append(j)
1045
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001046 def test_ifilter(self):
1047 a = []
1048 self.makecycle(ifilter(lambda x:True, [a]*2), a)
1049
1050 def test_ifilterfalse(self):
1051 a = []
1052 self.makecycle(ifilterfalse(lambda x:False, a), a)
1053
1054 def test_izip(self):
1055 a = []
1056 self.makecycle(izip([a]*2, [a]*3), a)
1057
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001058 def test_izip_longest(self):
1059 a = []
1060 self.makecycle(izip_longest([a]*2, [a]*3), a)
1061 b = [a, None]
1062 self.makecycle(izip_longest([a]*2, [a]*3, fillvalue=b), a)
1063
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001064 def test_imap(self):
1065 a = []
1066 self.makecycle(imap(lambda x:x, [a]*2), a)
1067
1068 def test_islice(self):
1069 a = []
1070 self.makecycle(islice([a]*2, None), a)
1071
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001072 def test_permutations(self):
1073 a = []
1074 self.makecycle(permutations([1,2,a,3], 3), a)
1075
1076 def test_product(self):
1077 a = []
1078 self.makecycle(product([1,2,a,3], repeat=3), a)
1079
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001080 def test_repeat(self):
1081 a = []
1082 self.makecycle(repeat(a), a)
1083
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001084 def test_starmap(self):
1085 a = []
1086 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1087
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001088 def test_takewhile(self):
1089 a = []
1090 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1091
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001092def R(seqn):
1093 'Regular generator'
1094 for i in seqn:
1095 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001096
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001097class G:
1098 'Sequence using __getitem__'
1099 def __init__(self, seqn):
1100 self.seqn = seqn
1101 def __getitem__(self, i):
1102 return self.seqn[i]
1103
1104class I:
1105 'Sequence using iterator protocol'
1106 def __init__(self, seqn):
1107 self.seqn = seqn
1108 self.i = 0
1109 def __iter__(self):
1110 return self
1111 def next(self):
1112 if self.i >= len(self.seqn): raise StopIteration
1113 v = self.seqn[self.i]
1114 self.i += 1
1115 return v
1116
1117class Ig:
1118 'Sequence using iterator protocol defined with a generator'
1119 def __init__(self, seqn):
1120 self.seqn = seqn
1121 self.i = 0
1122 def __iter__(self):
1123 for val in self.seqn:
1124 yield val
1125
1126class X:
1127 'Missing __getitem__ and __iter__'
1128 def __init__(self, seqn):
1129 self.seqn = seqn
1130 self.i = 0
1131 def next(self):
1132 if self.i >= len(self.seqn): raise StopIteration
1133 v = self.seqn[self.i]
1134 self.i += 1
1135 return v
1136
1137class N:
1138 'Iterator missing next()'
1139 def __init__(self, seqn):
1140 self.seqn = seqn
1141 self.i = 0
1142 def __iter__(self):
1143 return self
1144
1145class E:
1146 'Test propagation of exceptions'
1147 def __init__(self, seqn):
1148 self.seqn = seqn
1149 self.i = 0
1150 def __iter__(self):
1151 return self
1152 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001153 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001154
1155class S:
1156 'Test immediate stop'
1157 def __init__(self, seqn):
1158 pass
1159 def __iter__(self):
1160 return self
1161 def next(self):
1162 raise StopIteration
1163
1164def L(seqn):
1165 'Test multiple tiers of iterators'
1166 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
1167
1168
1169class TestVariousIteratorArgs(unittest.TestCase):
1170
1171 def test_chain(self):
1172 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1173 for g in (G, I, Ig, S, L, R):
1174 self.assertEqual(list(chain(g(s))), list(g(s)))
1175 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001176 self.assertRaises(TypeError, list, chain(X(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001177 self.assertRaises(TypeError, list, chain(N(s)))
1178 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1179
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001180 def test_compress(self):
1181 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1182 n = len(s)
1183 for g in (G, I, Ig, S, L, R):
1184 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1185 self.assertRaises(TypeError, compress, X(s), repeat(1))
1186 self.assertRaises(TypeError, list, compress(N(s), repeat(1)))
1187 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1188
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001189 def test_product(self):
1190 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1191 self.assertRaises(TypeError, product, X(s))
1192 self.assertRaises(TypeError, product, N(s))
1193 self.assertRaises(ZeroDivisionError, product, E(s))
1194
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001195 def test_cycle(self):
1196 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1197 for g in (G, I, Ig, S, L, R):
1198 tgtlen = len(s) * 3
1199 expected = list(g(s))*3
1200 actual = list(islice(cycle(g(s)), tgtlen))
1201 self.assertEqual(actual, expected)
1202 self.assertRaises(TypeError, cycle, X(s))
1203 self.assertRaises(TypeError, list, cycle(N(s)))
1204 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1205
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001206 def test_groupby(self):
1207 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1208 for g in (G, I, Ig, S, L, R):
1209 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1210 self.assertRaises(TypeError, groupby, X(s))
1211 self.assertRaises(TypeError, list, groupby(N(s)))
1212 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1213
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001214 def test_ifilter(self):
1215 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1216 for g in (G, I, Ig, S, L, R):
1217 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
1218 self.assertRaises(TypeError, ifilter, isEven, X(s))
1219 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
1220 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
1221
1222 def test_ifilterfalse(self):
1223 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1224 for g in (G, I, Ig, S, L, R):
1225 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
1226 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
1227 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
1228 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
1229
1230 def test_izip(self):
1231 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1232 for g in (G, I, Ig, S, L, R):
1233 self.assertEqual(list(izip(g(s))), zip(g(s)))
1234 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
1235 self.assertRaises(TypeError, izip, X(s))
1236 self.assertRaises(TypeError, list, izip(N(s)))
1237 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
1238
Raymond Hettingerd36862c2007-02-21 05:20:38 +00001239 def test_iziplongest(self):
1240 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1241 for g in (G, I, Ig, S, L, R):
1242 self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
1243 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
1244 self.assertRaises(TypeError, izip_longest, X(s))
1245 self.assertRaises(TypeError, list, izip_longest(N(s)))
1246 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
1247
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001248 def test_imap(self):
1249 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1250 for g in (G, I, Ig, S, L, R):
1251 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
1252 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
1253 self.assertRaises(TypeError, imap, onearg, X(s))
1254 self.assertRaises(TypeError, list, imap(onearg, N(s)))
1255 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
1256
1257 def test_islice(self):
1258 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1259 for g in (G, I, Ig, S, L, R):
1260 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1261 self.assertRaises(TypeError, islice, X(s), 10)
1262 self.assertRaises(TypeError, list, islice(N(s), 10))
1263 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1264
1265 def test_starmap(self):
1266 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1267 for g in (G, I, Ig, S, L, R):
1268 ss = zip(s, s)
1269 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
1270 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
1271 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
1272 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1273
1274 def test_takewhile(self):
1275 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1276 for g in (G, I, Ig, S, L, R):
1277 tgt = []
1278 for elem in g(s):
1279 if not isEven(elem): break
1280 tgt.append(elem)
1281 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1282 self.assertRaises(TypeError, takewhile, isEven, X(s))
1283 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
1284 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1285
1286 def test_dropwhile(self):
1287 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1288 for g in (G, I, Ig, S, L, R):
1289 tgt = []
1290 for elem in g(s):
1291 if not tgt and isOdd(elem): continue
1292 tgt.append(elem)
1293 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1294 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
1295 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
1296 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1297
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001298 def test_tee(self):
1299 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1300 for g in (G, I, Ig, S, L, R):
1301 it1, it2 = tee(g(s))
1302 self.assertEqual(list(it1), list(g(s)))
1303 self.assertEqual(list(it2), list(g(s)))
1304 self.assertRaises(TypeError, tee, X(s))
1305 self.assertRaises(TypeError, list, tee(N(s))[0])
1306 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1307
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001308class LengthTransparency(unittest.TestCase):
1309
1310 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001311 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001312 self.assertEqual(len(repeat(None, 50)), 50)
1313 self.assertRaises(TypeError, len, repeat(None))
1314
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001315class RegressionTests(unittest.TestCase):
1316
1317 def test_sf_793826(self):
1318 # Fix Armin Rigo's successful efforts to wreak havoc
1319
1320 def mutatingtuple(tuple1, f, tuple2):
1321 # this builds a tuple t which is a copy of tuple1,
1322 # then calls f(t), then mutates t to be equal to tuple2
1323 # (needs len(tuple1) == len(tuple2)).
1324 def g(value, first=[1]):
1325 if first:
1326 del first[:]
1327 f(z.next())
1328 return value
1329 items = list(tuple2)
1330 items[1:1] = list(tuple1)
1331 gen = imap(g, items)
1332 z = izip(*[gen]*len(tuple1))
1333 z.next()
1334
1335 def f(t):
1336 global T
1337 T = t
1338 first[:] = list(T)
1339
1340 first = []
1341 mutatingtuple((1,2,3), f, (4,5,6))
1342 second = list(T)
1343 self.assertEqual(first, second)
1344
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001345
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001346 def test_sf_950057(self):
1347 # Make sure that chain() and cycle() catch exceptions immediately
1348 # rather than when shifting between input sources
1349
1350 def gen1():
1351 hist.append(0)
1352 yield 1
1353 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001354 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001355 hist.append(2)
1356
1357 def gen2(x):
1358 hist.append(3)
1359 yield 2
1360 hist.append(4)
1361 if x:
1362 raise StopIteration
1363
1364 hist = []
1365 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1366 self.assertEqual(hist, [0,1])
1367
1368 hist = []
1369 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1370 self.assertEqual(hist, [0,1])
1371
1372 hist = []
1373 self.assertRaises(AssertionError, list, cycle(gen1()))
1374 self.assertEqual(hist, [0,1])
1375
Georg Brandlb84c1372007-01-21 10:28:43 +00001376class SubclassWithKwargsTest(unittest.TestCase):
1377 def test_keywords_in_subclass(self):
1378 # count is not subclassable...
1379 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001380 starmap, islice, takewhile, dropwhile, cycle, compress):
Georg Brandlb84c1372007-01-21 10:28:43 +00001381 class Subclass(cls):
1382 def __init__(self, newarg=None, *args):
1383 cls.__init__(self, *args)
1384 try:
1385 Subclass(newarg=1)
1386 except TypeError, err:
1387 # we expect type errors because of wrong argument count
Benjamin Peterson5c8da862009-06-30 22:57:08 +00001388 self.assertFalse("does not take keyword arguments" in err.args[0])
Georg Brandlb84c1372007-01-21 10:28:43 +00001389
1390
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001391libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001392
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001393
1394>>> amounts = [120.15, 764.05, 823.14]
1395>>> for checknum, amount in izip(count(1200), amounts):
1396... print 'Check %d is for $%.2f' % (checknum, amount)
1397...
1398Check 1200 is for $120.15
1399Check 1201 is for $764.05
1400Check 1202 is for $823.14
1401
1402>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001403>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1404... print cube
1405...
14061
14078
140827
1409
1410>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001411>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001412... print name.title()
1413...
1414Alex
1415Laura
1416Martin
1417Walter
1418Samuele
1419
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001420>>> from operator import itemgetter
1421>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Armin Rigoa3f09272006-05-28 19:13:17 +00001422>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001423>>> for k, g in groupby(di, itemgetter(1)):
1424... print k, map(itemgetter(0), g)
1425...
14261 ['a', 'c', 'e']
14272 ['b', 'd', 'f']
14283 ['g']
1429
Raymond Hettinger734fb572004-01-20 20:04:40 +00001430# Find runs of consecutive numbers using groupby. The key to the solution
1431# is differencing with a range so that consecutive numbers all appear in
1432# same group.
1433>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
Senthil Kumarance8e33a2010-01-08 19:04:16 +00001434>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
Raymond Hettinger734fb572004-01-20 20:04:40 +00001435... print map(operator.itemgetter(1), g)
1436...
1437[1]
1438[4, 5, 6]
1439[10]
1440[15, 16, 17, 18]
1441[22]
1442[25, 26, 27, 28]
1443
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001444>>> def take(n, iterable):
1445... "Return first n items of the iterable as a list"
1446... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001447
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001448>>> def enumerate(iterable, start=0):
1449... return izip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001450
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001451>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001452... "Return function(0), function(1), ..."
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001453... return imap(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001454
Raymond Hettingerf9bce832009-02-19 05:34:35 +00001455>>> def nth(iterable, n, default=None):
1456... "Returns the nth item or a default value"
1457... return next(islice(iterable, n, None), default)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001458
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001459>>> def quantify(iterable, pred=bool):
1460... "Count how many times the predicate is true"
1461... return sum(imap(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001462
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001463>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001464... "Returns the sequence elements and then returns None indefinitely"
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001465... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001466
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001467>>> def ncycles(iterable, n):
1468... "Returns the seqeuence elements n times"
1469... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001470
1471>>> def dotproduct(vec1, vec2):
1472... return sum(imap(operator.mul, vec1, vec2))
1473
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001474>>> def flatten(listOfLists):
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001475... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001476
1477>>> def repeatfunc(func, times=None, *args):
1478... "Repeat calls to func with specified arguments."
1479... " Example: repeatfunc(random.random)"
1480... if times is None:
1481... return starmap(func, repeat(args))
1482... else:
1483... return starmap(func, repeat(args, times))
1484
Raymond Hettingerd591f662003-10-26 15:34:50 +00001485>>> def pairwise(iterable):
1486... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1487... a, b = tee(iterable)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001488... for elem in b:
1489... break
Raymond Hettingerad983e72003-11-12 14:32:26 +00001490... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001491
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001492>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001493... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001494... args = [iter(iterable)] * n
Raymond Hettingerf080e6d2008-07-31 01:19:50 +00001495... return izip_longest(fillvalue=fillvalue, *args)
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001496
1497>>> def roundrobin(*iterables):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001498... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001499... # Recipe credited to George Sakkis
1500... pending = len(iterables)
1501... nexts = cycle(iter(it).next for it in iterables)
1502... while pending:
1503... try:
1504... for next in nexts:
1505... yield next()
1506... except StopIteration:
1507... pending -= 1
1508... nexts = cycle(islice(nexts, pending))
1509
1510>>> def powerset(iterable):
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001511... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1512... s = list(iterable)
1513... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001514
Raymond Hettinger44e15812009-01-02 21:26:45 +00001515>>> def unique_everseen(iterable, key=None):
1516... "List unique elements, preserving order. Remember all elements ever seen."
1517... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1518... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1519... seen = set()
1520... seen_add = seen.add
1521... if key is None:
1522... for element in iterable:
1523... if element not in seen:
1524... seen_add(element)
1525... yield element
1526... else:
1527... for element in iterable:
1528... k = key(element)
1529... if k not in seen:
1530... seen_add(k)
1531... yield element
1532
1533>>> def unique_justseen(iterable, key=None):
1534... "List unique elements, preserving order. Remember only the element just seen."
1535... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1536... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1537... return imap(next, imap(itemgetter(1), groupby(iterable, key)))
1538
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001539This is not part of the examples but it tests to make sure the definitions
1540perform as purported.
1541
Raymond Hettingera098b332003-09-08 23:58:40 +00001542>>> take(10, count())
1543[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1544
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001545>>> list(enumerate('abc'))
1546[(0, 'a'), (1, 'b'), (2, 'c')]
1547
1548>>> list(islice(tabulate(lambda x: 2*x), 4))
1549[0, 2, 4, 6]
1550
1551>>> nth('abcde', 3)
Raymond Hettingerd507afd2009-02-04 10:52:32 +00001552'd'
1553
1554>>> nth('abcde', 9) is None
1555True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001556
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001557>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000155850
1559
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001560>>> a = [[1, 2, 3], [4, 5, 6]]
1561>>> flatten(a)
1562[1, 2, 3, 4, 5, 6]
1563
1564>>> list(repeatfunc(pow, 5, 2, 3))
1565[8, 8, 8, 8, 8]
1566
1567>>> import random
1568>>> take(5, imap(int, repeatfunc(random.random)))
1569[0, 0, 0, 0, 0]
1570
Raymond Hettingerd591f662003-10-26 15:34:50 +00001571>>> list(pairwise('abcd'))
1572[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001573
Raymond Hettingerd591f662003-10-26 15:34:50 +00001574>>> list(pairwise([]))
1575[]
1576
1577>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001578[]
1579
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001580>>> list(islice(padnone('abc'), 0, 6))
1581['a', 'b', 'c', None, None, None]
1582
1583>>> list(ncycles('abc', 3))
1584['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1585
1586>>> dotproduct([1,2,3], [4,5,6])
158732
1588
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001589>>> list(grouper(3, 'abcdefg', 'x'))
1590[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1591
1592>>> list(roundrobin('abc', 'd', 'ef'))
1593['a', 'd', 'e', 'b', 'f', 'c']
1594
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001595>>> list(powerset([1,2,3]))
1596[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001597
Raymond Hettinger560f9a82009-01-27 13:26:35 +00001598>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1599True
1600
1601>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1602True
1603
Raymond Hettinger44e15812009-01-02 21:26:45 +00001604>>> list(unique_everseen('AAAABBBCCDAABBB'))
1605['A', 'B', 'C', 'D']
1606
1607>>> list(unique_everseen('ABBCcAD', str.lower))
1608['A', 'B', 'C', 'D']
1609
1610>>> list(unique_justseen('AAAABBBCCDAABBB'))
1611['A', 'B', 'C', 'D', 'A', 'B']
1612
1613>>> list(unique_justseen('ABBCcAD', str.lower))
1614['A', 'B', 'C', 'A', 'D']
1615
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001616"""
1617
1618__test__ = {'libreftest' : libreftest}
1619
1620def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001621 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Georg Brandlb84c1372007-01-21 10:28:43 +00001622 RegressionTests, LengthTransparency,
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001623 SubclassWithKwargsTest, TestExamples)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001624 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001625
1626 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001627 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001628 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001629 counts = [None] * 5
1630 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001631 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001632 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001633 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001634 print counts
1635
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001636 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001637 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001638
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001639if __name__ == "__main__":
1640 test_main(verbose=True)