python · hauntsaninja · Aug 17, 2025 · Aug 17, 2025 · Aug 17, 2025 · Aug 17, 2025
diff --git 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
@@ -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(