blob: 553d7239463fc4e10a89b8a8b05cfa0c5df713bf [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 Hettinger2012f172003-02-07 05:32:58 +00006import sys
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00007import operator
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00008import random
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +00009maxsize = test_support.MAX_Py_ssize_t
10minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000011
12def onearg(x):
13 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000014 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000015
16def errfunc(*args):
17 'Test function that raises an error'
18 raise ValueError
19
20def gen3():
21 'Non-restartable source sequence'
22 for i in (0, 1, 2):
23 yield i
24
25def isEven(x):
26 'Test predicate'
27 return x%2==0
28
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000029def isOdd(x):
30 'Test predicate'
31 return x%2==1
32
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000033class StopNow:
34 'Class emulating an empty iterable.'
35 def __iter__(self):
36 return self
37 def next(self):
38 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000039
Raymond Hettinger02420702003-06-29 20:36:23 +000040def take(n, seq):
41 'Convenience function for partially consuming a long of infinite iterable'
42 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000043
Raymond Hettingerd553d852008-03-04 04:17:08 +000044def prod(iterable):
45 return reduce(operator.mul, iterable, 1)
46
Raymond Hettinger93e804d2008-02-26 23:40:50 +000047def fact(n):
48 'Factorial'
Raymond Hettingerd553d852008-03-04 04:17:08 +000049 return prod(range(1, n+1))
50
Raymond Hettinger96ef8112003-02-01 00:10:11 +000051class TestBasicOps(unittest.TestCase):
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000052 def test_chain(self):
Raymond Hettingerad47fa12008-03-06 20:52:01 +000053
54 def chain2(*iterables):
55 'Pure python version in the docs'
56 for it in iterables:
57 for element in it:
58 yield element
59
60 for c in (chain, chain2):
61 self.assertEqual(list(c('abc', 'def')), list('abcdef'))
62 self.assertEqual(list(c('abc')), list('abc'))
63 self.assertEqual(list(c('')), [])
64 self.assertEqual(take(4, c('abc', 'def')), list('abcd'))
65 self.assertRaises(TypeError, list,c(2, 3))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000066
Raymond Hettingerb4cbc982008-02-28 22:46:41 +000067 def test_chain_from_iterable(self):
68 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
69 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
70 self.assertEqual(list(chain.from_iterable([''])), [])
71 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
72 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
73
Raymond Hettinger93e804d2008-02-26 23:40:50 +000074 def test_combinations(self):
Raymond Hettinger5b913e32009-01-08 06:39:04 +000075 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
Raymond Hettinger93e804d2008-02-26 23:40:50 +000076 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
Raymond Hettingerd553d852008-03-04 04:17:08 +000077 self.assertRaises(TypeError, combinations, None) # pool is not iterable
Raymond Hettinger93e804d2008-02-26 23:40:50 +000078 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
Raymond Hettinger5b913e32009-01-08 06:39:04 +000079 self.assertEqual(list(combinations('abc', 32)), []) # r > n
Raymond Hettinger93e804d2008-02-26 23:40:50 +000080 self.assertEqual(list(combinations(range(4), 3)),
81 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
Raymond Hettingerd553d852008-03-04 04:17:08 +000082
83 def combinations1(iterable, r):
84 'Pure python version shown in the docs'
85 pool = tuple(iterable)
86 n = len(pool)
Raymond Hettinger5b913e32009-01-08 06:39:04 +000087 if r > n:
88 return
Raymond Hettingerd553d852008-03-04 04:17:08 +000089 indices = range(r)
90 yield tuple(pool[i] for i in indices)
91 while 1:
92 for i in reversed(range(r)):
93 if indices[i] != i + n - r:
94 break
95 else:
96 return
97 indices[i] += 1
98 for j in range(i+1, r):
99 indices[j] = indices[j-1] + 1
100 yield tuple(pool[i] for i in indices)
101
102 def combinations2(iterable, r):
103 'Pure python version shown in the docs'
104 pool = tuple(iterable)
105 n = len(pool)
106 for indices in permutations(range(n), r):
107 if sorted(indices) == list(indices):
108 yield tuple(pool[i] for i in indices)
109
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000110 def combinations3(iterable, r):
111 'Pure python version from cwr()'
112 pool = tuple(iterable)
113 n = len(pool)
114 for indices in combinations_with_replacement(range(n), r):
115 if len(set(indices)) == r:
116 yield tuple(pool[i] for i in indices)
117
Raymond Hettingerd553d852008-03-04 04:17:08 +0000118 for n in range(7):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000119 values = [5*x-12 for x in range(n)]
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000120 for r in range(n+2):
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000121 result = list(combinations(values, r))
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000122 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 +0000123 self.assertEqual(len(result), len(set(result))) # no repeats
124 self.assertEqual(result, sorted(result)) # lexicographic order
125 for c in result:
126 self.assertEqual(len(c), r) # r-length combinations
127 self.assertEqual(len(set(c)), r) # no duplicate elements
128 self.assertEqual(list(c), sorted(c)) # keep original ordering
129 self.assert_(all(e in values for e in c)) # elements taken from input iterable
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000130 self.assertEqual(list(c),
131 [e for e in values if e in c]) # comb is a subsequence of the input iterable
Raymond Hettingerd553d852008-03-04 04:17:08 +0000132 self.assertEqual(result, list(combinations1(values, r))) # matches first pure python version
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000133 self.assertEqual(result, list(combinations2(values, r))) # matches second pure python version
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000134 self.assertEqual(result, list(combinations3(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000135
136 # Test implementation detail: tuple re-use
137 self.assertEqual(len(set(map(id, combinations('abcde', 3)))), 1)
138 self.assertNotEqual(len(set(map(id, list(combinations('abcde', 3))))), 1)
139
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000140 def test_combinations_with_replacement(self):
141 cwr = combinations_with_replacement
142 self.assertRaises(TypeError, cwr, 'abc') # missing r argument
143 self.assertRaises(TypeError, cwr, 'abc', 2, 1) # too many arguments
144 self.assertRaises(TypeError, cwr, None) # pool is not iterable
145 self.assertRaises(ValueError, cwr, 'abc', -2) # r is negative
146 self.assertEqual(list(cwr('ABC', 2)),
147 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
148
149 def cwr1(iterable, r):
150 'Pure python version shown in the docs'
151 # number items returned: (n+r-1)! / r! / (n-1)! when n>0
152 pool = tuple(iterable)
153 n = len(pool)
154 if not n and r:
155 return
156 indices = [0] * r
157 yield tuple(pool[i] for i in indices)
158 while 1:
159 for i in reversed(range(r)):
160 if indices[i] != n - 1:
161 break
162 else:
163 return
164 indices[i:] = [indices[i] + 1] * (r - i)
165 yield tuple(pool[i] for i in indices)
166
167 def cwr2(iterable, r):
168 'Pure python version shown in the docs'
169 pool = tuple(iterable)
170 n = len(pool)
171 for indices in product(range(n), repeat=r):
172 if sorted(indices) == list(indices):
173 yield tuple(pool[i] for i in indices)
174
175 def numcombs(n, r):
176 if not n:
177 return 0 if r else 1
178 return fact(n+r-1) / fact(r)/ fact(n-1)
179
180 for n in range(7):
181 values = [5*x-12 for x in range(n)]
182 for r in range(n+2):
183 result = list(cwr(values, r))
184
185 self.assertEqual(len(result), numcombs(n, r)) # right number of combs
186 self.assertEqual(len(result), len(set(result))) # no repeats
187 self.assertEqual(result, sorted(result)) # lexicographic order
188
189 regular_combs = list(combinations(values, r)) # compare to combs without replacement
190 if n == 0 or r <= 1:
191 self.assertEquals(result, regular_combs) # cases that should be identical
192 else:
193 self.assert_(set(result) >= set(regular_combs)) # rest should be supersets of regular combs
194
195 for c in result:
196 self.assertEqual(len(c), r) # r-length combinations
197 noruns = [k for k,v in groupby(c)] # combo without consecutive repeats
198 self.assertEqual(len(noruns), len(set(noruns))) # no repeats other than consecutive
199 self.assertEqual(list(c), sorted(c)) # keep original ordering
200 self.assert_(all(e in values for e in c)) # elements taken from input iterable
201 self.assertEqual(noruns,
202 [e for e in values if e in c]) # comb is a subsequence of the input iterable
203 self.assertEqual(result, list(cwr1(values, r))) # matches first pure python version
204 self.assertEqual(result, list(cwr2(values, r))) # matches second pure python version
205
206 # Test implementation detail: tuple re-use
207 self.assertEqual(len(set(map(id, cwr('abcde', 3)))), 1)
208 self.assertNotEqual(len(set(map(id, list(cwr('abcde', 3))))), 1)
209
Raymond Hettingerd553d852008-03-04 04:17:08 +0000210 def test_permutations(self):
211 self.assertRaises(TypeError, permutations) # too few arguments
212 self.assertRaises(TypeError, permutations, 'abc', 2, 1) # too many arguments
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000213 self.assertRaises(TypeError, permutations, None) # pool is not iterable
214 self.assertRaises(ValueError, permutations, 'abc', -2) # r is negative
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000215 self.assertEqual(list(permutations('abc', 32)), []) # r > n
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000216 self.assertRaises(TypeError, permutations, 'abc', 's') # r is not an int or None
Raymond Hettingerd553d852008-03-04 04:17:08 +0000217 self.assertEqual(list(permutations(range(3), 2)),
218 [(0,1), (0,2), (1,0), (1,2), (2,0), (2,1)])
219
220 def permutations1(iterable, r=None):
221 'Pure python version shown in the docs'
222 pool = tuple(iterable)
223 n = len(pool)
224 r = n if r is None else r
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000225 if r > n:
226 return
Raymond Hettingerd553d852008-03-04 04:17:08 +0000227 indices = range(n)
Raymond Hettingere70bb8d2008-03-23 00:55:46 +0000228 cycles = range(n, n-r, -1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000229 yield tuple(pool[i] for i in indices[:r])
230 while n:
231 for i in reversed(range(r)):
232 cycles[i] -= 1
233 if cycles[i] == 0:
234 indices[i:] = indices[i+1:] + indices[i:i+1]
235 cycles[i] = n - i
236 else:
237 j = cycles[i]
238 indices[i], indices[-j] = indices[-j], indices[i]
239 yield tuple(pool[i] for i in indices[:r])
240 break
241 else:
242 return
243
244 def permutations2(iterable, r=None):
245 'Pure python version shown in the docs'
246 pool = tuple(iterable)
247 n = len(pool)
248 r = n if r is None else r
249 for indices in product(range(n), repeat=r):
250 if len(set(indices)) == r:
251 yield tuple(pool[i] for i in indices)
252
253 for n in range(7):
254 values = [5*x-12 for x in range(n)]
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000255 for r in range(n+2):
Raymond Hettingerd553d852008-03-04 04:17:08 +0000256 result = list(permutations(values, r))
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000257 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 +0000258 self.assertEqual(len(result), len(set(result))) # no repeats
259 self.assertEqual(result, sorted(result)) # lexicographic order
260 for p in result:
261 self.assertEqual(len(p), r) # r-length permutations
262 self.assertEqual(len(set(p)), r) # no duplicate elements
263 self.assert_(all(e in values for e in p)) # elements taken from input iterable
264 self.assertEqual(result, list(permutations1(values, r))) # matches first pure python version
Raymond Hettinger5b913e32009-01-08 06:39:04 +0000265 self.assertEqual(result, list(permutations2(values, r))) # matches second pure python version
Raymond Hettingerd553d852008-03-04 04:17:08 +0000266 if r == n:
267 self.assertEqual(result, list(permutations(values, None))) # test r as None
268 self.assertEqual(result, list(permutations(values))) # test default r
269
270 # Test implementation detail: tuple re-use
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000271 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000272 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
Raymond Hettinger93e804d2008-02-26 23:40:50 +0000273
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000274 def test_combinatorics(self):
275 # Test relationships between product(), permutations(),
276 # combinations() and combinations_with_replacement().
277
Raymond Hettinger2f6c2e02009-01-27 10:36:14 +0000278 for n in range(6):
279 s = 'ABCDEFG'[:n]
280 for r in range(8):
281 prod = list(product(s, repeat=r))
282 cwr = list(combinations_with_replacement(s, r))
283 perm = list(permutations(s, r))
284 comb = list(combinations(s, r))
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000285
Raymond Hettinger2f6c2e02009-01-27 10:36:14 +0000286 # Check size
287 self.assertEquals(len(prod), n**r)
288 self.assertEquals(len(cwr), (fact(n+r-1) / fact(r)/ fact(n-1)) if n else (not r))
289 self.assertEquals(len(perm), 0 if r>n else fact(n) / fact(n-r))
290 self.assertEquals(len(comb), 0 if r>n else fact(n) / fact(r) / fact(n-r))
291
292 # Check lexicographic order without repeated tuples
293 self.assertEquals(prod, sorted(set(prod)))
294 self.assertEquals(cwr, sorted(set(cwr)))
295 self.assertEquals(perm, sorted(set(perm)))
296 self.assertEquals(comb, sorted(set(comb)))
297
298 # Check interrelationships
299 self.assertEquals(cwr, [t for t in prod if sorted(t)==list(t)]) # cwr: prods which are sorted
300 self.assertEquals(perm, [t for t in prod if len(set(t))==r]) # perm: prods with no dups
301 self.assertEqual(comb, [t for t in perm if sorted(t)==list(t)]) # comb: perms that are sorted
302 self.assertEqual(comb, [t for t in cwr if len(set(t))==r]) # comb: cwrs without dups
303 self.assertEqual(comb, filter(set(cwr).__contains__, perm)) # comb: perm that is a cwr
304 self.assertEqual(comb, filter(set(perm).__contains__, cwr)) # comb: cwr that is a perm
305 self.assertEqual(comb, sorted(set(cwr) & set(perm))) # comb: both a cwr and a perm
Raymond Hettinger2976aaa2009-01-27 09:33:06 +0000306
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000307 def test_compress(self):
308 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
309 self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
310 self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
311 self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
312 self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
313 n = 10000
314 data = chain.from_iterable(repeat(range(6), n))
315 selectors = chain.from_iterable(repeat((0, 1)))
316 self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
317 self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
318 self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
319 self.assertRaises(TypeError, compress, range(6)) # too few args
320 self.assertRaises(TypeError, compress, range(6), None) # too many args
321
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000322 def test_count(self):
323 self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
324 self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000325 self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000326 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
327 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000328 self.assertRaises(TypeError, count, 2, 3, 4)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000329 self.assertRaises(TypeError, count, 'a')
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000330 self.assertEqual(list(islice(count(maxsize-5), 10)), range(maxsize-5, maxsize+5))
331 self.assertEqual(list(islice(count(-maxsize-5), 10)), range(-maxsize-5, -maxsize+5))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000332 c = count(3)
333 self.assertEqual(repr(c), 'count(3)')
334 c.next()
335 self.assertEqual(repr(c), 'count(4)')
Jack Diederich36234e82006-09-21 17:50:26 +0000336 c = count(-9)
337 self.assertEqual(repr(c), 'count(-9)')
338 c.next()
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000339 self.assertEqual(repr(count(10.25)), 'count(10.25)')
Jack Diederich36234e82006-09-21 17:50:26 +0000340 self.assertEqual(c.next(), -8)
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000341 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 +0000342 # Test repr (ignoring the L in longs)
343 r1 = repr(count(i)).replace('L', '')
344 r2 = 'count(%r)'.__mod__(i).replace('L', '')
345 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000346
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000347 def test_count_with_stride(self):
348 self.assertEqual(zip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
349 self.assertEqual(zip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
350 self.assertEqual(zip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
351 self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3)))
352 self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3)))
353 self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])
Raymond Hettingeraa044612009-02-12 12:04:26 +0000354 self.assertEqual(take(3, count(Decimal('1.1'), Decimal('.1'))),
355 [Decimal('1.1'), Decimal('1.2'), Decimal('1.3')])
Raymond Hettinger31c769c2009-02-12 05:39:46 +0000356 self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0]))
357 c = count(3, 5)
358 self.assertEqual(repr(c), 'count(3, 5)')
359 c.next()
360 self.assertEqual(repr(c), 'count(8, 5)')
361 c = count(-9, 0)
362 self.assertEqual(repr(c), 'count(-9, 0)')
363 c.next()
364 self.assertEqual(repr(c), 'count(-9, 0)')
365 c = count(-9, -3)
366 self.assertEqual(repr(c), 'count(-9, -3)')
367 c.next()
368 self.assertEqual(repr(c), 'count(-12, -3)')
369 self.assertEqual(repr(c), 'count(-12, -3)')
370 self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)')
371 self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int
372 self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0
373 for i in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 10, sys.maxint-5, sys.maxint+5):
374 for j in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 1, 10, sys.maxint-5, sys.maxint+5):
375 # Test repr (ignoring the L in longs)
376 r1 = repr(count(i, j)).replace('L', '')
377 if j == 1:
378 r2 = ('count(%r)' % i).replace('L', '')
379 else:
380 r2 = ('count(%r, %r)' % (i, j)).replace('L', '')
381 self.assertEqual(r1, r2)
382
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000383 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000384 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000385 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000386 self.assertRaises(TypeError, cycle)
387 self.assertRaises(TypeError, cycle, 5)
388 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000389
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000390 def test_groupby(self):
391 # Check whether it accepts arguments correctly
392 self.assertEqual([], list(groupby([])))
393 self.assertEqual([], list(groupby([], key=id)))
394 self.assertRaises(TypeError, list, groupby('abc', []))
395 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000396 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000397
398 # Check normal input
399 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
400 (2,15,22), (3,16,23), (3,17,23)]
401 dup = []
402 for k, g in groupby(s, lambda r:r[0]):
403 for elem in g:
404 self.assertEqual(k, elem[0])
405 dup.append(elem)
406 self.assertEqual(s, dup)
407
408 # Check nested case
409 dup = []
410 for k, g in groupby(s, lambda r:r[0]):
411 for ik, ig in groupby(g, lambda r:r[2]):
412 for elem in ig:
413 self.assertEqual(k, elem[0])
414 self.assertEqual(ik, elem[2])
415 dup.append(elem)
416 self.assertEqual(s, dup)
417
418 # Check case where inner iterator is not used
419 keys = [k for k, g in groupby(s, lambda r:r[0])]
420 expectedkeys = set([r[0] for r in s])
421 self.assertEqual(set(keys), expectedkeys)
422 self.assertEqual(len(keys), len(expectedkeys))
423
424 # Exercise pipes and filters style
425 s = 'abracadabra'
426 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000427 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000428 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
429 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000430 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000431 self.assertEqual(r, ['a', 'b', 'r'])
432 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000433 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000434 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
435 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000436 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000437 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
438
439 # iter.next failure
440 class ExpectedError(Exception):
441 pass
442 def delayed_raise(n=0):
443 for i in range(n):
444 yield 'yo'
445 raise ExpectedError
446 def gulp(iterable, keyp=None, func=list):
447 return [func(g) for k, g in groupby(iterable, keyp)]
448
449 # iter.next failure on outer object
450 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
451 # iter.next failure on inner object
452 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
453
454 # __cmp__ failure
455 class DummyCmp:
456 def __cmp__(self, dst):
457 raise ExpectedError
458 s = [DummyCmp(), DummyCmp(), None]
459
460 # __cmp__ failure on outer object
461 self.assertRaises(ExpectedError, gulp, s, func=id)
462 # __cmp__ failure on inner object
463 self.assertRaises(ExpectedError, gulp, s)
464
465 # keyfunc failure
466 def keyfunc(obj):
467 if keyfunc.skip > 0:
468 keyfunc.skip -= 1
469 return obj
470 else:
471 raise ExpectedError
472
473 # keyfunc failure on outer object
474 keyfunc.skip = 0
475 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
476 keyfunc.skip = 1
477 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
478
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000479 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000480 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000481 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000482 self.assertEqual(list(ifilter(bool, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000483 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000484 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000485 self.assertRaises(TypeError, ifilter, lambda x:x)
486 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000487 self.assertRaises(TypeError, ifilter, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000488 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000489
490 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000491 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
492 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000493 self.assertEqual(list(ifilterfalse(bool, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000494 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000495 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000496 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
497 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000498 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000499 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000500
501 def test_izip(self):
502 ans = [(x,y) for x, y in izip('abc',count())]
503 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000504 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
505 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
Raymond Hettinger02420702003-06-29 20:36:23 +0000506 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000507 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000508 self.assertEqual(list(izip()), zip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000509 self.assertRaises(TypeError, izip, 3)
510 self.assertRaises(TypeError, izip, range(3), 3)
511 # Check tuple re-use (implementation detail)
512 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
513 zip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000514 self.assertEqual([pair for pair in izip('abc', 'def')],
515 zip('abc', 'def'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000516 ids = map(id, izip('abc', 'def'))
517 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000518 ids = map(id, list(izip('abc', 'def')))
519 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000520
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000521 def test_iziplongest(self):
522 for args in [
523 ['abc', range(6)],
524 [range(6), 'abc'],
525 [range(1000), range(2000,2100), range(3000,3050)],
526 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
527 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
528 ]:
529 target = map(None, *args)
530 self.assertEqual(list(izip_longest(*args)), target)
531 self.assertEqual(list(izip_longest(*args, **{})), target)
532 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
533 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Tim Petersea5962f2007-03-12 18:07:52 +0000534
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000535 self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
536
537 self.assertEqual(list(izip_longest()), zip())
538 self.assertEqual(list(izip_longest([])), zip([]))
539 self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
Tim Petersea5962f2007-03-12 18:07:52 +0000540
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000541 self.assertEqual(list(izip_longest('abc', 'defg', **{})), map(None, 'abc', 'defg')) # empty keyword dict
542 self.assertRaises(TypeError, izip_longest, 3)
543 self.assertRaises(TypeError, izip_longest, range(3), 3)
544
545 for stmt in [
546 "izip_longest('abc', fv=1)",
Tim Petersea5962f2007-03-12 18:07:52 +0000547 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000548 ]:
549 try:
550 eval(stmt, globals(), locals())
551 except TypeError:
552 pass
553 else:
554 self.fail('Did not raise Type in: ' + stmt)
Tim Petersea5962f2007-03-12 18:07:52 +0000555
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000556 # Check tuple re-use (implementation detail)
557 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
558 zip('abc', 'def'))
559 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
560 zip('abc', 'def'))
561 ids = map(id, izip_longest('abc', 'def'))
562 self.assertEqual(min(ids), max(ids))
563 ids = map(id, list(izip_longest('abc', 'def')))
564 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
565
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000566 def test_product(self):
567 for args, result in [
Raymond Hettingerd553d852008-03-04 04:17:08 +0000568 ([], [()]), # zero iterables
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000569 (['ab'], [('a',), ('b',)]), # one iterable
570 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
571 ([range(0), range(2), range(3)], []), # first iterable with zero length
572 ([range(2), range(0), range(3)], []), # middle iterable with zero length
573 ([range(2), range(3), range(0)], []), # last iterable with zero length
574 ]:
575 self.assertEqual(list(product(*args)), result)
Raymond Hettinger08ff6822008-02-29 02:21:48 +0000576 for r in range(4):
577 self.assertEqual(list(product(*(args*r))),
578 list(product(*args, **dict(repeat=r))))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000579 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
580 self.assertRaises(TypeError, product, range(6), None)
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000581
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000582 def product1(*args, **kwds):
583 pools = map(tuple, args) * kwds.get('repeat', 1)
584 n = len(pools)
585 if n == 0:
586 yield ()
587 return
588 if any(len(pool) == 0 for pool in pools):
589 return
590 indices = [0] * n
591 yield tuple(pool[i] for pool, i in zip(pools, indices))
592 while 1:
593 for i in reversed(range(n)): # right to left
594 if indices[i] == len(pools[i]) - 1:
595 continue
596 indices[i] += 1
597 for j in range(i+1, n):
598 indices[j] = 0
599 yield tuple(pool[i] for pool, i in zip(pools, indices))
600 break
601 else:
602 return
603
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000604 def product2(*args, **kwds):
605 'Pure python version used in docs'
606 pools = map(tuple, args) * kwds.get('repeat', 1)
607 result = [[]]
608 for pool in pools:
609 result = [x+[y] for x in result for y in pool]
610 for prod in result:
611 yield tuple(prod)
612
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000613 argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
614 set('abcdefg'), range(11), tuple(range(13))]
615 for i in range(100):
616 args = [random.choice(argtypes) for j in range(random.randrange(5))]
Raymond Hettingerd553d852008-03-04 04:17:08 +0000617 expected_len = prod(map(len, args))
618 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +0000619 self.assertEqual(list(product(*args)), list(product1(*args)))
Raymond Hettinger66f91ea2008-03-05 20:59:58 +0000620 self.assertEqual(list(product(*args)), list(product2(*args)))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000621 args = map(iter, args)
Raymond Hettingerd553d852008-03-04 04:17:08 +0000622 self.assertEqual(len(list(product(*args))), expected_len)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000623
Raymond Hettinger73d79632008-02-23 02:20:41 +0000624 # Test implementation detail: tuple re-use
625 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
626 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000627
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000628 def test_repeat(self):
629 self.assertEqual(zip(xrange(3),repeat('a')),
630 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000631 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000632 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000633 self.assertEqual(list(repeat('a', 0)), [])
634 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000635 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000636 self.assertRaises(TypeError, repeat, None, 3, 4)
637 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000638 r = repeat(1+0j)
639 self.assertEqual(repr(r), 'repeat((1+0j))')
640 r = repeat(1+0j, 5)
641 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
642 list(r)
643 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000644
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000645 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000646 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
647 [0**1, 1**2, 2**3])
648 self.assertEqual(list(imap(None, 'abc', range(5))),
649 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000650 self.assertEqual(list(imap(None, 'abc', count())),
651 [('a',0),('b',1),('c',2)])
652 self.assertEqual(take(2,imap(None, 'abc', count())),
653 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000654 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000655 self.assertRaises(TypeError, imap)
656 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000657 self.assertRaises(TypeError, imap(10, range(5)).next)
658 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
659 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000660
661 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000662 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
663 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000664 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
665 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000666 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger47317092008-01-17 03:02:14 +0000667 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
668 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000669 self.assertRaises(TypeError, starmap)
670 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
671 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
672 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
673 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000674
675 def test_islice(self):
676 for args in [ # islice(args) should agree with range(args)
677 (10, 20, 3),
678 (10, 3, 20),
679 (10, 20),
680 (10, 3),
681 (20,)
682 ]:
683 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
684
685 for args, tgtargs in [ # Stop when seqn is exhausted
686 ((10, 110, 3), ((10, 100, 3))),
687 ((10, 110), ((10, 100))),
688 ((110,), (100,))
689 ]:
690 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
691
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000692 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000693 self.assertEqual(list(islice(xrange(10), None)), range(10))
Raymond Hettingerb2594052004-12-05 09:25:51 +0000694 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
695 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000696 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
697 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
698
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000699 # Test number of items consumed SF #1171417
700 it = iter(range(10))
701 self.assertEqual(list(islice(it, 3)), range(3))
702 self.assertEqual(list(it), range(3, 10))
703
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000704 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000705 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000706 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
707 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
708 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
709 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
710 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000711 self.assertRaises(ValueError, islice, xrange(10), 'a')
712 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
713 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
714 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
715 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +0000716 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000717
718 def test_takewhile(self):
719 data = [1, 3, 5, 20, 2, 4, 6, 8]
720 underten = lambda x: x<10
721 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000722 self.assertEqual(list(takewhile(underten, [])), [])
723 self.assertRaises(TypeError, takewhile)
724 self.assertRaises(TypeError, takewhile, operator.pow)
725 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
726 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
727 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000728 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
729 self.assertEqual(list(t), [1, 1, 1])
730 self.assertRaises(StopIteration, t.next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000731
732 def test_dropwhile(self):
733 data = [1, 3, 5, 20, 2, 4, 6, 8]
734 underten = lambda x: x<10
735 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000736 self.assertEqual(list(dropwhile(underten, [])), [])
737 self.assertRaises(TypeError, dropwhile)
738 self.assertRaises(TypeError, dropwhile, operator.pow)
739 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
740 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
741 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000742
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000743 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000744 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000745 def irange(n):
746 for i in xrange(n):
747 yield i
748
749 a, b = tee([]) # test empty iterator
750 self.assertEqual(list(a), [])
751 self.assertEqual(list(b), [])
752
753 a, b = tee(irange(n)) # test 100% interleaved
754 self.assertEqual(zip(a,b), zip(range(n),range(n)))
755
756 a, b = tee(irange(n)) # test 0% interleaved
757 self.assertEqual(list(a), range(n))
758 self.assertEqual(list(b), range(n))
759
760 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000761 for i in xrange(100):
762 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000763 del a
764 self.assertEqual(list(b), range(n))
765
766 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000767 for i in xrange(100):
768 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000769 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000770 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000771
772 for j in xrange(5): # test randomly interleaved
773 order = [0]*n + [1]*n
774 random.shuffle(order)
775 lists = ([], [])
776 its = tee(irange(n))
777 for i in order:
778 value = its[i].next()
779 lists[i].append(value)
780 self.assertEqual(lists[0], range(n))
781 self.assertEqual(lists[1], range(n))
782
Raymond Hettingerad983e72003-11-12 14:32:26 +0000783 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000784 self.assertRaises(TypeError, tee)
785 self.assertRaises(TypeError, tee, 3)
786 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000787 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000788
Raymond Hettingerad983e72003-11-12 14:32:26 +0000789 # tee object should be instantiable
790 a, b = tee('abc')
791 c = type(a)('def')
792 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000793
Raymond Hettingerad983e72003-11-12 14:32:26 +0000794 # test long-lagged and multi-way split
795 a, b, c = tee(xrange(2000), 3)
796 for i in xrange(100):
797 self.assertEqual(a.next(), i)
798 self.assertEqual(list(b), range(2000))
799 self.assertEqual([c.next(), c.next()], range(2))
800 self.assertEqual(list(a), range(100,2000))
801 self.assertEqual(list(c), range(2,2000))
802
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000803 # test values of n
804 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Neal Norwitz69e88972006-09-02 02:58:13 +0000805 self.assertRaises(ValueError, tee, [], -1)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000806 for n in xrange(5):
807 result = tee('abc', n)
808 self.assertEqual(type(result), tuple)
809 self.assertEqual(len(result), n)
810 self.assertEqual(map(list, result), [list('abc')]*n)
811
Raymond Hettingerad983e72003-11-12 14:32:26 +0000812 # tee pass-through to copyable iterator
813 a, b = tee('abc')
814 c, d = tee(a)
815 self.assert_(a is c)
816
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000817 # test tee_new
818 t1, t2 = tee('abc')
819 tnew = type(t1)
820 self.assertRaises(TypeError, tnew)
821 self.assertRaises(TypeError, tnew, 10)
822 t3 = tnew(t1)
823 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000824
Raymond Hettingera9f60922004-10-17 16:40:14 +0000825 # test that tee objects are weak referencable
826 a, b = tee(xrange(10))
827 p = proxy(a)
828 self.assertEqual(getattr(p, '__class__'), type(b))
829 del a
830 self.assertRaises(ReferenceError, getattr, p, '__class__')
831
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000832 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000833 self.assertRaises(StopIteration, izip().next)
834
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000835 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000836 self.assertRaises(StopIteration, f([]).next)
837 self.assertRaises(StopIteration, f(StopNow()).next)
838
839 self.assertRaises(StopIteration, islice([], None).next)
840 self.assertRaises(StopIteration, islice(StopNow(), None).next)
841
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000842 p, q = tee([])
843 self.assertRaises(StopIteration, p.next)
844 self.assertRaises(StopIteration, q.next)
845 p, q = tee(StopNow())
846 self.assertRaises(StopIteration, p.next)
847 self.assertRaises(StopIteration, q.next)
848
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000849 self.assertRaises(StopIteration, repeat(None, 0).next)
850
851 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
852 self.assertRaises(StopIteration, f(lambda x:x, []).next)
853 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
854
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000855class TestExamples(unittest.TestCase):
856
857 def test_chain(self):
858 self.assertEqual(''.join(chain('ABC', 'DEF')), 'ABCDEF')
859
860 def test_chain_from_iterable(self):
861 self.assertEqual(''.join(chain.from_iterable(['ABC', 'DEF'])), 'ABCDEF')
862
863 def test_combinations(self):
864 self.assertEqual(list(combinations('ABCD', 2)),
865 [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')])
866 self.assertEqual(list(combinations(range(4), 3)),
867 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
868
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000869 def test_combinations_with_replacement(self):
870 self.assertEqual(list(combinations_with_replacement('ABC', 2)),
871 [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')])
872
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000873 def test_compress(self):
874 self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
875
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000876 def test_count(self):
877 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
878
879 def test_cycle(self):
880 self.assertEqual(list(islice(cycle('ABCD'), 12)), list('ABCDABCDABCD'))
881
882 def test_dropwhile(self):
883 self.assertEqual(list(dropwhile(lambda x: x<5, [1,4,6,4,1])), [6,4,1])
884
885 def test_groupby(self):
886 self.assertEqual([k for k, g in groupby('AAAABBBCCDAABBB')],
887 list('ABCDAB'))
888 self.assertEqual([(list(g)) for k, g in groupby('AAAABBBCCD')],
889 [list('AAAA'), list('BBB'), list('CC'), list('D')])
890
891 def test_ifilter(self):
892 self.assertEqual(list(ifilter(lambda x: x%2, range(10))), [1,3,5,7,9])
893
894 def test_ifilterfalse(self):
895 self.assertEqual(list(ifilterfalse(lambda x: x%2, range(10))), [0,2,4,6,8])
896
897 def test_imap(self):
898 self.assertEqual(list(imap(pow, (2,3,10), (5,2,3))), [32, 9, 1000])
899
900 def test_islice(self):
901 self.assertEqual(list(islice('ABCDEFG', 2)), list('AB'))
902 self.assertEqual(list(islice('ABCDEFG', 2, 4)), list('CD'))
903 self.assertEqual(list(islice('ABCDEFG', 2, None)), list('CDEFG'))
904 self.assertEqual(list(islice('ABCDEFG', 0, None, 2)), list('ACEG'))
905
906 def test_izip(self):
907 self.assertEqual(list(izip('ABCD', 'xy')), [('A', 'x'), ('B', 'y')])
908
909 def test_izip_longest(self):
910 self.assertEqual(list(izip_longest('ABCD', 'xy', fillvalue='-')),
911 [('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')])
912
913 def test_permutations(self):
914 self.assertEqual(list(permutations('ABCD', 2)),
915 map(tuple, 'AB AC AD BA BC BD CA CB CD DA DB DC'.split()))
916 self.assertEqual(list(permutations(range(3))),
917 [(0,1,2), (0,2,1), (1,0,2), (1,2,0), (2,0,1), (2,1,0)])
918
919 def test_product(self):
920 self.assertEqual(list(product('ABCD', 'xy')),
921 map(tuple, 'Ax Ay Bx By Cx Cy Dx Dy'.split()))
922 self.assertEqual(list(product(range(2), repeat=3)),
923 [(0,0,0), (0,0,1), (0,1,0), (0,1,1),
924 (1,0,0), (1,0,1), (1,1,0), (1,1,1)])
925
926 def test_repeat(self):
927 self.assertEqual(list(repeat(10, 3)), [10, 10, 10])
928
929 def test_stapmap(self):
930 self.assertEqual(list(starmap(pow, [(2,5), (3,2), (10,3)])),
931 [32, 9, 1000])
932
933 def test_takewhile(self):
934 self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
935
936
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000937class TestGC(unittest.TestCase):
938
939 def makecycle(self, iterator, container):
940 container.append(iterator)
941 iterator.next()
942 del container, iterator
943
944 def test_chain(self):
945 a = []
946 self.makecycle(chain(a), a)
947
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000948 def test_chain_from_iterable(self):
949 a = []
950 self.makecycle(chain.from_iterable([a]), a)
951
952 def test_combinations(self):
953 a = []
954 self.makecycle(combinations([1,2,a,3], 3), a)
955
Raymond Hettingerd081abc2009-01-27 02:58:49 +0000956 def test_combinations_with_replacement(self):
957 a = []
958 self.makecycle(combinations_with_replacement([1,2,a,3], 3), a)
959
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +0000960 def test_compress(self):
961 a = []
962 self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
963
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000964 def test_cycle(self):
965 a = []
966 self.makecycle(cycle([a]*2), a)
967
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000968 def test_dropwhile(self):
969 a = []
970 self.makecycle(dropwhile(bool, [0, a, a]), a)
971
972 def test_groupby(self):
973 a = []
974 self.makecycle(groupby([a]*2, lambda x:x), a)
975
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000976 def test_issue2246(self):
977 # Issue 2246 -- the _grouper iterator was not included in GC
978 n = 10
979 keyfunc = lambda x: x
980 for i, j in groupby(xrange(n), key=keyfunc):
981 keyfunc.__dict__.setdefault('x',[]).append(j)
982
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000983 def test_ifilter(self):
984 a = []
985 self.makecycle(ifilter(lambda x:True, [a]*2), a)
986
987 def test_ifilterfalse(self):
988 a = []
989 self.makecycle(ifilterfalse(lambda x:False, a), a)
990
991 def test_izip(self):
992 a = []
993 self.makecycle(izip([a]*2, [a]*3), a)
994
Raymond Hettingerad47fa12008-03-06 20:52:01 +0000995 def test_izip_longest(self):
996 a = []
997 self.makecycle(izip_longest([a]*2, [a]*3), a)
998 b = [a, None]
999 self.makecycle(izip_longest([a]*2, [a]*3, fillvalue=b), a)
1000
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001001 def test_imap(self):
1002 a = []
1003 self.makecycle(imap(lambda x:x, [a]*2), a)
1004
1005 def test_islice(self):
1006 a = []
1007 self.makecycle(islice([a]*2, None), a)
1008
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001009 def test_permutations(self):
1010 a = []
1011 self.makecycle(permutations([1,2,a,3], 3), a)
1012
1013 def test_product(self):
1014 a = []
1015 self.makecycle(product([1,2,a,3], repeat=3), a)
1016
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001017 def test_repeat(self):
1018 a = []
1019 self.makecycle(repeat(a), a)
1020
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001021 def test_starmap(self):
1022 a = []
1023 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
1024
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +00001025 def test_takewhile(self):
1026 a = []
1027 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
1028
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001029def R(seqn):
1030 'Regular generator'
1031 for i in seqn:
1032 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001033
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001034class G:
1035 'Sequence using __getitem__'
1036 def __init__(self, seqn):
1037 self.seqn = seqn
1038 def __getitem__(self, i):
1039 return self.seqn[i]
1040
1041class I:
1042 'Sequence using iterator protocol'
1043 def __init__(self, seqn):
1044 self.seqn = seqn
1045 self.i = 0
1046 def __iter__(self):
1047 return self
1048 def next(self):
1049 if self.i >= len(self.seqn): raise StopIteration
1050 v = self.seqn[self.i]
1051 self.i += 1
1052 return v
1053
1054class Ig:
1055 'Sequence using iterator protocol defined with a generator'
1056 def __init__(self, seqn):
1057 self.seqn = seqn
1058 self.i = 0
1059 def __iter__(self):
1060 for val in self.seqn:
1061 yield val
1062
1063class X:
1064 'Missing __getitem__ and __iter__'
1065 def __init__(self, seqn):
1066 self.seqn = seqn
1067 self.i = 0
1068 def next(self):
1069 if self.i >= len(self.seqn): raise StopIteration
1070 v = self.seqn[self.i]
1071 self.i += 1
1072 return v
1073
1074class N:
1075 'Iterator missing next()'
1076 def __init__(self, seqn):
1077 self.seqn = seqn
1078 self.i = 0
1079 def __iter__(self):
1080 return self
1081
1082class E:
1083 'Test propagation of exceptions'
1084 def __init__(self, seqn):
1085 self.seqn = seqn
1086 self.i = 0
1087 def __iter__(self):
1088 return self
1089 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +00001090 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001091
1092class S:
1093 'Test immediate stop'
1094 def __init__(self, seqn):
1095 pass
1096 def __iter__(self):
1097 return self
1098 def next(self):
1099 raise StopIteration
1100
1101def L(seqn):
1102 'Test multiple tiers of iterators'
1103 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
1104
1105
1106class TestVariousIteratorArgs(unittest.TestCase):
1107
1108 def test_chain(self):
1109 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1110 for g in (G, I, Ig, S, L, R):
1111 self.assertEqual(list(chain(g(s))), list(g(s)))
1112 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001113 self.assertRaises(TypeError, list, chain(X(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001114 self.assertRaises(TypeError, list, chain(N(s)))
1115 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
1116
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001117 def test_compress(self):
1118 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1119 n = len(s)
1120 for g in (G, I, Ig, S, L, R):
1121 self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
1122 self.assertRaises(TypeError, compress, X(s), repeat(1))
1123 self.assertRaises(TypeError, list, compress(N(s), repeat(1)))
1124 self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
1125
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001126 def test_product(self):
1127 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1128 self.assertRaises(TypeError, product, X(s))
1129 self.assertRaises(TypeError, product, N(s))
1130 self.assertRaises(ZeroDivisionError, product, E(s))
1131
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001132 def test_cycle(self):
1133 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1134 for g in (G, I, Ig, S, L, R):
1135 tgtlen = len(s) * 3
1136 expected = list(g(s))*3
1137 actual = list(islice(cycle(g(s)), tgtlen))
1138 self.assertEqual(actual, expected)
1139 self.assertRaises(TypeError, cycle, X(s))
1140 self.assertRaises(TypeError, list, cycle(N(s)))
1141 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
1142
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001143 def test_groupby(self):
1144 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1145 for g in (G, I, Ig, S, L, R):
1146 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
1147 self.assertRaises(TypeError, groupby, X(s))
1148 self.assertRaises(TypeError, list, groupby(N(s)))
1149 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
1150
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001151 def test_ifilter(self):
1152 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1153 for g in (G, I, Ig, S, L, R):
1154 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
1155 self.assertRaises(TypeError, ifilter, isEven, X(s))
1156 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
1157 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
1158
1159 def test_ifilterfalse(self):
1160 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1161 for g in (G, I, Ig, S, L, R):
1162 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
1163 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
1164 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
1165 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
1166
1167 def test_izip(self):
1168 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1169 for g in (G, I, Ig, S, L, R):
1170 self.assertEqual(list(izip(g(s))), zip(g(s)))
1171 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
1172 self.assertRaises(TypeError, izip, X(s))
1173 self.assertRaises(TypeError, list, izip(N(s)))
1174 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
1175
Raymond Hettingerd36862c2007-02-21 05:20:38 +00001176 def test_iziplongest(self):
1177 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1178 for g in (G, I, Ig, S, L, R):
1179 self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
1180 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
1181 self.assertRaises(TypeError, izip_longest, X(s))
1182 self.assertRaises(TypeError, list, izip_longest(N(s)))
1183 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
1184
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001185 def test_imap(self):
1186 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1187 for g in (G, I, Ig, S, L, R):
1188 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
1189 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
1190 self.assertRaises(TypeError, imap, onearg, X(s))
1191 self.assertRaises(TypeError, list, imap(onearg, N(s)))
1192 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
1193
1194 def test_islice(self):
1195 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1196 for g in (G, I, Ig, S, L, R):
1197 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
1198 self.assertRaises(TypeError, islice, X(s), 10)
1199 self.assertRaises(TypeError, list, islice(N(s), 10))
1200 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
1201
1202 def test_starmap(self):
1203 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
1204 for g in (G, I, Ig, S, L, R):
1205 ss = zip(s, s)
1206 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
1207 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
1208 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
1209 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
1210
1211 def test_takewhile(self):
1212 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1213 for g in (G, I, Ig, S, L, R):
1214 tgt = []
1215 for elem in g(s):
1216 if not isEven(elem): break
1217 tgt.append(elem)
1218 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
1219 self.assertRaises(TypeError, takewhile, isEven, X(s))
1220 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
1221 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
1222
1223 def test_dropwhile(self):
1224 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
1225 for g in (G, I, Ig, S, L, R):
1226 tgt = []
1227 for elem in g(s):
1228 if not tgt and isOdd(elem): continue
1229 tgt.append(elem)
1230 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
1231 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
1232 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
1233 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
1234
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001235 def test_tee(self):
1236 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
1237 for g in (G, I, Ig, S, L, R):
1238 it1, it2 = tee(g(s))
1239 self.assertEqual(list(it1), list(g(s)))
1240 self.assertEqual(list(it2), list(g(s)))
1241 self.assertRaises(TypeError, tee, X(s))
1242 self.assertRaises(TypeError, list, tee(N(s))[0])
1243 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
1244
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001245class LengthTransparency(unittest.TestCase):
1246
1247 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00001248 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00001249 self.assertEqual(len(repeat(None, 50)), 50)
1250 self.assertRaises(TypeError, len, repeat(None))
1251
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001252class RegressionTests(unittest.TestCase):
1253
1254 def test_sf_793826(self):
1255 # Fix Armin Rigo's successful efforts to wreak havoc
1256
1257 def mutatingtuple(tuple1, f, tuple2):
1258 # this builds a tuple t which is a copy of tuple1,
1259 # then calls f(t), then mutates t to be equal to tuple2
1260 # (needs len(tuple1) == len(tuple2)).
1261 def g(value, first=[1]):
1262 if first:
1263 del first[:]
1264 f(z.next())
1265 return value
1266 items = list(tuple2)
1267 items[1:1] = list(tuple1)
1268 gen = imap(g, items)
1269 z = izip(*[gen]*len(tuple1))
1270 z.next()
1271
1272 def f(t):
1273 global T
1274 T = t
1275 first[:] = list(T)
1276
1277 first = []
1278 mutatingtuple((1,2,3), f, (4,5,6))
1279 second = list(T)
1280 self.assertEqual(first, second)
1281
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001282
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001283 def test_sf_950057(self):
1284 # Make sure that chain() and cycle() catch exceptions immediately
1285 # rather than when shifting between input sources
1286
1287 def gen1():
1288 hist.append(0)
1289 yield 1
1290 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +00001291 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001292 hist.append(2)
1293
1294 def gen2(x):
1295 hist.append(3)
1296 yield 2
1297 hist.append(4)
1298 if x:
1299 raise StopIteration
1300
1301 hist = []
1302 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
1303 self.assertEqual(hist, [0,1])
1304
1305 hist = []
1306 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
1307 self.assertEqual(hist, [0,1])
1308
1309 hist = []
1310 self.assertRaises(AssertionError, list, cycle(gen1()))
1311 self.assertEqual(hist, [0,1])
1312
Georg Brandlb84c1372007-01-21 10:28:43 +00001313class SubclassWithKwargsTest(unittest.TestCase):
1314 def test_keywords_in_subclass(self):
1315 # count is not subclassable...
1316 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
Raymond Hettinger2bcb8e92009-01-25 21:04:14 +00001317 starmap, islice, takewhile, dropwhile, cycle, compress):
Georg Brandlb84c1372007-01-21 10:28:43 +00001318 class Subclass(cls):
1319 def __init__(self, newarg=None, *args):
1320 cls.__init__(self, *args)
1321 try:
1322 Subclass(newarg=1)
1323 except TypeError, err:
1324 # we expect type errors because of wrong argument count
1325 self.failIf("does not take keyword arguments" in err.args[0])
1326
1327
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001328libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001329
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001330
1331>>> amounts = [120.15, 764.05, 823.14]
1332>>> for checknum, amount in izip(count(1200), amounts):
1333... print 'Check %d is for $%.2f' % (checknum, amount)
1334...
1335Check 1200 is for $120.15
1336Check 1201 is for $764.05
1337Check 1202 is for $823.14
1338
1339>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001340>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
1341... print cube
1342...
13431
13448
134527
1346
1347>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +00001348>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001349... print name.title()
1350...
1351Alex
1352Laura
1353Martin
1354Walter
1355Samuele
1356
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001357>>> from operator import itemgetter
1358>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Armin Rigoa3f09272006-05-28 19:13:17 +00001359>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00001360>>> for k, g in groupby(di, itemgetter(1)):
1361... print k, map(itemgetter(0), g)
1362...
13631 ['a', 'c', 'e']
13642 ['b', 'd', 'f']
13653 ['g']
1366
Raymond Hettinger734fb572004-01-20 20:04:40 +00001367# Find runs of consecutive numbers using groupby. The key to the solution
1368# is differencing with a range so that consecutive numbers all appear in
1369# same group.
1370>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
1371>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
1372... print map(operator.itemgetter(1), g)
1373...
1374[1]
1375[4, 5, 6]
1376[10]
1377[15, 16, 17, 18]
1378[22]
1379[25, 26, 27, 28]
1380
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001381>>> def take(n, iterable):
1382... "Return first n items of the iterable as a list"
1383... return list(islice(iterable, n))
Raymond Hettingera098b332003-09-08 23:58:40 +00001384
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001385>>> def enumerate(iterable, start=0):
1386... return izip(count(start), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001387
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001388>>> def tabulate(function, start=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001389... "Return function(0), function(1), ..."
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001390... return imap(function, count(start))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001391
1392>>> def nth(iterable, n):
Raymond Hettingerd507afd2009-02-04 10:52:32 +00001393... "Returns the nth item or None"
1394... return next(islice(iterable, n, None), None)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001395
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001396>>> def quantify(iterable, pred=bool):
1397... "Count how many times the predicate is true"
1398... return sum(imap(pred, iterable))
Raymond Hettinger60eca932003-02-09 06:40:58 +00001399
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001400>>> def padnone(iterable):
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001401... "Returns the sequence elements and then returns None indefinitely"
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001402... return chain(iterable, repeat(None))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001403
Raymond Hettingerf1f46f02008-07-19 23:58:47 +00001404>>> def ncycles(iterable, n):
1405... "Returns the seqeuence elements n times"
1406... return chain(*repeat(iterable, n))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001407
1408>>> def dotproduct(vec1, vec2):
1409... return sum(imap(operator.mul, vec1, vec2))
1410
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001411>>> def flatten(listOfLists):
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001412... return list(chain.from_iterable(listOfLists))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001413
1414>>> def repeatfunc(func, times=None, *args):
1415... "Repeat calls to func with specified arguments."
1416... " Example: repeatfunc(random.random)"
1417... if times is None:
1418... return starmap(func, repeat(args))
1419... else:
1420... return starmap(func, repeat(args, times))
1421
Raymond Hettingerd591f662003-10-26 15:34:50 +00001422>>> def pairwise(iterable):
1423... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1424... a, b = tee(iterable)
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001425... for elem in b:
1426... break
Raymond Hettingerad983e72003-11-12 14:32:26 +00001427... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001428
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001429>>> def grouper(n, iterable, fillvalue=None):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001430... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
Raymond Hettinger38fb9be2008-03-07 01:33:20 +00001431... args = [iter(iterable)] * n
Raymond Hettingerf080e6d2008-07-31 01:19:50 +00001432... return izip_longest(fillvalue=fillvalue, *args)
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001433
1434>>> def roundrobin(*iterables):
Raymond Hettingerefdf7062008-07-30 07:27:30 +00001435... "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001436... # Recipe credited to George Sakkis
1437... pending = len(iterables)
1438... nexts = cycle(iter(it).next for it in iterables)
1439... while pending:
1440... try:
1441... for next in nexts:
1442... yield next()
1443... except StopIteration:
1444... pending -= 1
1445... nexts = cycle(islice(nexts, pending))
1446
1447>>> def powerset(iterable):
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001448... "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
1449... s = list(iterable)
1450... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001451
Raymond Hettinger44e15812009-01-02 21:26:45 +00001452>>> def unique_everseen(iterable, key=None):
1453... "List unique elements, preserving order. Remember all elements ever seen."
1454... # unique_everseen('AAAABBBCCDAABBB') --> A B C D
1455... # unique_everseen('ABBCcAD', str.lower) --> A B C D
1456... seen = set()
1457... seen_add = seen.add
1458... if key is None:
1459... for element in iterable:
1460... if element not in seen:
1461... seen_add(element)
1462... yield element
1463... else:
1464... for element in iterable:
1465... k = key(element)
1466... if k not in seen:
1467... seen_add(k)
1468... yield element
1469
1470>>> def unique_justseen(iterable, key=None):
1471... "List unique elements, preserving order. Remember only the element just seen."
1472... # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
1473... # unique_justseen('ABBCcAD', str.lower) --> A B C A D
1474... return imap(next, imap(itemgetter(1), groupby(iterable, key)))
1475
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001476This is not part of the examples but it tests to make sure the definitions
1477perform as purported.
1478
Raymond Hettingera098b332003-09-08 23:58:40 +00001479>>> take(10, count())
1480[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1481
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001482>>> list(enumerate('abc'))
1483[(0, 'a'), (1, 'b'), (2, 'c')]
1484
1485>>> list(islice(tabulate(lambda x: 2*x), 4))
1486[0, 2, 4, 6]
1487
1488>>> nth('abcde', 3)
Raymond Hettingerd507afd2009-02-04 10:52:32 +00001489'd'
1490
1491>>> nth('abcde', 9) is None
1492True
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001493
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001494>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000149550
1496
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001497>>> a = [[1, 2, 3], [4, 5, 6]]
1498>>> flatten(a)
1499[1, 2, 3, 4, 5, 6]
1500
1501>>> list(repeatfunc(pow, 5, 2, 3))
1502[8, 8, 8, 8, 8]
1503
1504>>> import random
1505>>> take(5, imap(int, repeatfunc(random.random)))
1506[0, 0, 0, 0, 0]
1507
Raymond Hettingerd591f662003-10-26 15:34:50 +00001508>>> list(pairwise('abcd'))
1509[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001510
Raymond Hettingerd591f662003-10-26 15:34:50 +00001511>>> list(pairwise([]))
1512[]
1513
1514>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001515[]
1516
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001517>>> list(islice(padnone('abc'), 0, 6))
1518['a', 'b', 'c', None, None, None]
1519
1520>>> list(ncycles('abc', 3))
1521['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1522
1523>>> dotproduct([1,2,3], [4,5,6])
152432
1525
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001526>>> list(grouper(3, 'abcdefg', 'x'))
1527[('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
1528
1529>>> list(roundrobin('abc', 'd', 'ef'))
1530['a', 'd', 'e', 'b', 'f', 'c']
1531
Raymond Hettinger68d919e2009-01-25 21:31:47 +00001532>>> list(powerset([1,2,3]))
1533[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001534
Raymond Hettinger560f9a82009-01-27 13:26:35 +00001535>>> all(len(list(powerset(range(n)))) == 2**n for n in range(18))
1536True
1537
1538>>> list(powerset('abcde')) == sorted(sorted(set(powerset('abcde'))), key=len)
1539True
1540
Raymond Hettinger44e15812009-01-02 21:26:45 +00001541>>> list(unique_everseen('AAAABBBCCDAABBB'))
1542['A', 'B', 'C', 'D']
1543
1544>>> list(unique_everseen('ABBCcAD', str.lower))
1545['A', 'B', 'C', 'D']
1546
1547>>> list(unique_justseen('AAAABBBCCDAABBB'))
1548['A', 'B', 'C', 'D', 'A', 'B']
1549
1550>>> list(unique_justseen('ABBCcAD', str.lower))
1551['A', 'B', 'C', 'A', 'D']
1552
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001553"""
1554
1555__test__ = {'libreftest' : libreftest}
1556
1557def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001558 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Georg Brandlb84c1372007-01-21 10:28:43 +00001559 RegressionTests, LengthTransparency,
Raymond Hettingerad47fa12008-03-06 20:52:01 +00001560 SubclassWithKwargsTest, TestExamples)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001561 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001562
1563 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001564 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001565 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001566 counts = [None] * 5
1567 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001568 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001569 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001570 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001571 print counts
1572
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001573 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001574 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001575
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001576if __name__ == "__main__":
1577 test_main(verbose=True)