blob: 087570c93f1a46222fdf5b7b1cbffa32f6ba98c1 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001import unittest
2from test import test_support
3from itertools import *
Raymond Hettingera9f60922004-10-17 16:40:14 +00004from weakref import proxy
Raymond Hettinger2012f172003-02-07 05:32:58 +00005import sys
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00006import operator
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00007import random
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +00008maxsize = test_support.MAX_Py_ssize_t
9minsize = -maxsize-1
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000010
11def onearg(x):
12 'Test function of one argument'
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000013 return 2*x
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000014
15def errfunc(*args):
16 'Test function that raises an error'
17 raise ValueError
18
19def gen3():
20 'Non-restartable source sequence'
21 for i in (0, 1, 2):
22 yield i
23
24def isEven(x):
25 'Test predicate'
26 return x%2==0
27
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000028def isOdd(x):
29 'Test predicate'
30 return x%2==1
31
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000032class StopNow:
33 'Class emulating an empty iterable.'
34 def __iter__(self):
35 return self
36 def next(self):
37 raise StopIteration
Raymond Hettinger96ef8112003-02-01 00:10:11 +000038
Raymond Hettinger02420702003-06-29 20:36:23 +000039def take(n, seq):
40 'Convenience function for partially consuming a long of infinite iterable'
41 return list(islice(seq, n))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +000042
Raymond Hettinger93e804d2008-02-26 23:40:50 +000043def fact(n):
44 'Factorial'
45 return reduce(operator.mul, range(1, n+1), 1)
46
Raymond Hettinger96ef8112003-02-01 00:10:11 +000047class TestBasicOps(unittest.TestCase):
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000048 def test_chain(self):
49 self.assertEqual(list(chain('abc', 'def')), list('abcdef'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000050 self.assertEqual(list(chain('abc')), list('abc'))
51 self.assertEqual(list(chain('')), [])
Raymond Hettinger02420702003-06-29 20:36:23 +000052 self.assertEqual(take(4, chain('abc', 'def')), list('abcd'))
Raymond Hettinger05bf6332008-02-28 22:30:42 +000053 self.assertRaises(TypeError, list,chain(2, 3))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000054
Raymond Hettingerb4cbc982008-02-28 22:46:41 +000055 def test_chain_from_iterable(self):
56 self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
57 self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
58 self.assertEqual(list(chain.from_iterable([''])), [])
59 self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
60 self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
61
Raymond Hettinger93e804d2008-02-26 23:40:50 +000062 def test_combinations(self):
63 self.assertRaises(TypeError, combinations, 'abc') # missing r argument
64 self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
65 self.assertRaises(ValueError, combinations, 'abc', -2) # r is negative
66 self.assertRaises(ValueError, combinations, 'abc', 32) # r is too big
67 self.assertEqual(list(combinations(range(4), 3)),
68 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
Raymond Hettinger2ab05522008-02-27 01:08:30 +000069 for n in range(8):
Raymond Hettinger93e804d2008-02-26 23:40:50 +000070 values = [5*x-12 for x in range(n)]
71 for r in range(n+1):
72 result = list(combinations(values, r))
73 self.assertEqual(len(result), fact(n) / fact(r) / fact(n-r)) # right number of combs
74 self.assertEqual(len(result), len(set(result))) # no repeats
75 self.assertEqual(result, sorted(result)) # lexicographic order
76 for c in result:
77 self.assertEqual(len(c), r) # r-length combinations
78 self.assertEqual(len(set(c)), r) # no duplicate elements
79 self.assertEqual(list(c), sorted(c)) # keep original ordering
80 self.assert_(all(e in values for e in c)) # elements taken from input iterable
81
Raymond Hettinger96ef8112003-02-01 00:10:11 +000082 def test_count(self):
83 self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
84 self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
Raymond Hettinger02420702003-06-29 20:36:23 +000085 self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
Raymond Hettinger50e90e22007-10-04 00:20:27 +000086 self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)])
87 self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)])
Raymond Hettinger96ef8112003-02-01 00:10:11 +000088 self.assertRaises(TypeError, count, 2, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +000089 self.assertRaises(TypeError, count, 'a')
Raymond Hettinger50e90e22007-10-04 00:20:27 +000090 self.assertEqual(list(islice(count(maxsize-5), 10)), range(maxsize-5, maxsize+5))
91 self.assertEqual(list(islice(count(-maxsize-5), 10)), range(-maxsize-5, -maxsize+5))
Raymond Hettinger4cda01e2004-09-28 04:45:28 +000092 c = count(3)
93 self.assertEqual(repr(c), 'count(3)')
94 c.next()
95 self.assertEqual(repr(c), 'count(4)')
Jack Diederich36234e82006-09-21 17:50:26 +000096 c = count(-9)
97 self.assertEqual(repr(c), 'count(-9)')
98 c.next()
99 self.assertEqual(c.next(), -8)
Raymond Hettinger50e90e22007-10-04 00:20:27 +0000100 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 +0000101 # Test repr (ignoring the L in longs)
102 r1 = repr(count(i)).replace('L', '')
103 r2 = 'count(%r)'.__mod__(i).replace('L', '')
104 self.assertEqual(r1, r2)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000105
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000106 def test_cycle(self):
Raymond Hettinger02420702003-06-29 20:36:23 +0000107 self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000108 self.assertEqual(list(cycle('')), [])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000109 self.assertRaises(TypeError, cycle)
110 self.assertRaises(TypeError, cycle, 5)
111 self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000112
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000113 def test_groupby(self):
114 # Check whether it accepts arguments correctly
115 self.assertEqual([], list(groupby([])))
116 self.assertEqual([], list(groupby([], key=id)))
117 self.assertRaises(TypeError, list, groupby('abc', []))
118 self.assertRaises(TypeError, groupby, None)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000119 self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000120
121 # Check normal input
122 s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
123 (2,15,22), (3,16,23), (3,17,23)]
124 dup = []
125 for k, g in groupby(s, lambda r:r[0]):
126 for elem in g:
127 self.assertEqual(k, elem[0])
128 dup.append(elem)
129 self.assertEqual(s, dup)
130
131 # Check nested case
132 dup = []
133 for k, g in groupby(s, lambda r:r[0]):
134 for ik, ig in groupby(g, lambda r:r[2]):
135 for elem in ig:
136 self.assertEqual(k, elem[0])
137 self.assertEqual(ik, elem[2])
138 dup.append(elem)
139 self.assertEqual(s, dup)
140
141 # Check case where inner iterator is not used
142 keys = [k for k, g in groupby(s, lambda r:r[0])]
143 expectedkeys = set([r[0] for r in s])
144 self.assertEqual(set(keys), expectedkeys)
145 self.assertEqual(len(keys), len(expectedkeys))
146
147 # Exercise pipes and filters style
148 s = 'abracadabra'
149 # sort s | uniq
Raymond Hettinger64958a12003-12-17 20:43:33 +0000150 r = [k for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000151 self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
152 # sort s | uniq -d
Raymond Hettinger64958a12003-12-17 20:43:33 +0000153 r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000154 self.assertEqual(r, ['a', 'b', 'r'])
155 # sort s | uniq -c
Raymond Hettinger64958a12003-12-17 20:43:33 +0000156 r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000157 self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
158 # sort s | uniq -c | sort -rn | head -3
Raymond Hettinger64958a12003-12-17 20:43:33 +0000159 r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000160 self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
161
162 # iter.next failure
163 class ExpectedError(Exception):
164 pass
165 def delayed_raise(n=0):
166 for i in range(n):
167 yield 'yo'
168 raise ExpectedError
169 def gulp(iterable, keyp=None, func=list):
170 return [func(g) for k, g in groupby(iterable, keyp)]
171
172 # iter.next failure on outer object
173 self.assertRaises(ExpectedError, gulp, delayed_raise(0))
174 # iter.next failure on inner object
175 self.assertRaises(ExpectedError, gulp, delayed_raise(1))
176
177 # __cmp__ failure
178 class DummyCmp:
179 def __cmp__(self, dst):
180 raise ExpectedError
181 s = [DummyCmp(), DummyCmp(), None]
182
183 # __cmp__ failure on outer object
184 self.assertRaises(ExpectedError, gulp, s, func=id)
185 # __cmp__ failure on inner object
186 self.assertRaises(ExpectedError, gulp, s)
187
188 # keyfunc failure
189 def keyfunc(obj):
190 if keyfunc.skip > 0:
191 keyfunc.skip -= 1
192 return obj
193 else:
194 raise ExpectedError
195
196 # keyfunc failure on outer object
197 keyfunc.skip = 0
198 self.assertRaises(ExpectedError, gulp, [None], keyfunc)
199 keyfunc.skip = 1
200 self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
201
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000202 def test_ifilter(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000203 self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000204 self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000205 self.assertEqual(list(ifilter(bool, [0,1,0,2,0])), [1,2])
Raymond Hettinger02420702003-06-29 20:36:23 +0000206 self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000207 self.assertRaises(TypeError, ifilter)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000208 self.assertRaises(TypeError, ifilter, lambda x:x)
209 self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000210 self.assertRaises(TypeError, ifilter, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000211 self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000212
213 def test_ifilterfalse(self):
Raymond Hettinger60eca932003-02-09 06:40:58 +0000214 self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
215 self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger9d638372008-02-25 22:42:32 +0000216 self.assertEqual(list(ifilterfalse(bool, [0,1,0,2,0])), [0,0,0])
Raymond Hettinger02420702003-06-29 20:36:23 +0000217 self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
Raymond Hettinger60eca932003-02-09 06:40:58 +0000218 self.assertRaises(TypeError, ifilterfalse)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000219 self.assertRaises(TypeError, ifilterfalse, lambda x:x)
220 self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000221 self.assertRaises(TypeError, ifilterfalse, isEven, 3)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000222 self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000223
224 def test_izip(self):
225 ans = [(x,y) for x, y in izip('abc',count())]
226 self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000227 self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
228 self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
Raymond Hettinger02420702003-06-29 20:36:23 +0000229 self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000230 self.assertEqual(list(izip('abcdef')), zip('abcdef'))
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000231 self.assertEqual(list(izip()), zip())
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000232 self.assertRaises(TypeError, izip, 3)
233 self.assertRaises(TypeError, izip, range(3), 3)
234 # Check tuple re-use (implementation detail)
235 self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
236 zip('abc', 'def'))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000237 self.assertEqual([pair for pair in izip('abc', 'def')],
238 zip('abc', 'def'))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000239 ids = map(id, izip('abc', 'def'))
240 self.assertEqual(min(ids), max(ids))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000241 ids = map(id, list(izip('abc', 'def')))
242 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000243
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000244 def test_iziplongest(self):
245 for args in [
246 ['abc', range(6)],
247 [range(6), 'abc'],
248 [range(1000), range(2000,2100), range(3000,3050)],
249 [range(1000), range(0), range(3000,3050), range(1200), range(1500)],
250 [range(1000), range(0), range(3000,3050), range(1200), range(1500), range(0)],
251 ]:
252 target = map(None, *args)
253 self.assertEqual(list(izip_longest(*args)), target)
254 self.assertEqual(list(izip_longest(*args, **{})), target)
255 target = [tuple((e is None and 'X' or e) for e in t) for t in target] # Replace None fills with 'X'
256 self.assertEqual(list(izip_longest(*args, **dict(fillvalue='X'))), target)
Tim Petersea5962f2007-03-12 18:07:52 +0000257
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000258 self.assertEqual(take(3,izip_longest('abcdef', count())), zip('abcdef', range(3))) # take 3 from infinite input
259
260 self.assertEqual(list(izip_longest()), zip())
261 self.assertEqual(list(izip_longest([])), zip([]))
262 self.assertEqual(list(izip_longest('abcdef')), zip('abcdef'))
Tim Petersea5962f2007-03-12 18:07:52 +0000263
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000264 self.assertEqual(list(izip_longest('abc', 'defg', **{})), map(None, 'abc', 'defg')) # empty keyword dict
265 self.assertRaises(TypeError, izip_longest, 3)
266 self.assertRaises(TypeError, izip_longest, range(3), 3)
267
268 for stmt in [
269 "izip_longest('abc', fv=1)",
Tim Petersea5962f2007-03-12 18:07:52 +0000270 "izip_longest('abc', fillvalue=1, bogus_keyword=None)",
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000271 ]:
272 try:
273 eval(stmt, globals(), locals())
274 except TypeError:
275 pass
276 else:
277 self.fail('Did not raise Type in: ' + stmt)
Tim Petersea5962f2007-03-12 18:07:52 +0000278
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000279 # Check tuple re-use (implementation detail)
280 self.assertEqual([tuple(list(pair)) for pair in izip_longest('abc', 'def')],
281 zip('abc', 'def'))
282 self.assertEqual([pair for pair in izip_longest('abc', 'def')],
283 zip('abc', 'def'))
284 ids = map(id, izip_longest('abc', 'def'))
285 self.assertEqual(min(ids), max(ids))
286 ids = map(id, list(izip_longest('abc', 'def')))
287 self.assertEqual(len(dict.fromkeys(ids)), len(ids))
288
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000289 def test_product(self):
290 for args, result in [
291 ([], []), # zero iterables ??? is this correct
292 (['ab'], [('a',), ('b',)]), # one iterable
293 ([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
294 ([range(0), range(2), range(3)], []), # first iterable with zero length
295 ([range(2), range(0), range(3)], []), # middle iterable with zero length
296 ([range(2), range(3), range(0)], []), # last iterable with zero length
297 ]:
298 self.assertEqual(list(product(*args)), result)
Raymond Hettinger08ff6822008-02-29 02:21:48 +0000299 for r in range(4):
300 self.assertEqual(list(product(*(args*r))),
301 list(product(*args, **dict(repeat=r))))
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000302 self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
303 self.assertRaises(TypeError, product, range(6), None)
304 argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
305 set('abcdefg'), range(11), tuple(range(13))]
306 for i in range(100):
307 args = [random.choice(argtypes) for j in range(random.randrange(5))]
308 n = reduce(operator.mul, map(len, args), 1) if args else 0
309 self.assertEqual(len(list(product(*args))), n)
310 args = map(iter, args)
311 self.assertEqual(len(list(product(*args))), n)
312
Raymond Hettinger73d79632008-02-23 02:20:41 +0000313 # Test implementation detail: tuple re-use
314 self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
315 self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000316
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000317 def test_repeat(self):
318 self.assertEqual(zip(xrange(3),repeat('a')),
319 [(0, 'a'), (1, 'a'), (2, 'a')])
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000320 self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
Raymond Hettinger02420702003-06-29 20:36:23 +0000321 self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000322 self.assertEqual(list(repeat('a', 0)), [])
323 self.assertEqual(list(repeat('a', -3)), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000324 self.assertRaises(TypeError, repeat)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000325 self.assertRaises(TypeError, repeat, None, 3, 4)
326 self.assertRaises(TypeError, repeat, None, 'a')
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000327 r = repeat(1+0j)
328 self.assertEqual(repr(r), 'repeat((1+0j))')
329 r = repeat(1+0j, 5)
330 self.assertEqual(repr(r), 'repeat((1+0j), 5)')
331 list(r)
332 self.assertEqual(repr(r), 'repeat((1+0j), 0)')
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000333
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000334 def test_imap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000335 self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
336 [0**1, 1**2, 2**3])
337 self.assertEqual(list(imap(None, 'abc', range(5))),
338 [('a',0),('b',1),('c',2)])
Raymond Hettinger02420702003-06-29 20:36:23 +0000339 self.assertEqual(list(imap(None, 'abc', count())),
340 [('a',0),('b',1),('c',2)])
341 self.assertEqual(take(2,imap(None, 'abc', count())),
342 [('a',0),('b',1)])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000343 self.assertEqual(list(imap(operator.pow, [])), [])
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000344 self.assertRaises(TypeError, imap)
345 self.assertRaises(TypeError, imap, operator.neg)
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000346 self.assertRaises(TypeError, imap(10, range(5)).next)
347 self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
348 self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000349
350 def test_starmap(self):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000351 self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
352 [0**1, 1**2, 2**3])
Raymond Hettinger02420702003-06-29 20:36:23 +0000353 self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
354 [0**1, 1**2, 2**3])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000355 self.assertEqual(list(starmap(operator.pow, [])), [])
Raymond Hettinger47317092008-01-17 03:02:14 +0000356 self.assertEqual(list(starmap(operator.pow, [iter([4,5])])), [4**5])
357 self.assertRaises(TypeError, list, starmap(operator.pow, [None]))
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000358 self.assertRaises(TypeError, starmap)
359 self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
360 self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
361 self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
362 self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000363
364 def test_islice(self):
365 for args in [ # islice(args) should agree with range(args)
366 (10, 20, 3),
367 (10, 3, 20),
368 (10, 20),
369 (10, 3),
370 (20,)
371 ]:
372 self.assertEqual(list(islice(xrange(100), *args)), range(*args))
373
374 for args, tgtargs in [ # Stop when seqn is exhausted
375 ((10, 110, 3), ((10, 100, 3))),
376 ((10, 110), ((10, 100))),
377 ((110,), (100,))
378 ]:
379 self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
380
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000381 # Test stop=None
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000382 self.assertEqual(list(islice(xrange(10), None)), range(10))
Raymond Hettingerb2594052004-12-05 09:25:51 +0000383 self.assertEqual(list(islice(xrange(10), None, None)), range(10))
384 self.assertEqual(list(islice(xrange(10), None, None, None)), range(10))
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000385 self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
386 self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
387
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000388 # Test number of items consumed SF #1171417
389 it = iter(range(10))
390 self.assertEqual(list(islice(it, 3)), range(3))
391 self.assertEqual(list(it), range(3, 10))
392
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000393 # Test invalid arguments
Raymond Hettinger341deb72003-05-02 19:44:20 +0000394 self.assertRaises(TypeError, islice, xrange(10))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000395 self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
396 self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
397 self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
398 self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
399 self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000400 self.assertRaises(ValueError, islice, xrange(10), 'a')
401 self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
402 self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
403 self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
404 self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
Kristján Valur Jónsson170eee92007-05-03 20:09:56 +0000405 self.assertEqual(len(list(islice(count(), 1, 10, maxsize))), 1)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000406
407 def test_takewhile(self):
408 data = [1, 3, 5, 20, 2, 4, 6, 8]
409 underten = lambda x: x<10
410 self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000411 self.assertEqual(list(takewhile(underten, [])), [])
412 self.assertRaises(TypeError, takewhile)
413 self.assertRaises(TypeError, takewhile, operator.pow)
414 self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
415 self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
416 self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000417 t = takewhile(bool, [1, 1, 1, 0, 0, 0])
418 self.assertEqual(list(t), [1, 1, 1])
419 self.assertRaises(StopIteration, t.next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000420
421 def test_dropwhile(self):
422 data = [1, 3, 5, 20, 2, 4, 6, 8]
423 underten = lambda x: x<10
424 self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +0000425 self.assertEqual(list(dropwhile(underten, [])), [])
426 self.assertRaises(TypeError, dropwhile)
427 self.assertRaises(TypeError, dropwhile, operator.pow)
428 self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
429 self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
430 self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000431
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000432 def test_tee(self):
Raymond Hettingerad983e72003-11-12 14:32:26 +0000433 n = 200
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000434 def irange(n):
435 for i in xrange(n):
436 yield i
437
438 a, b = tee([]) # test empty iterator
439 self.assertEqual(list(a), [])
440 self.assertEqual(list(b), [])
441
442 a, b = tee(irange(n)) # test 100% interleaved
443 self.assertEqual(zip(a,b), zip(range(n),range(n)))
444
445 a, b = tee(irange(n)) # test 0% interleaved
446 self.assertEqual(list(a), range(n))
447 self.assertEqual(list(b), range(n))
448
449 a, b = tee(irange(n)) # test dealloc of leading iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000450 for i in xrange(100):
451 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000452 del a
453 self.assertEqual(list(b), range(n))
454
455 a, b = tee(irange(n)) # test dealloc of trailing iterator
Raymond Hettingerad983e72003-11-12 14:32:26 +0000456 for i in xrange(100):
457 self.assertEqual(a.next(), i)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000458 del b
Raymond Hettingerad983e72003-11-12 14:32:26 +0000459 self.assertEqual(list(a), range(100, n))
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000460
461 for j in xrange(5): # test randomly interleaved
462 order = [0]*n + [1]*n
463 random.shuffle(order)
464 lists = ([], [])
465 its = tee(irange(n))
466 for i in order:
467 value = its[i].next()
468 lists[i].append(value)
469 self.assertEqual(lists[0], range(n))
470 self.assertEqual(lists[1], range(n))
471
Raymond Hettingerad983e72003-11-12 14:32:26 +0000472 # test argument format checking
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000473 self.assertRaises(TypeError, tee)
474 self.assertRaises(TypeError, tee, 3)
475 self.assertRaises(TypeError, tee, [1,2], 'x')
Raymond Hettingerad983e72003-11-12 14:32:26 +0000476 self.assertRaises(TypeError, tee, [1,2], 3, 'x')
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000477
Raymond Hettingerad983e72003-11-12 14:32:26 +0000478 # tee object should be instantiable
479 a, b = tee('abc')
480 c = type(a)('def')
481 self.assertEqual(list(c), list('def'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000482
Raymond Hettingerad983e72003-11-12 14:32:26 +0000483 # test long-lagged and multi-way split
484 a, b, c = tee(xrange(2000), 3)
485 for i in xrange(100):
486 self.assertEqual(a.next(), i)
487 self.assertEqual(list(b), range(2000))
488 self.assertEqual([c.next(), c.next()], range(2))
489 self.assertEqual(list(a), range(100,2000))
490 self.assertEqual(list(c), range(2,2000))
491
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000492 # test values of n
493 self.assertRaises(TypeError, tee, 'abc', 'invalid')
Neal Norwitz69e88972006-09-02 02:58:13 +0000494 self.assertRaises(ValueError, tee, [], -1)
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000495 for n in xrange(5):
496 result = tee('abc', n)
497 self.assertEqual(type(result), tuple)
498 self.assertEqual(len(result), n)
499 self.assertEqual(map(list, result), [list('abc')]*n)
500
Raymond Hettingerad983e72003-11-12 14:32:26 +0000501 # tee pass-through to copyable iterator
502 a, b = tee('abc')
503 c, d = tee(a)
504 self.assert_(a is c)
505
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000506 # test tee_new
507 t1, t2 = tee('abc')
508 tnew = type(t1)
509 self.assertRaises(TypeError, tnew)
510 self.assertRaises(TypeError, tnew, 10)
511 t3 = tnew(t1)
512 self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000513
Raymond Hettingera9f60922004-10-17 16:40:14 +0000514 # test that tee objects are weak referencable
515 a, b = tee(xrange(10))
516 p = proxy(a)
517 self.assertEqual(getattr(p, '__class__'), type(b))
518 del a
519 self.assertRaises(ReferenceError, getattr, p, '__class__')
520
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000521 def test_StopIteration(self):
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000522 self.assertRaises(StopIteration, izip().next)
523
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000524 for f in (chain, cycle, izip, groupby):
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000525 self.assertRaises(StopIteration, f([]).next)
526 self.assertRaises(StopIteration, f(StopNow()).next)
527
528 self.assertRaises(StopIteration, islice([], None).next)
529 self.assertRaises(StopIteration, islice(StopNow(), None).next)
530
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000531 p, q = tee([])
532 self.assertRaises(StopIteration, p.next)
533 self.assertRaises(StopIteration, q.next)
534 p, q = tee(StopNow())
535 self.assertRaises(StopIteration, p.next)
536 self.assertRaises(StopIteration, q.next)
537
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000538 self.assertRaises(StopIteration, repeat(None, 0).next)
539
540 for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
541 self.assertRaises(StopIteration, f(lambda x:x, []).next)
542 self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
543
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000544class TestGC(unittest.TestCase):
545
546 def makecycle(self, iterator, container):
547 container.append(iterator)
548 iterator.next()
549 del container, iterator
550
551 def test_chain(self):
552 a = []
553 self.makecycle(chain(a), a)
554
555 def test_cycle(self):
556 a = []
557 self.makecycle(cycle([a]*2), a)
558
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000559 def test_dropwhile(self):
560 a = []
561 self.makecycle(dropwhile(bool, [0, a, a]), a)
562
563 def test_groupby(self):
564 a = []
565 self.makecycle(groupby([a]*2, lambda x:x), a)
566
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000567 def test_ifilter(self):
568 a = []
569 self.makecycle(ifilter(lambda x:True, [a]*2), a)
570
571 def test_ifilterfalse(self):
572 a = []
573 self.makecycle(ifilterfalse(lambda x:False, a), a)
574
575 def test_izip(self):
576 a = []
577 self.makecycle(izip([a]*2, [a]*3), a)
578
579 def test_imap(self):
580 a = []
581 self.makecycle(imap(lambda x:x, [a]*2), a)
582
583 def test_islice(self):
584 a = []
585 self.makecycle(islice([a]*2, None), a)
586
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000587 def test_repeat(self):
588 a = []
589 self.makecycle(repeat(a), a)
590
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000591 def test_starmap(self):
592 a = []
593 self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
594
Raymond Hettingerff5dc0e2004-09-29 11:40:50 +0000595 def test_takewhile(self):
596 a = []
597 self.makecycle(takewhile(bool, [1, 0, a, a]), a)
598
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000599def R(seqn):
600 'Regular generator'
601 for i in seqn:
602 yield i
Raymond Hettinger8fd3f872003-05-02 22:38:07 +0000603
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000604class G:
605 'Sequence using __getitem__'
606 def __init__(self, seqn):
607 self.seqn = seqn
608 def __getitem__(self, i):
609 return self.seqn[i]
610
611class I:
612 'Sequence using iterator protocol'
613 def __init__(self, seqn):
614 self.seqn = seqn
615 self.i = 0
616 def __iter__(self):
617 return self
618 def next(self):
619 if self.i >= len(self.seqn): raise StopIteration
620 v = self.seqn[self.i]
621 self.i += 1
622 return v
623
624class Ig:
625 'Sequence using iterator protocol defined with a generator'
626 def __init__(self, seqn):
627 self.seqn = seqn
628 self.i = 0
629 def __iter__(self):
630 for val in self.seqn:
631 yield val
632
633class X:
634 'Missing __getitem__ and __iter__'
635 def __init__(self, seqn):
636 self.seqn = seqn
637 self.i = 0
638 def next(self):
639 if self.i >= len(self.seqn): raise StopIteration
640 v = self.seqn[self.i]
641 self.i += 1
642 return v
643
644class N:
645 'Iterator missing next()'
646 def __init__(self, seqn):
647 self.seqn = seqn
648 self.i = 0
649 def __iter__(self):
650 return self
651
652class E:
653 'Test propagation of exceptions'
654 def __init__(self, seqn):
655 self.seqn = seqn
656 self.i = 0
657 def __iter__(self):
658 return self
659 def next(self):
Raymond Hettingerffdb8bb2004-09-27 15:29:05 +0000660 3 // 0
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000661
662class S:
663 'Test immediate stop'
664 def __init__(self, seqn):
665 pass
666 def __iter__(self):
667 return self
668 def next(self):
669 raise StopIteration
670
671def L(seqn):
672 'Test multiple tiers of iterators'
673 return chain(imap(lambda x:x, R(Ig(G(seqn)))))
674
675
676class TestVariousIteratorArgs(unittest.TestCase):
677
678 def test_chain(self):
679 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
680 for g in (G, I, Ig, S, L, R):
681 self.assertEqual(list(chain(g(s))), list(g(s)))
682 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
Raymond Hettinger05bf6332008-02-28 22:30:42 +0000683 self.assertRaises(TypeError, list, chain(X(s)))
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000684 self.assertRaises(TypeError, list, chain(N(s)))
685 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
686
Raymond Hettinger50986cc2008-02-22 03:16:42 +0000687 def test_product(self):
688 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
689 self.assertRaises(TypeError, product, X(s))
690 self.assertRaises(TypeError, product, N(s))
691 self.assertRaises(ZeroDivisionError, product, E(s))
692
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000693 def test_cycle(self):
694 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
695 for g in (G, I, Ig, S, L, R):
696 tgtlen = len(s) * 3
697 expected = list(g(s))*3
698 actual = list(islice(cycle(g(s)), tgtlen))
699 self.assertEqual(actual, expected)
700 self.assertRaises(TypeError, cycle, X(s))
701 self.assertRaises(TypeError, list, cycle(N(s)))
702 self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
703
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000704 def test_groupby(self):
705 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
706 for g in (G, I, Ig, S, L, R):
707 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
708 self.assertRaises(TypeError, groupby, X(s))
709 self.assertRaises(TypeError, list, groupby(N(s)))
710 self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
711
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000712 def test_ifilter(self):
713 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
714 for g in (G, I, Ig, S, L, R):
715 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
716 self.assertRaises(TypeError, ifilter, isEven, X(s))
717 self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
718 self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
719
720 def test_ifilterfalse(self):
721 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
722 for g in (G, I, Ig, S, L, R):
723 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
724 self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
725 self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
726 self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
727
728 def test_izip(self):
729 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
730 for g in (G, I, Ig, S, L, R):
731 self.assertEqual(list(izip(g(s))), zip(g(s)))
732 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
733 self.assertRaises(TypeError, izip, X(s))
734 self.assertRaises(TypeError, list, izip(N(s)))
735 self.assertRaises(ZeroDivisionError, list, izip(E(s)))
736
Raymond Hettingerd36862c2007-02-21 05:20:38 +0000737 def test_iziplongest(self):
738 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
739 for g in (G, I, Ig, S, L, R):
740 self.assertEqual(list(izip_longest(g(s))), zip(g(s)))
741 self.assertEqual(list(izip_longest(g(s), g(s))), zip(g(s), g(s)))
742 self.assertRaises(TypeError, izip_longest, X(s))
743 self.assertRaises(TypeError, list, izip_longest(N(s)))
744 self.assertRaises(ZeroDivisionError, list, izip_longest(E(s)))
745
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000746 def test_imap(self):
747 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
748 for g in (G, I, Ig, S, L, R):
749 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
750 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
751 self.assertRaises(TypeError, imap, onearg, X(s))
752 self.assertRaises(TypeError, list, imap(onearg, N(s)))
753 self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
754
755 def test_islice(self):
756 for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
757 for g in (G, I, Ig, S, L, R):
758 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
759 self.assertRaises(TypeError, islice, X(s), 10)
760 self.assertRaises(TypeError, list, islice(N(s), 10))
761 self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
762
763 def test_starmap(self):
764 for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
765 for g in (G, I, Ig, S, L, R):
766 ss = zip(s, s)
767 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
768 self.assertRaises(TypeError, starmap, operator.pow, X(ss))
769 self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
770 self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
771
772 def test_takewhile(self):
773 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
774 for g in (G, I, Ig, S, L, R):
775 tgt = []
776 for elem in g(s):
777 if not isEven(elem): break
778 tgt.append(elem)
779 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
780 self.assertRaises(TypeError, takewhile, isEven, X(s))
781 self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
782 self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
783
784 def test_dropwhile(self):
785 for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
786 for g in (G, I, Ig, S, L, R):
787 tgt = []
788 for elem in g(s):
789 if not tgt and isOdd(elem): continue
790 tgt.append(elem)
791 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
792 self.assertRaises(TypeError, dropwhile, isOdd, X(s))
793 self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
794 self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
795
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000796 def test_tee(self):
797 for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
798 for g in (G, I, Ig, S, L, R):
799 it1, it2 = tee(g(s))
800 self.assertEqual(list(it1), list(g(s)))
801 self.assertEqual(list(it2), list(g(s)))
802 self.assertRaises(TypeError, tee, X(s))
803 self.assertRaises(TypeError, list, tee(N(s))[0])
804 self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
805
Raymond Hettinger5cab2e32004-02-10 09:25:40 +0000806class LengthTransparency(unittest.TestCase):
807
808 def test_repeat(self):
Raymond Hettinger6b27cda2005-09-24 21:23:05 +0000809 from test.test_iterlen import len
Raymond Hettinger5cab2e32004-02-10 09:25:40 +0000810 self.assertEqual(len(repeat(None, 50)), 50)
811 self.assertRaises(TypeError, len, repeat(None))
812
Raymond Hettingera56f6b62003-08-29 23:09:58 +0000813class RegressionTests(unittest.TestCase):
814
815 def test_sf_793826(self):
816 # Fix Armin Rigo's successful efforts to wreak havoc
817
818 def mutatingtuple(tuple1, f, tuple2):
819 # this builds a tuple t which is a copy of tuple1,
820 # then calls f(t), then mutates t to be equal to tuple2
821 # (needs len(tuple1) == len(tuple2)).
822 def g(value, first=[1]):
823 if first:
824 del first[:]
825 f(z.next())
826 return value
827 items = list(tuple2)
828 items[1:1] = list(tuple1)
829 gen = imap(g, items)
830 z = izip(*[gen]*len(tuple1))
831 z.next()
832
833 def f(t):
834 global T
835 T = t
836 first[:] = list(T)
837
838 first = []
839 mutatingtuple((1,2,3), f, (4,5,6))
840 second = list(T)
841 self.assertEqual(first, second)
842
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000843
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000844 def test_sf_950057(self):
845 # Make sure that chain() and cycle() catch exceptions immediately
846 # rather than when shifting between input sources
847
848 def gen1():
849 hist.append(0)
850 yield 1
851 hist.append(1)
Tim Petersbeb7c0c2004-07-18 17:34:03 +0000852 raise AssertionError
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000853 hist.append(2)
854
855 def gen2(x):
856 hist.append(3)
857 yield 2
858 hist.append(4)
859 if x:
860 raise StopIteration
861
862 hist = []
863 self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
864 self.assertEqual(hist, [0,1])
865
866 hist = []
867 self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
868 self.assertEqual(hist, [0,1])
869
870 hist = []
871 self.assertRaises(AssertionError, list, cycle(gen1()))
872 self.assertEqual(hist, [0,1])
873
Georg Brandlb84c1372007-01-21 10:28:43 +0000874class SubclassWithKwargsTest(unittest.TestCase):
875 def test_keywords_in_subclass(self):
876 # count is not subclassable...
877 for cls in (repeat, izip, ifilter, ifilterfalse, chain, imap,
878 starmap, islice, takewhile, dropwhile, cycle):
879 class Subclass(cls):
880 def __init__(self, newarg=None, *args):
881 cls.__init__(self, *args)
882 try:
883 Subclass(newarg=1)
884 except TypeError, err:
885 # we expect type errors because of wrong argument count
886 self.failIf("does not take keyword arguments" in err.args[0])
887
888
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +0000889libreftest = """ Doctest for examples in the library reference: libitertools.tex
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000890
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000891
892>>> amounts = [120.15, 764.05, 823.14]
893>>> for checknum, amount in izip(count(1200), amounts):
894... print 'Check %d is for $%.2f' % (checknum, amount)
895...
896Check 1200 is for $120.15
897Check 1201 is for $764.05
898Check 1202 is for $823.14
899
900>>> import operator
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000901>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
902... print cube
903...
9041
9058
90627
907
908>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +0000909>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000910... print name.title()
911...
912Alex
913Laura
914Martin
915Walter
916Samuele
917
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000918>>> from operator import itemgetter
919>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Armin Rigoa3f09272006-05-28 19:13:17 +0000920>>> di = sorted(sorted(d.iteritems()), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000921>>> for k, g in groupby(di, itemgetter(1)):
922... print k, map(itemgetter(0), g)
923...
9241 ['a', 'c', 'e']
9252 ['b', 'd', 'f']
9263 ['g']
927
Raymond Hettinger734fb572004-01-20 20:04:40 +0000928# Find runs of consecutive numbers using groupby. The key to the solution
929# is differencing with a range so that consecutive numbers all appear in
930# same group.
931>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
932>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
933... print map(operator.itemgetter(1), g)
934...
935[1]
936[4, 5, 6]
937[10]
938[15, 16, 17, 18]
939[22]
940[25, 26, 27, 28]
941
Raymond Hettingera098b332003-09-08 23:58:40 +0000942>>> def take(n, seq):
943... return list(islice(seq, n))
944
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000945>>> def enumerate(iterable):
946... return izip(count(), iterable)
947
948>>> def tabulate(function):
949... "Return function(0), function(1), ..."
950... return imap(function, count())
951
952>>> def iteritems(mapping):
953... return izip(mapping.iterkeys(), mapping.itervalues())
954
955>>> def nth(iterable, n):
956... "Returns the nth item"
Raymond Hettinger60eca932003-02-09 06:40:58 +0000957... return list(islice(iterable, n, n+1))
958
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000959>>> def all(seq, pred=None):
960... "Returns True if pred(x) is true for every element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000961... for elem in ifilterfalse(pred, seq):
962... return False
963... return True
Raymond Hettinger60eca932003-02-09 06:40:58 +0000964
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000965>>> def any(seq, pred=None):
966... "Returns True if pred(x) is true for at least one element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000967... for elem in ifilter(pred, seq):
968... return True
969... return False
Raymond Hettinger60eca932003-02-09 06:40:58 +0000970
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000971>>> def no(seq, pred=None):
972... "Returns True if pred(x) is false for every element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000973... for elem in ifilter(pred, seq):
974... return False
975... return True
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000976
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000977>>> def quantify(seq, pred=None):
978... "Count how many times the predicate is true in the sequence"
Raymond Hettinger77fe69b2003-08-08 04:33:19 +0000979... return sum(imap(pred, seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000980
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000981>>> def padnone(seq):
982... "Returns the sequence elements and then returns None indefinitely"
983... return chain(seq, repeat(None))
984
985>>> def ncycles(seq, n):
986... "Returns the sequence elements n times"
987... return chain(*repeat(seq, n))
988
989>>> def dotproduct(vec1, vec2):
990... return sum(imap(operator.mul, vec1, vec2))
991
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000992>>> def flatten(listOfLists):
993... return list(chain(*listOfLists))
994
995>>> def repeatfunc(func, times=None, *args):
996... "Repeat calls to func with specified arguments."
997... " Example: repeatfunc(random.random)"
998... if times is None:
999... return starmap(func, repeat(args))
1000... else:
1001... return starmap(func, repeat(args, times))
1002
Raymond Hettingerd591f662003-10-26 15:34:50 +00001003>>> def pairwise(iterable):
1004... "s -> (s0,s1), (s1,s2), (s2, s3), ..."
1005... a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +00001006... try:
1007... b.next()
1008... except StopIteration:
1009... pass
1010... return izip(a, b)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001011
1012This is not part of the examples but it tests to make sure the definitions
1013perform as purported.
1014
Raymond Hettingera098b332003-09-08 23:58:40 +00001015>>> take(10, count())
1016[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1017
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001018>>> list(enumerate('abc'))
1019[(0, 'a'), (1, 'b'), (2, 'c')]
1020
1021>>> list(islice(tabulate(lambda x: 2*x), 4))
1022[0, 2, 4, 6]
1023
1024>>> nth('abcde', 3)
1025['d']
1026
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001027>>> all([2, 4, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001028True
1029
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001030>>> all([2, 3, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001031False
1032
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001033>>> any([2, 4, 6, 8], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001034True
1035
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001036>>> any([1, 3, 5, 9], lambda x: x%2==0,)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001037False
1038
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001039>>> no([1, 3, 5, 9], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001040True
1041
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001042>>> no([1, 2, 5, 9], lambda x: x%2==0)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001043False
1044
Raymond Hettingerdbe3d282003-10-05 16:47:36 +00001045>>> quantify(xrange(99), lambda x: x%2==0)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000104650
1047
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001048>>> a = [[1, 2, 3], [4, 5, 6]]
1049>>> flatten(a)
1050[1, 2, 3, 4, 5, 6]
1051
1052>>> list(repeatfunc(pow, 5, 2, 3))
1053[8, 8, 8, 8, 8]
1054
1055>>> import random
1056>>> take(5, imap(int, repeatfunc(random.random)))
1057[0, 0, 0, 0, 0]
1058
Raymond Hettingerd591f662003-10-26 15:34:50 +00001059>>> list(pairwise('abcd'))
1060[('a', 'b'), ('b', 'c'), ('c', 'd')]
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001061
Raymond Hettingerd591f662003-10-26 15:34:50 +00001062>>> list(pairwise([]))
1063[]
1064
1065>>> list(pairwise('a'))
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001066[]
1067
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001068>>> list(islice(padnone('abc'), 0, 6))
1069['a', 'b', 'c', None, None, None]
1070
1071>>> list(ncycles('abc', 3))
1072['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
1073
1074>>> dotproduct([1,2,3], [4,5,6])
107532
1076
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001077"""
1078
1079__test__ = {'libreftest' : libreftest}
1080
1081def test_main(verbose=None):
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001082 test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
Georg Brandlb84c1372007-01-21 10:28:43 +00001083 RegressionTests, LengthTransparency,
1084 SubclassWithKwargsTest)
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001085 test_support.run_unittest(*test_classes)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001086
1087 # verify reference counting
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001088 if verbose and hasattr(sys, "gettotalrefcount"):
Raymond Hettinger02420702003-06-29 20:36:23 +00001089 import gc
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001090 counts = [None] * 5
1091 for i in xrange(len(counts)):
Raymond Hettingerf0fa1c02003-05-29 07:18:57 +00001092 test_support.run_unittest(*test_classes)
Raymond Hettinger02420702003-06-29 20:36:23 +00001093 gc.collect()
Raymond Hettinger8fd3f872003-05-02 22:38:07 +00001094 counts[i] = sys.gettotalrefcount()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001095 print counts
1096
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001097 # doctest the examples in the library reference
Raymond Hettinger929f06c2003-05-16 23:16:36 +00001098 test_support.run_doctest(sys.modules[__name__], verbose)
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001099
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001100if __name__ == "__main__":
1101 test_main(verbose=True)