Support yield statement for creating coroutines (#652)

* Support `yield` statement for creating coroutines

* Here you have it

* Forgot this lad

Author: chrisrink10

CHANGELOG.md

@@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
  * Added support for hierarchies (#633)
  * Added support for several more utility Namespace and Var utility functions (#636)
  * Added `basilisp.io` namespace with polymorphic reader and writer functions (#645)
+ * Added support for coroutines and generators using `yield` syntax (#652)
 
 ### Changed
  * PyTest is now an optional extra dependency, rather than a required dependency (#622)

src/basilisp/io.lpy

@@ -15,7 +15,7 @@
 
 (extend-protocol Coercions
   nil
-  (as-path [f] nil)
+  (as-path [_] nil)
 
   python/str
   (as-path [f] (pathlib/Path f))

src/basilisp/lang/compiler/analyzer.py

@@ -121,7 +121,11 @@
 from basilisp.lang.compiler.nodes import Set as SetNode
 from basilisp.lang.compiler.nodes import SetBang, SpecialFormNode, Throw, Try, VarRef
 from basilisp.lang.compiler.nodes import Vector as VectorNode
-from basilisp.lang.compiler.nodes import WithMeta, deftype_or_reify_python_member_names
+from basilisp.lang.compiler.nodes import (
+    WithMeta,
+    Yield,
+    deftype_or_reify_python_member_names,
+)
 from basilisp.lang.interfaces import IMeta, IRecord, ISeq, IType, IWithMeta
 from basilisp.lang.runtime import Var
 from basilisp.lang.typing import CompilerOpts, LispForm, ReaderForm
@@ -2903,6 +2907,32 @@ def _var_ast(form: ISeq, ctx: AnalyzerContext) -> VarRef:
     )
 
 
+def _yield_ast(form: ISeq, ctx: AnalyzerContext) -> Yield:
+    assert form.first == SpecialForm.YIELD
+
+    if ctx.func_ctx is None:
+        raise AnalyzerException(
+            "yield forms may not appear in function context", form=form
+        )
+
+    nelems = count(form)
+    if nelems not in {1, 2}:
+        raise AnalyzerException(
+            "yield forms must contain 1 or 2 elements, as in: (yield [expr])", form=form
+        )
+
+    if nelems == 2:
+        with ctx.expr_pos():
+            expr = _analyze_form(runtime.nth(form, 1), ctx)
+        return Yield(
+            form=form,
+            expr=expr,
+            env=ctx.get_node_env(pos=ctx.syntax_position),
+        )
+    else:
+        return Yield.expressionless(form, ctx.get_node_env(pos=ctx.syntax_position))
+
+
 SpecialFormHandler = Callable[[ISeq, AnalyzerContext], SpecialFormNode]
 _SPECIAL_FORM_HANDLERS: Mapping[sym.Symbol, SpecialFormHandler] = {
     SpecialForm.AWAIT: _await_ast,
@@ -2924,6 +2954,7 @@ def _var_ast(form: ISeq, ctx: AnalyzerContext) -> VarRef:
     SpecialForm.THROW: _throw_ast,
     SpecialForm.TRY: _try_ast,
     SpecialForm.VAR: _var_ast,
+    SpecialForm.YIELD: _yield_ast,
 }
 
 

src/basilisp/lang/compiler/constants.py

@@ -25,6 +25,7 @@ class SpecialForm:
     THROW = sym.symbol("throw")
     TRY = sym.symbol("try")
     VAR = sym.symbol("var")
+    YIELD = sym.symbol("yield")
 
 
 AMPERSAND = sym.symbol("&")

src/basilisp/lang/compiler/generator.py

@@ -93,7 +93,7 @@
 from basilisp.lang.compiler.nodes import Set as SetNode
 from basilisp.lang.compiler.nodes import SetBang, Throw, Try, VarRef
 from basilisp.lang.compiler.nodes import Vector as VectorNode
-from basilisp.lang.compiler.nodes import WithMeta
+from basilisp.lang.compiler.nodes import WithMeta, Yield
 from basilisp.lang.interfaces import IMeta, IRecord, ISeq, ISeqable, IType
 from basilisp.lang.runtime import CORE_NS
 from basilisp.lang.runtime import NS_VAR_NAME as LISP_NS_VAR
@@ -710,8 +710,6 @@ def statementize(e: ast.AST) -> ast.AST:
             ast.With,
             ast.FunctionDef,
             ast.Return,
-            ast.Yield,
-            ast.YieldFrom,
             ast.Global,
             ast.ClassDef,
             ast.AsyncFunctionDef,
@@ -2815,6 +2813,17 @@ def _try_to_py_ast(ctx: GeneratorContext, node: Try) -> GeneratedPyAST:
     )
 
 
+@_with_ast_loc_deps
+def _yield_to_py_ast(ctx: GeneratorContext, node: Yield) -> GeneratedPyAST:
+    assert node.op == NodeOp.YIELD
+    if node.expr is None:
+        return GeneratedPyAST(node=ast.Yield(value=None))
+    expr_ast = gen_py_ast(ctx, node.expr)
+    return GeneratedPyAST(
+        node=ast.Yield(value=expr_ast.node), dependencies=expr_ast.dependencies
+    )
+
+
 ##########
 # Symbols
 ##########
@@ -3607,6 +3616,7 @@ def _const_node_to_py_ast(ctx: GeneratorContext, lisp_ast: Const) -> GeneratedPy
     NodeOp.SET_BANG: _set_bang_to_py_ast,
     NodeOp.THROW: _throw_to_py_ast,
     NodeOp.TRY: _try_to_py_ast,
+    NodeOp.YIELD: _yield_to_py_ast,
     NodeOp.VAR: _var_sym_to_py_ast,
     NodeOp.VECTOR: _vec_to_py_ast,
     NodeOp.WITH_META: _with_meta_to_py_ast,  # type: ignore

src/basilisp/lang/compiler/nodes.py

@@ -104,6 +104,7 @@ class NodeOp(Enum):
     VAR = kw.keyword("var")
     VECTOR = kw.keyword("vector")
     WITH_META = kw.keyword("with-meta")
+    YIELD = kw.keyword("yield")
 
 
 T = TypeVar("T")
@@ -937,6 +938,21 @@ class WithMeta(Node[LispForm]):
     raw_forms: IPersistentVector[LispForm] = vec.PersistentVector.empty()
 
 
+@attr.s(auto_attribs=True, frozen=True, slots=True)
+class Yield(Node[SpecialForm]):
+    form: SpecialForm
+    expr: Optional[Node]
+    env: NodeEnv
+    children: Sequence[kw.Keyword] = vec.v(EXPR)
+    op: NodeOp = NodeOp.YIELD
+    top_level: bool = False
+    raw_forms: IPersistentVector[LispForm] = vec.PersistentVector.empty()
+
+    @classmethod
+    def expressionless(cls, form: SpecialForm, env: NodeEnv):
+        return cls(form=form, expr=None, env=env, children=vec.PersistentVector.empty())
+
+
 SpecialFormNode = Union[
     Await,
     Def,
@@ -959,4 +975,5 @@ class WithMeta(Node[LispForm]):
     Throw,
     Try,
     VarRef,
+    Yield,
 ]

src/basilisp/lang/runtime.py

@@ -113,6 +113,7 @@
 _THROW = sym.symbol("throw")
 _TRY = sym.symbol("try")
 _VAR = sym.symbol("var")
+_YIELD = sym.symbol("yield")
 _SPECIAL_FORMS = lset.s(
     _AWAIT,
     _CATCH,
@@ -136,6 +137,7 @@
     _THROW,
     _TRY,
     _VAR,
+    _YIELD,
 )
 
 

tests/basilisp/compiler_test.py

@@ -5676,3 +5676,107 @@ def test_var_reader_literal(self, lcompile: CompileFn, ns: runtime.Namespace):
         v = lcompile(code)
         assert v == Var.find_in_ns(sym.symbol(ns_name), sym.symbol("some-var"))
         assert v.value == "a value"
+
+
+class TestYield:
+    @pytest.mark.parametrize(
+        "code",
+        [
+            "(yield)",
+            "(yield :a)",
+            "(let* [v :a] (yield v))",
+        ],
+    )
+    def test_yield_must_be_in_fn_context(self, lcompile: CompileFn, code: str):
+        with pytest.raises(compiler.CompilerException):
+            lcompile(code)
+
+    @pytest.mark.parametrize(
+        "code",
+        [
+            "(fn [] (yield :two :items))",
+            "(fn [] (yield :three :items :now))",
+        ],
+    )
+    def test_yield_num_elems(self, lcompile: CompileFn, code: str):
+        with pytest.raises(compiler.CompilerException):
+            lcompile(code)
+
+    def test_yield_control_only(self, lcompile: CompileFn, ns: runtime.Namespace):
+        f = lcompile(
+            """
+            (def state (atom nil))
+            (fn []
+              (reset! state :started)
+              (yield)
+              (reset! state :done))
+            """
+        )
+        state = ns.find(sym.symbol("state")).value
+        coro = f()
+        assert None is state.deref()
+        assert None is next(coro)
+        assert kw.keyword("started") == state.deref()
+        assert None is next(coro, None)
+        assert kw.keyword("done") == state.deref()
+
+    def test_yield_value(self, lcompile: CompileFn, ns: runtime.Namespace):
+        f = lcompile(
+            """
+            (def state (atom nil))
+            (fn []
+              (reset! state :started)
+              (yield :yielding)
+              (reset! state :done))
+            """
+        )
+        state = ns.find(sym.symbol("state")).value
+        coro = f()
+        assert None is state.deref()
+        assert kw.keyword("yielding") == next(coro)
+        assert kw.keyword("started") == state.deref()
+        assert None is next(coro, None)
+        assert kw.keyword("done") == state.deref()
+
+    def test_yield_as_coroutine(self, lcompile: CompileFn, ns: runtime.Namespace):
+        f = lcompile(
+            """
+            (def state (atom nil))
+            (fn []
+              (reset! state :started)
+              (let [v (yield :yielding)]
+                (reset! state v)))
+            """
+        )
+        state = ns.find(sym.symbol("state")).value
+        coro = f()
+        assert None is state.deref()
+        assert kw.keyword("yielding") == next(coro)
+        assert kw.keyword("started") == state.deref()
+        with pytest.raises(StopIteration):
+            coro.send(kw.keyword("coroutine-value"))
+        assert kw.keyword("coroutine-value") == state.deref()
+
+    def test_yield_as_coroutine_with_multiple_yields(
+        self, lcompile: CompileFn, ns: runtime.Namespace
+    ):
+        f = lcompile(
+            """
+            (def state (atom nil))
+            (fn []
+              (reset! state :started)
+              (let [v (yield :yielding)]
+                (reset! state v)
+                (yield)
+                (reset! state :done)))
+            """
+        )
+        state = ns.find(sym.symbol("state")).value
+        coro = f()
+        assert None is state.deref()
+        assert kw.keyword("yielding") == next(coro)
+        assert kw.keyword("started") == state.deref()
+        assert None is coro.send(kw.keyword("coroutine-value"))
+        assert kw.keyword("coroutine-value") == state.deref()
+        assert None is next(coro, None)
+        assert kw.keyword("done") == state.deref()