Add test suite for iterators.

1 files changed, 246 insertions, 0 deletions

diff --git a/Lib/test/test_iter.py b/Lib/test/test_iter.py
new file mode 100644
index 0000000..d4b8736
--- /dev/null
+++ b/Lib/test/test_iter.py
@@ -0,0 +1,246 @@
+# 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)