[pypy-commit] pypy stacklet: Tests for generators.

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: stacklet
Changeset: r46590:fcadccd37211
Date: 2011-08-18 09:03 +0200
http://bitbucket.org/pypy/pypy/changeset/fcadccd37211/

Log:	Tests for generators.

diff --git a/pypy/module/_continuation/app_continuation.py b/pypy/module/_continuation/app_continuation.py
--- a/pypy/module/_continuation/app_continuation.py
+++ b/pypy/module/_continuation/app_continuation.py
@@ -25,9 +25,11 @@
 
     def next(self, value=None):
         res = self.switch(value)
-        if self.is_alive():
+        if self.is_pending():
             return res
         else:
+            if res is not None:
+                raise TypeError("_continuation.generator must return None")
             raise StopIteration
 
     send = next
diff --git a/pypy/module/_continuation/test/test_generator.py b/pypy/module/_continuation/test/test_generator.py
new file mode 100644
--- /dev/null
+++ b/pypy/module/_continuation/test/test_generator.py
@@ -0,0 +1,72 @@
+from pypy.conftest import gettestobjspace
+
+
+class AppTestGenerator:
+    def setup_class(cls):
+        cls.space = gettestobjspace(usemodules=['_continuation'])
+
+    def test_simple(self):
+        from _continuation import generator
+        #
+        @generator
+        def f(gen, n):
+            gen.switch(n+1)
+            f2(gen, n+2)
+            gen.switch(n+3)
+        #
+        def f2(gen, m):
+            gen.switch(m*2)
+        #
+        g = f(10)
+        res = g.next()
+        assert res == 11
+        res = g.next()
+        assert res == 24
+        res = g.next()
+        assert res == 13
+        raises(StopIteration, g.next)
+
+    def test_iterator(self):
+        from _continuation import generator
+        #
+        @generator
+        def f(gen, n):
+            gen.switch(n+1)
+            f2(gen, n+2)
+            gen.switch(n+3)
+        #
+        def f2(gen, m):
+            gen.switch(m*2)
+        #
+        res = list(f(10))
+        assert res == [11, 24, 13]
+        g = f(20)
+        assert iter(g) is g
+
+    def test_bound_method(self):
+        from _continuation import generator
+        #
+        class A(object):
+            def __init__(self, m):
+                self.m = m
+            #
+            @generator
+            def f(self, gen, n):
+                gen.switch(n - self.m)
+        #
+        a = A(10)
+        res = list(a.f(25))
+        assert res == [15]
+
+    def test_must_return_None(self):
+        from _continuation import generator
+        #
+        @generator
+        def f(gen, n):
+            gen.switch(n+1)
+            return "foo"
+        #
+        g = f(10)
+        res = g.next()
+        assert res == 11
+        raises(TypeError, g.next)