247 lines
7.2 KiB
Python
247 lines
7.2 KiB
Python
|
# Test iterators.
|
||
|
|
|
||
|
|
import unittest
|
||
|
|
from test_support import run_unittest, TESTFN, unlink
|
||
|
|
|
||
|
|
# Test result of triple loop (too big to inline)
|
||
|
|
TRIPLETS = [(0, 0, 0), (0, 0, 1), (0, 0, 2),
|
||
|
|
(0, 1, 0), (0, 1, 1), (0, 1, 2),
|
||
|
|
(0, 2, 0), (0, 2, 1), (0, 2, 2),
|
||
|
|
|
||
|
|
(1, 0, 0), (1, 0, 1), (1, 0, 2),
|
||
|
|
(1, 1, 0), (1, 1, 1), (1, 1, 2),
|
||
|
|
(1, 2, 0), (1, 2, 1), (1, 2, 2),
|
||
|
|
|
||
|
|
(2, 0, 0), (2, 0, 1), (2, 0, 2),
|
||
|
|
(2, 1, 0), (2, 1, 1), (2, 1, 2),
|
||
|
|
(2, 2, 0), (2, 2, 1), (2, 2, 2)]
|
||
|
|
|
||
|
|
# Helper classes
|
||
|
|
|
||
|
|
class BasicIterClass:
|
||
|
|
def __init__(self, n):
|
||
|
|
self.n = n
|
||
|
|
self.i = 0
|
||
|
|
def next(self):
|
||
|
|
res = self.i
|
||
|
|
if res >= self.n:
|
||
|
|
raise StopIteration
|
||
|
|
self.i = res + 1
|
||
|
|
return res
|
||
|
|
|
||
|
|
class IteratingSequenceClass:
|
||
|
|
def __init__(self, n):
|
||
|
|
self.n = n
|
||
|
|
def __iter__(self):
|
||
|
|
return BasicIterClass(self.n)
|
||
|
|
|
||
|
|
class SequenceClass:
|
||
|
|
def __init__(self, n):
|
||
|
|
self.n = n
|
||
|
|
def __getitem__(self, i):
|
||
|
|
if 0 <= i < self.n:
|
||
|
|
return i
|
||
|
|
else:
|
||
|
|
raise IndexError
|
||
|
|
|
||
|
|
# Main test suite
|
||
|
|
|
||
|
|
class TestCase(unittest.TestCase):
|
||
|
|
|
||
|
|
# Helper to check that an iterator returns a given sequence
|
||
|
|
def check_iterator(self, it, seq):
|
||
|
|
res = []
|
||
|
|
while 1:
|
||
|
|
try:
|
||
|
|
val = it.next()
|
||
|
|
except StopIteration:
|
||
|
|
break
|
||
|
|
res.append(val)
|
||
|
|
self.assertEqual(res, seq)
|
||
|
|
|
||
|
|
# Helper to check that a for loop generates a given sequence
|
||
|
|
def check_for_loop(self, expr, seq):
|
||
|
|
res = []
|
||
|
|
for val in expr:
|
||
|
|
res.append(val)
|
||
|
|
self.assertEqual(res, seq)
|
||
|
|
|
||
|
|
# Test basic use of iter() function
|
||
|
|
def test_iter_basic(self):
|
||
|
|
self.check_iterator(iter(range(10)), range(10))
|
||
|
|
|
||
|
|
# Test that iter(iter(x)) is the same as iter(x)
|
||
|
|
def test_iter_idempotency(self):
|
||
|
|
seq = range(10)
|
||
|
|
it = iter(seq)
|
||
|
|
it2 = iter(it)
|
||
|
|
self.assert_(it is it2)
|
||
|
|
|
||
|
|
# Test that for loops over iterators work
|
||
|
|
def test_iter_for_loop(self):
|
||
|
|
self.check_for_loop(iter(range(10)), range(10))
|
||
|
|
|
||
|
|
# Test several independent iterators over the same list
|
||
|
|
def test_iter_independence(self):
|
||
|
|
seq = range(3)
|
||
|
|
res = []
|
||
|
|
for i in iter(seq):
|
||
|
|
for j in iter(seq):
|
||
|
|
for k in iter(seq):
|
||
|
|
res.append((i, j, k))
|
||
|
|
self.assertEqual(res, TRIPLETS)
|
||
|
|
|
||
|
|
# Test triple list comprehension using iterators
|
||
|
|
def test_nested_comprehensions_iter(self):
|
||
|
|
seq = range(3)
|
||
|
|
res = [(i, j, k)
|
||
|
|
for i in iter(seq) for j in iter(seq) for k in iter(seq)]
|
||
|
|
self.assertEqual(res, TRIPLETS)
|
||
|
|
|
||
|
|
# Test triple list comprehension without iterators
|
||
|
|
def test_nested_comprehensions_for(self):
|
||
|
|
seq = range(3)
|
||
|
|
res = [(i, j, k) for i in seq for j in seq for k in seq]
|
||
|
|
self.assertEqual(res, TRIPLETS)
|
||
|
|
|
||
|
|
# Test a class with __iter__ in a for loop
|
||
|
|
def test_iter_class_for(self):
|
||
|
|
self.check_for_loop(IteratingSequenceClass(10), range(10))
|
||
|
|
|
||
|
|
# Test a class with __iter__ with explicit iter()
|
||
|
|
def test_iter_class_iter(self):
|
||
|
|
self.check_iterator(iter(IteratingSequenceClass(10)), range(10))
|
||
|
|
|
||
|
|
# Test for loop on a sequence class without __iter__
|
||
|
|
def test_seq_class_for(self):
|
||
|
|
self.check_for_loop(SequenceClass(10), range(10))
|
||
|
|
|
||
|
|
# Test iter() on a sequence class without __iter__
|
||
|
|
def test_seq_class_iter(self):
|
||
|
|
self.check_iterator(iter(SequenceClass(10)), range(10))
|
||
|
|
|
||
|
|
# Test two-argument iter() with callable instance
|
||
|
|
def test_iter_callable(self):
|
||
|
|
class C:
|
||
|
|
def __init__(self):
|
||
|
|
self.i = 0
|
||
|
|
def __call__(self):
|
||
|
|
i = self.i
|
||
|
|
self.i = i + 1
|
||
|
|
if i > 100:
|
||
|
|
raise IndexError # Emergency stop
|
||
|
|
return i
|
||
|
|
self.check_iterator(iter(C(), 10), range(10))
|
||
|
|
|
||
|
|
# Test two-argument iter() with function
|
||
|
|
def test_iter_function(self):
|
||
|
|
def spam(state=[0]):
|
||
|
|
i = state[0]
|
||
|
|
state[0] = i+1
|
||
|
|
return i
|
||
|
|
self.check_iterator(iter(spam, 10), range(10))
|
||
|
|
|
||
|
|
# Test two-argument iter() with function that raises StopIteration
|
||
|
|
def test_iter_function_stop(self):
|
||
|
|
def spam(state=[0]):
|
||
|
|
i = state[0]
|
||
|
|
if i == 10:
|
||
|
|
raise StopIteration
|
||
|
|
state[0] = i+1
|
||
|
|
return i
|
||
|
|
self.check_iterator(iter(spam, 20), range(10))
|
||
|
|
|
||
|
|
# Test exception propagation through function iterator
|
||
|
|
def test_exception_function(self):
|
||
|
|
def spam(state=[0]):
|
||
|
|
i = state[0]
|
||
|
|
state[0] = i+1
|
||
|
|
if i == 10:
|
||
|
|
raise RuntimeError
|
||
|
|
return i
|
||
|
|
res = []
|
||
|
|
try:
|
||
|
|
for x in iter(spam, 20):
|
||
|
|
res.append(x)
|
||
|
|
except RuntimeError:
|
||
|
|
self.assertEqual(res, range(10))
|
||
|
|
else:
|
||
|
|
self.fail("should have raised RuntimeError")
|
||
|
|
|
||
|
|
# Test exception propagation through sequence iterator
|
||
|
|
def test_exception_sequence(self):
|
||
|
|
class MySequenceClass(SequenceClass):
|
||
|
|
def __getitem__(self, i):
|
||
|
|
if i == 10:
|
||
|
|
raise RuntimeError
|
||
|
|
return SequenceClass.__getitem__(self, i)
|
||
|
|
res = []
|
||
|
|
try:
|
||
|
|
for x in MySequenceClass(20):
|
||
|
|
res.append(x)
|
||
|
|
except RuntimeError:
|
||
|
|
self.assertEqual(res, range(10))
|
||
|
|
else:
|
||
|
|
self.fail("should have raised RuntimeError")
|
||
|
|
|
||
|
|
# Test for StopIteration from __getitem__
|
||
|
|
def test_stop_sequence(self):
|
||
|
|
class MySequenceClass(SequenceClass):
|
||
|
|
def __getitem__(self, i):
|
||
|
|
if i == 10:
|
||
|
|
raise StopIteration
|
||
|
|
return SequenceClass.__getitem__(self, i)
|
||
|
|
self.check_for_loop(MySequenceClass(20), range(10))
|
||
|
|
|
||
|
|
# Test a big range
|
||
|
|
def test_iter_big_range(self):
|
||
|
|
self.check_for_loop(iter(range(10000)), range(10000))
|
||
|
|
|
||
|
|
# Test an empty list
|
||
|
|
def test_iter_empty(self):
|
||
|
|
self.check_for_loop(iter([]), [])
|
||
|
|
|
||
|
|
# Test a tuple
|
||
|
|
def test_iter_tuple(self):
|
||
|
|
self.check_for_loop(iter((0,1,2,3,4,5,6,7,8,9)), range(10))
|
||
|
|
|
||
|
|
# Test an xrange
|
||
|
|
def test_iter_xrange(self):
|
||
|
|
self.check_for_loop(iter(xrange(10)), range(10))
|
||
|
|
|
||
|
|
# Test a string
|
||
|
|
def test_iter_string(self):
|
||
|
|
self.check_for_loop(iter("abcde"), ["a", "b", "c", "d", "e"])
|
||
|
|
|
||
|
|
# Test a Unicode string
|
||
|
|
def test_iter_unicode(self):
|
||
|
|
self.check_for_loop(iter(u"abcde"), [u"a", u"b", u"c", u"d", u"e"])
|
||
|
|
|
||
|
|
# Test a directory
|
||
|
|
def test_iter_dict(self):
|
||
|
|
dict = {}
|
||
|
|
for i in range(10):
|
||
|
|
dict[i] = None
|
||
|
|
self.check_for_loop(dict, dict.keys())
|
||
|
|
|
||
|
|
# Test a file
|
||
|
|
def test_iter_file(self):
|
||
|
|
f = open(TESTFN, "w")
|
||
|
|
try:
|
||
|
|
for i in range(5):
|
||
|
|
f.write("%d\n" % i)
|
||
|
|
finally:
|
||
|
|
f.close()
|
||
|
|
f = open(TESTFN, "r")
|
||
|
|
try:
|
||
|
|
self.check_for_loop(f, ["0\n", "1\n", "2\n", "3\n", "4\n"])
|
||
|
|
self.check_for_loop(f, [])
|
||
|
|
finally:
|
||
|
|
f.close()
|
||
|
|
try:
|
||
|
|
unlink(TESTFN)
|
||
|
|
except OSError:
|
||
|
|
pass
|
||
|
|
|
||
|
|
run_unittest(TestCase)
|