[mypyc] add support for contextlib.contextmanager

mypyc/irbuild/statement.py

@@ -49,6 +49,7 @@
     YieldExpr,
     YieldFromExpr,
 )
+from mypy.types import Instance, get_proper_type
 from mypyc.common import (
     GENERATOR_HELPER_NAME,
     KEEP_ALIVE_SHORT_LIVED,
@@ -114,6 +115,10 @@
     AssignmentTargetTuple,
 )
 from mypyc.primitives.exc_ops import (
+    contextmanager_bind_op,
+    contextmanager_enter_op,
+    contextmanager_exit_exception_op,
+    contextmanager_exit_normal_op,
     error_catch_op,
     exc_matches_op,
     get_exc_info_op,
@@ -1125,18 +1130,44 @@ def transform_with(
     is_async: bool,
     line: int,
 ) -> None:
+
+    if isinstance(expr, mypy.nodes.CallExpr) and isinstance(expr.callee, mypy.nodes.RefExpr):
+        callee = expr.callee
+        ctm_def_node: mypy.nodes.SymbolNode | None = callee.node
+        self_expr: Expression | None = None
+        if isinstance(callee, mypy.nodes.MemberExpr) and not builder.is_module_member_expr(callee):
+            ctm_def_node = None
+            # Since we construct the call to the method ourselves in _transform_with_contextmanager,
+            # we're a little paranoid about doing things in the wrong order in weird cases
+            if not find_walrus_targets(expr) and not expr_has_suspend(expr):
+                self_type = get_proper_type(builder.types[callee.expr])
+                if isinstance(self_type, Instance):
+                    symbol = self_type.type.get(callee.name)
+                    if symbol is not None:
+                        ctm_def_node = symbol.node
+                        self_expr = callee.expr
+
+        if (
+            not is_async
+            and isinstance(dec := ctm_def_node, mypy.nodes.Decorator)
+            and not (self_expr is not None and (dec.func.is_static or dec.func.is_class))
+            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, self_expr)
+
     # 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,
-    # but that is not our common case at all, so.
 
     al = "a" if is_async else ""
 
     mgr_v = builder.accept(expr)
 
     if not is_async and mgr_v.type == lock_rprimitive:
         transform_with_lock(builder, mgr_v, target, body, line)
-        return
+        return None
     # The methods may be inherited from a non-native class, and then they can't be
     # called natively.
     is_native = (
@@ -1240,6 +1271,87 @@ def finally_body() -> None:
     transform_try_finally_stmt(builder, try_body, finally_body, line)
 
 
+def _transform_with_contextmanager(
+    builder: IRBuilder,
+    expr: mypy.nodes.CallExpr,
+    target: Lvalue | None,
+    with_body: GenFunc,
+    line: int,
+    receiver_expr: Expression | None,
+) -> None:
+
+    # wrapped_call = {expr.callee}.__wrapped__(*args, **kwargs)
+    # where {expr.callee} is either a function or a method
+
+    assert isinstance(expr.callee, mypy.nodes.RefExpr)
+    if receiver_expr is None:
+        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)
+    else:
+        assert isinstance(expr.callee, mypy.nodes.MemberExpr)
+        receiver = builder.accept(receiver_expr)
+        method = builder.py_get_attr(receiver, expr.callee.name, line)
+        wrapped = builder.py_get_attr(method, "__wrapped__", line)
+        wrapped = builder.call_c(contextmanager_bind_op, [method, wrapped], line)
+        args = [builder.accept(arg) for arg in expr.args]
+        gen = builder.py_call(wrapped, args, line, expr.arg_kinds, expr.arg_names)
+
+    # try:
+    #     mgr_target = next(gen)
+    # except StopIteration:
+    #     raise RuntimeError("generator didn't yield") from None
+    mgr_target = builder.call_c(contextmanager_enter_op, [gen], line)
+
+    # try:
+    #     {body}
+
+    def try_body() -> None:
+        if target:
+            builder.assign(builder.get_assignment_target(target), mgr_target, line)
+        with_body()
+
+    # except BaseException as e:
+    #     {contextmanager_exit_exception_op}
+    # finally:
+    #     {contextmanager_exit_normal_op}
+
+    exc = builder.ensure_register(builder.true())
+
+    def except_body() -> None:
+        builder.assign(exc, builder.false(), line)
+        suppressed = builder.call_c(
+            contextmanager_exit_exception_op, [builder.read(gen), *get_sys_exc_info(builder)], line
+        )
+        out_block, reraise_block = BasicBlock(), BasicBlock()
+        builder.add_bool_branch(suppressed, out_block, reraise_block)
+        builder.activate_block(reraise_block)
+        builder.call_c(reraise_exception_op, [], NO_TRACEBACK_LINE_NO)
+        builder.add(Unreachable())
+        builder.activate_block(out_block)
+
+    handlers = [(None, None, except_body)]
+
+    def finally_body() -> None:
+        out_block, exit_block = BasicBlock(), BasicBlock()
+        builder.add(Branch(builder.read(exc), exit_block, out_block, Branch.BOOL))
+        builder.activate_block(exit_block)
+        builder.call_c(contextmanager_exit_normal_op, [builder.read(gen)], line)
+        builder.goto_and_activate(out_block)
+
+    transform_try_finally_stmt(
+        builder,
+        lambda: transform_try_except(builder, try_body, handlers, None, line),
+        finally_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:

mypyc/lib-rt/CPy.h

@@ -1071,6 +1071,11 @@ bool CPy_ExceptionMatches(PyObject *type);
 PyObject *CPy_GetExcValue(void);
 tuple_T3OOO CPy_GetExcInfo(void);
 void _CPy_GetExcInfo(PyObject **p_type, PyObject **p_value, PyObject **p_traceback);
+PyObject *CPyContextManager_Bind(PyObject *callable, PyObject *wrapped);
+PyObject *CPyContextManager_Enter(PyObject *gen);
+bool CPyContextManager_ExitNormal(PyObject *gen);
+int CPyContextManager_ExitException(PyObject *gen, PyObject *type,
+                                  PyObject *value, PyObject *traceback);
 void CPyError_OutOfMemory(void);
 void CPy_TypeError(const char *expected, PyObject *value);
 void CPy_AddTraceback(const char *filename, const char *funcname, int line, PyObject *globals);

mypyc/lib-rt/exc_ops.c

@@ -163,6 +163,206 @@ tuple_T3OOO CPy_GetExcInfo(void) {
     return ret;
 }
 
+// Give a contextmanager's generator factory the same binding as its wrapper.
+// The returned callable is owned; static methods remain unbound.
+PyObject *CPyContextManager_Bind(PyObject *callable, PyObject *wrapped) {
+    if (PyMethod_Check(callable)) {
+        // Use the actual bound instance or class, not the attribute lookup's receiver.
+        return PyMethod_New(wrapped, PyMethod_GET_SELF(callable));
+    }
+    Py_INCREF(wrapped);
+    return wrapped;
+}
+
+// Implement contextlib's generator entry protocol. The returned value is owned.
+PyObject *CPyContextManager_Enter(PyObject *gen) {
+    if (!PyIter_Check(gen)) {
+        PyErr_Format(PyExc_TypeError, "'%.200s' object is not an iterator", Py_TYPE(gen)->tp_name);
+        return NULL;
+    }
+    PyObject *value = Py_TYPE(gen)->tp_iternext(gen);
+    if (value != NULL) {
+        return value;
+    }
+    if (!PyErr_Occurred()) {
+        // Iterators may signal exhaustion without setting StopIteration.
+        PyErr_SetNone(PyExc_StopIteration);
+    }
+    if (PyErr_ExceptionMatches(PyExc_StopIteration)) {
+        // raise RuntimeError(...) from None, retaining the suppressed context.
+        tuple_T3OOO old_exc = CPy_CatchError();
+        PyObject *error = PyObject_CallFunction(PyExc_RuntimeError, "s", "generator didn't yield");
+        if (error != NULL) {
+            CPy_RaiseFrom(error, Py_None);
+            Py_DECREF(error);
+        }
+        CPy_RestoreExcInfo(old_exc);
+        Py_DECREF(old_exc.f0);
+        Py_DECREF(old_exc.f1);
+        Py_DECREF(old_exc.f2);
+    }
+    return NULL;
+}
+
+// An extra yield is a protocol error. Always leave an exception pending.
+static void contextmanager_exit_error(PyObject *gen, const char *message) {
+    PyErr_SetString(PyExc_RuntimeError, message);
+#if PY_VERSION_HEX >= 0x030B0000
+    // Since Python 3.11, contextlib closes the generator in a finally block
+    tuple_T3OOO old_exc = CPy_CatchError();
+    PyObject *type, *value, *traceback;
+    PyErr_GetExcInfo(&type, &value, &traceback);
+    PyObject *closed = PyObject_CallMethod(gen, "close", NULL);
+    if (closed != NULL) {
+        Py_DECREF(closed);
+        // A native generator can change the handled exception during close()
+        // Reraise the saved protocol error, as a Python finally block does
+        PyErr_Restore(type, value, traceback);
+    } else {
+        Py_XDECREF(type);
+        Py_XDECREF(value);
+        Py_XDECREF(traceback);
+    }
+    CPy_RestoreExcInfo(old_exc);
+    Py_DECREF(old_exc.f0);
+    Py_DECREF(old_exc.f1);
+    Py_DECREF(old_exc.f2);
+#endif
+}
+
+// Finish a contextmanager whose body succeeded. Return false on error.
+bool CPyContextManager_ExitNormal(PyObject *gen) {
+    // try:
+    //     next(self.gen)
+    if (!PyIter_Check(gen)) {
+        PyErr_Format(PyExc_TypeError, "'%.200s' object is not an iterator", Py_TYPE(gen)->tp_name);
+        return false;
+    }
+    PyObject *value = Py_TYPE(gen)->tp_iternext(gen);
+
+    // except StopIteration:
+    //     return False
+    if (value == NULL) {
+        if (!PyErr_Occurred() || PyErr_ExceptionMatches(PyExc_StopIteration)) {
+            PyErr_Clear();
+            return true;
+        }
+        return false;
+    }
+
+    // else:
+    //     try:
+    //         raise RuntimeError("generator didn't stop")
+    //     finally:
+    //         self.gen.close()
+    Py_DECREF(value);
+    contextmanager_exit_error(gen, "generator didn't stop");
+    return false;
+}
+
+// Match contextlib._GeneratorContextManager.__exit__ for a failed with body.
+int CPyContextManager_ExitException(PyObject *gen, PyObject *type,
+                                    PyObject *value, PyObject *traceback) {
+
+    // Native generators can change the handled exception when throw() unwinds
+    // an except block. Preserve the caller's state before interpreting the result
+    PyObject *saved_type, *saved_value, *saved_traceback;
+    PyErr_GetExcInfo(&saved_type, &saved_value, &saved_traceback);
+
+    // try:
+    //     self.gen.throw(typ, value, traceback)
+#if PY_VERSION_HEX >= 0x030C0000
+    PyObject *yielded = PyObject_CallMethod(gen, "throw", "(O)", value);
+#else
+    PyObject *yielded = PyObject_CallMethod(gen, "throw", "(OOO)", type, value, traceback);
+#endif
+    PyErr_SetExcInfo(saved_type, saved_value, saved_traceback);
+
+    // ... if we didn't get an exception from throw, in particular, StopIteration ...
+    if (yielded != NULL) {
+        // try:
+        //     raise RuntimeError("generator didn't stop after throw()")
+        // finally:
+        //     self.gen.close()
+        Py_DECREF(yielded);
+        contextmanager_exit_error(gen, "generator didn't stop after throw()");
+        return -1;
+    }
+
+    tuple_T3OOO old_exc = CPy_CatchError();
+    PyObject *error = CPy_GetExcValue();
+    int result = -1;
+
+    // # Note the `return False` in the code comments is how contextlib tells Python to
+    // # propagate the exception being handled by __exit__
+
+    // except StopIteration as exc:
+    //     return exc is not value
+    if (CPy_ExceptionMatches(PyExc_StopIteration)) {
+        // A body StopIteration must not be mistaken for generator completion.
+        result = error != value;
+        goto done;
+    }
+    // except RuntimeError as exc:
+    //     if exc is value:
+    //         exc.__traceback__ = traceback
+    //         return False
+    //     if (
+    //         isinstance(value, StopIteration)
+    //         and exc.__cause__ is value
+    //     ):
+    //         value.__traceback__ = traceback
+    //         return False
+    //     raise
+    if (CPy_ExceptionMatches(PyExc_RuntimeError)) {
+        if (error == value) {
+            goto original;
+        }
+        // PEP 479 can wrap the original StopIteration (including subclasses).
+        int is_stop = PyObject_IsInstance(value, PyExc_StopIteration);
+        if (is_stop < 0) {
+            goto done;
+        }
+        if (is_stop) {
+            PyObject *cause = PyObject_GetAttrString(error, "__cause__");
+            if (cause == NULL) {
+                goto done;
+            }
+            bool is_original = cause == value;
+            Py_DECREF(cause);
+            if (is_original) {
+                goto original;
+            }
+        }
+    }
+    // except BaseException as exc:
+    //     if exc is not value:
+    //         raise
+    //     exc.__traceback__ = traceback
+    //     return False
+    else if (error == value) {
+        goto original;
+    }
+    CPy_Reraise();
+    goto done;
+
+original:
+#if PY_VERSION_HEX >= 0x030B0000
+    if (PyObject_SetAttrString(value, "__traceback__", traceback) < 0) {
+        goto done;
+    }
+#endif
+    result = 0;
+
+done:
+    Py_DECREF(error);
+    CPy_RestoreExcInfo(old_exc);
+    Py_DECREF(old_exc.f0);
+    Py_DECREF(old_exc.f1);
+    Py_DECREF(old_exc.f2);
+    return result;
+}
+
 void CPyError_OutOfMemory(void) {
     fprintf(stderr, "fatal: out of memory\n");
     fflush(stderr);