diff --git a/mypyc/irbuild/statement.py b/mypyc/irbuild/statement.py index eeeb40ac672f..9066ed1deeac 100644 --- a/mypyc/irbuild/statement.py +++ b/mypyc/irbuild/statement.py @@ -101,6 +101,7 @@ ) from mypyc.primitives.exc_ops import ( error_catch_op, + error_clear_op, exc_matches_op, get_exc_info_op, get_exc_value_op, @@ -887,6 +888,19 @@ def transform_with( is_async: bool, line: int, ) -> None: + + if ( + not is_async + and isinstance(expr, mypy.nodes.CallExpr) + and isinstance(expr.callee, mypy.nodes.RefExpr) + and isinstance(dec := expr.callee.node, mypy.nodes.Decorator) + and len(dec.decorators) == 1 + and isinstance(dec1 := dec.decorators[0], mypy.nodes.RefExpr) + and dec1.node + and dec1.node.fullname == "contextlib.contextmanager" + ): + return _transform_with_contextmanager(builder, expr, target, body, line) + # This is basically a straight transcription of the Python code in PEP 343. # I don't actually understand why a bunch of it is the way it is. # We could probably optimize the case where the manager is compiled by us, @@ -964,6 +978,133 @@ def finally_body() -> None: ) +def _transform_with_contextmanager( + builder: IRBuilder, + expr: mypy.nodes.CallExpr, + target: Lvalue | None, + with_body: GenFunc, + line: int, +) -> None: + assert isinstance(expr.callee, mypy.nodes.RefExpr) + dec = expr.callee.node + assert isinstance(dec, mypy.nodes.Decorator) + + # mgrv = ctx.__wrapped__(*args, **kwargs) + wrapped_call = mypy.nodes.CallExpr( + mypy.nodes.MemberExpr(expr.callee, "__wrapped__"), + expr.args, + expr.arg_kinds, + expr.arg_names, + ) + wrapped_call.line = line + gen = builder.accept(wrapped_call) + + # try: + # target = next(gen) + # except StopIteration: + # raise RuntimeError("generator didn't yield") from None + mgr_target = builder.call_c(next_raw_op, [gen], line) + + runtime_block, main_block = BasicBlock(), BasicBlock() + builder.add(Branch(mgr_target, runtime_block, main_block, Branch.IS_ERROR)) + + builder.activate_block(runtime_block) + builder.add( + RaiseStandardError(RaiseStandardError.RUNTIME_ERROR, "generator didn't yield", line) + ) + builder.add(Unreachable()) + + builder.activate_block(main_block) + + # try: + # {body} + + def try_body() -> None: + if target: + builder.assign(builder.get_assignment_target(target), mgr_target, line) + with_body() + + # except Exception as e: + # try: + # gen.throw(e) + # except StopIteration as e2: + # if e2 is not e: + # raise + # return + # except RuntimeError: + # # TODO: check the traceback munging + # raise + # except BaseException: + # # approximately + # raise + + def except_body() -> None: + exc_original = builder.call_c(get_exc_value_op, [], line) + + error_block, no_error_block = BasicBlock(), BasicBlock() + + builder.builder.push_error_handler(error_block) + builder.goto_and_activate(BasicBlock()) + builder.py_call(builder.py_get_attr(gen, "throw", line), [exc_original], line) + builder.goto(no_error_block) + builder.builder.pop_error_handler() + + builder.activate_block(no_error_block) + builder.add( + RaiseStandardError( + RaiseStandardError.RUNTIME_ERROR, "generator didn't stop after throw()", line + ) + ) + builder.add(Unreachable()) + + builder.activate_block(error_block) + stop_iteration = builder.call_c(check_stop_op, [], line) + is_same_exc = builder.binary_op(stop_iteration, exc_original, "==", line) + + suppress_block, propagate_block = BasicBlock(), BasicBlock() + builder.add(Branch(is_same_exc, suppress_block, propagate_block, Branch.BOOL)) + + builder.activate_block(propagate_block) + builder.call_c(keep_propagating_op, [], line) + builder.add(Unreachable()) + + builder.activate_block(suppress_block) + builder.call_c(error_clear_op, [], -1) + + # TODO: actually do the exceptions + handlers = [(None, None, except_body)] + + # else: + # try: + # next(gen) + # except StopIteration: + # pass + # else: + # try: + # raise RuntimeError("generator didn't stop") + # finally: + # gen.close() + + def else_body() -> None: + value = builder.call_c(next_raw_op, [builder.read(gen)], line) + stop_block, close_block = BasicBlock(), BasicBlock() + builder.add(Branch(value, stop_block, close_block, Branch.IS_ERROR)) + + builder.activate_block(close_block) + # TODO: this isn't exactly the right order + builder.py_call(builder.py_get_attr(gen, "close", line), [], line) + builder.add( + RaiseStandardError(RaiseStandardError.RUNTIME_ERROR, "generator didn't stop", line) + ) + builder.add(Unreachable()) + + builder.activate_block(stop_block) + # TODO: should check for StopIteration + builder.call_c(error_clear_op, [], -1) + + transform_try_except(builder, try_body, handlers, else_body, line) + + def transform_with_stmt(builder: IRBuilder, o: WithStmt) -> None: # Generate separate logic for each expr in it, left to right def generate(i: int) -> None: diff --git a/mypyc/primitives/exc_ops.py b/mypyc/primitives/exc_ops.py index e1234f807afa..01258712d3ad 100644 --- a/mypyc/primitives/exc_ops.py +++ b/mypyc/primitives/exc_ops.py @@ -80,6 +80,11 @@ arg_types=[], return_type=exc_rtuple, c_function_name="CPy_CatchError", error_kind=ERR_NEVER ) +# Clear the current exception. +error_clear_op = custom_op( + arg_types=[], return_type=void_rtype, c_function_name="PyErr_Clear", error_kind=ERR_NEVER +) + # Restore an old "currently handled exception" returned from. # error_catch (by sticking it into sys.exc_info()) restore_exc_info_op = custom_op(