Implement close method for generators

Author: MaxShvets

mypyc/irbuild/generator.py

@@ -15,12 +15,16 @@
 from mypyc.common import SELF_NAME, NEXT_LABEL_ATTR_NAME, ENV_ATTR_NAME
 from mypyc.ir.ops import (
     BasicBlock, Call, Return, Goto, Integer, SetAttr, Unreachable, RaiseStandardError,
-    Value, Register
+    Value, Register, MethodCall, TupleSet, Branch, NO_TRACEBACK_LINE_NO
 )
 from mypyc.ir.rtypes import RInstance, int_rprimitive, object_rprimitive
 from mypyc.ir.func_ir import FuncIR, FuncDecl, FuncSignature, RuntimeArg
 from mypyc.ir.class_ir import ClassIR
-from mypyc.primitives.exc_ops import raise_exception_with_tb_op
+from mypyc.irbuild.nonlocalcontrol import ExceptNonlocalControl
+from mypyc.primitives.exc_ops import (
+    raise_exception_with_tb_op, error_catch_op, exc_matches_op, reraise_exception_op,
+    restore_exc_info_op
+)
 from mypyc.irbuild.env_class import (
     add_args_to_env, load_outer_env, load_env_registers, finalize_env_class
 )
@@ -220,11 +224,46 @@ def add_throw_to_generator_class(builder: IRBuilder,
 
 def add_close_to_generator_class(builder: IRBuilder, fn_info: FuncInfo) -> None:
     """Generates the '__close__' method for a generator class."""
-    # TODO: Currently this method just triggers a runtime error.
-    #       We should fill this out (https://github.com/mypyc/mypyc/issues/790).
     with builder.enter_method(fn_info.generator_class.ir, 'close', object_rprimitive, fn_info):
+        except_block, else_block = BasicBlock(), BasicBlock()
+        builder.builder.push_error_handler(except_block)
+        builder.goto_and_activate(BasicBlock())
+        generator_exit = builder.load_module_attr_by_fullname('builtins.GeneratorExit',
+                                                              fn_info.fitem.line)
+        builder.add(MethodCall(
+            builder.self(),
+            'throw',
+            [generator_exit, builder.none_object(), builder.none_object()]))
+        builder.goto(else_block)
+        builder.builder.pop_error_handler()
+
+        builder.activate_block(except_block)
+        old_exc = builder.call_c(error_catch_op, [], fn_info.fitem.line)
+        builder.nonlocal_control.append(
+            ExceptNonlocalControl(builder.nonlocal_control[-1], old_exc))
+        stop_iteration = builder.load_module_attr_by_fullname('builtins.StopIteration',
+                                                              fn_info.fitem.line)
+        exceptions = builder.add(
+            TupleSet([generator_exit, stop_iteration], fn_info.fitem.line))
+        matches = builder.call_c(
+            exc_matches_op, [exceptions], fn_info.fitem.line)
+
+        match_block, non_match_block = BasicBlock(), BasicBlock()
+        builder.add(Branch(matches, match_block, non_match_block, Branch.BOOL))
+
+        builder.activate_block(match_block)
+        builder.call_c(restore_exc_info_op, [builder.read(old_exc)], fn_info.fitem.line)
+        builder.add(Return(builder.none_object()))
+
+        builder.activate_block(non_match_block)
+        builder.call_c(reraise_exception_op, [], NO_TRACEBACK_LINE_NO)
+        builder.add(Unreachable())
+
+        builder.nonlocal_control.pop()
+
+        builder.activate_block(else_block)
         builder.add(RaiseStandardError(RaiseStandardError.RUNTIME_ERROR,
-                                       'close method on generator classes unimplemented',
+                                       'generator ignored GeneratorExit',
                                        fn_info.fitem.line))
         builder.add(Unreachable())
 

mypyc/test-data/fixtures/ir.py

@@ -276,6 +276,8 @@ class ArithmeticError(Exception): pass
 
 class ZeroDivisionError(Exception): pass
 
+class GeneratorExit(BaseException): pass
+
 def any(i: Iterable[T]) -> bool: pass
 def all(i: Iterable[T]) -> bool: pass
 def sum(i: Iterable[T]) -> int: pass

mypyc/test-data/run-generators.test

@@ -516,3 +516,68 @@ class E:
 
 [file driver.py]
 # really I only care it builds
+
+[case testCloseStopIterationRaised]
+def g() -> object:
+    try:
+        yield 1
+    except GeneratorExit:
+        raise
+
+[file driver.py]
+from native import g
+
+gen = g()
+next(gen)
+gen.close()
+
+[case testCloseGeneratorExitRaised]
+def g() -> object:
+    yield 1
+
+[file driver.py]
+from native import g
+
+gen = g()
+next(gen)
+gen.close()
+
+[case testCloseGeneratorExitIgnored]
+def g() -> object:
+    try:
+        yield 1
+    except GeneratorExit:
+        pass
+
+    yield 2
+
+[file driver.py]
+from native import g
+
+gen = g()
+next(gen)
+try:
+    gen.close()
+except RuntimeError as e:
+    assert str(e) == 'generator ignored GeneratorExit'
+else:
+    assert False
+
+[case testCloseGeneratorRaisesAnotherException]
+def g() -> object:
+    try:
+        yield 1
+    except GeneratorExit:
+        raise RuntimeError("error")
+
+[file driver.py]
+from native import g
+
+gen = g()
+next(gen)
+try:
+    gen.close()
+except RuntimeError as e:
+    assert str(e) == 'error'
+else:
+    assert False