allow yield statements (#136)

* allow yield and yield from statements

docs/CHANGES.rst

@@ -4,7 +4,8 @@ Changes
 4.0b6 (unreleased)
 ------------------
 
-- Nothing changed yet.
+- Allow `yield` and `yield from` statements.
+  Generator functions would now work in RestrictedPython.
 
 
 4.0b5 (2018-09-05)

src/RestrictedPython/transformer.py

@@ -1365,12 +1365,12 @@ def visit_Return(self, node):
         return self.node_contents_visit(node)
 
     def visit_Yield(self, node):
-        """Deny `yield` unconditionally."""
-        self.not_allowed(node)
+        """Allow `yield`statements without restrictions."""
+        return self.node_contents_visit(node)
 
     def visit_YieldFrom(self, node):
-        """Deny `yield from` unconditionally."""
-        self.not_allowed(node)
+        """Allow `yield`statements without restrictions."""
+        return self.node_contents_visit(node)
 
     def visit_Global(self, node):
         """Allow `global` statements without restrictions."""

tests/transformer/test_yield.py

@@ -1,25 +1,59 @@
 from RestrictedPython._compat import IS_PY3
+from RestrictedPython._compat import IS_PY35_OR_GREATER
 from tests import c_exec
 
 import pytest
 
 
 YIELD_EXAMPLE = """\
-def no_yield():
+def test_generator():
     yield 42
 """
 
 
 @pytest.mark.parametrize(*c_exec)
 def test_yield(c_exec):
-    """It prevents using the `yield` statement."""
+    """`yield` statement should be allowed."""
     result = c_exec(YIELD_EXAMPLE)
-    assert result.errors == ("Line 2: Yield statements are not allowed.",)
-    assert result.code is None
+    assert result.errors == ()
+    assert result.code is not None
+    local = {}
+    exec(result.code, {}, local)
+    test_generator = local['test_generator']
+    exec_result = list(test_generator())
+    assert exec_result == [42]
 
 
-# Modified Example from http://stackabuse.com/python-async-await-tutorial/
 YIELD_FORM_EXAMPLE = """
+def reader_wapper(input):
+    yield from input
+"""
+
+
+@pytest.mark.skipif(
+    not IS_PY3,
+    reason="`yield from` statement was first introduced in Python 3.3")
+@pytest.mark.parametrize(*c_exec)
+def test_yield_from(c_exec):
+    """`yield from` statement should be allowed."""
+    result = c_exec(YIELD_FORM_EXAMPLE)
+    assert result.errors == ()
+    assert result.code is not None
+
+    def my_external_generator():
+        my_list = [1, 2, 3, 4, 5]
+        for elem in my_list:
+            yield(elem)
+
+    local = {}
+    exec(result.code, {}, local)
+    reader_wapper = local['reader_wapper']
+    exec_result = list(reader_wapper(my_external_generator()))
+    assert exec_result == [1, 2, 3, 4, 5]
+
+
+# Modified Example from http://stackabuse.com/python-async-await-tutorial/
+ASYNCIO_YIELD_FORM_EXAMPLE = """
 import asyncio
 
 @asyncio.coroutine
@@ -32,7 +66,29 @@ def get_json(client, url):
     not IS_PY3,
     reason="`yield from` statement was first introduced in Python 3.3")
 @pytest.mark.parametrize(*c_exec)
-def test_yield_from(c_exec):
-    result = c_exec(YIELD_FORM_EXAMPLE)
-    assert result.errors == ('Line 6: YieldFrom statements are not allowed.',)
+def test_asyncio_yield_from(c_exec):
+    """`yield from` statement should be allowed."""
+    result = c_exec(ASYNCIO_YIELD_FORM_EXAMPLE)
+    assert result.errors == ()
+    assert result.code is not None
+
+
+ASYNC_YIELD_FORM_EXAMPLE = """
+import asyncio
+
+async def get_json(client, url):
+    file_content = yield from load_file('data.ini')
+"""
+
+
+@pytest.mark.skipif(
+    not IS_PY35_OR_GREATER,
+    reason="`yield from` statement was first introduced in Python 3.3")
+@pytest.mark.parametrize(*c_exec)
+def test_async_yield_from(c_exec):
+    """`yield from` statement should be allowed."""
+    result = c_exec(ASYNC_YIELD_FORM_EXAMPLE)
+    assert result.errors == (
+        'Line 4: AsyncFunctionDef statements are not allowed.',
+    )
     assert result.code is None