[mypyc] Fix function wrapper descriptor for static/class methods (#21811)

Calling an `async` class method in interpreted context might fail on
incorrect number of arguments because the class `T` on the lhs of the
member expression `T.f(a,b,c)` is not bound to the `cls` argument.
Instead the first regular argument is taken as `cls` and all other are
shifted.

Expand the `tp_descr_get` slot of the function wrapper type object to
correctly handle class methods.

Also modify the function lowering logic to take the method from the type
dictionary instead of through `getattr` when calling a decorator on the
method. `getattr` goes through the descriptor and binds the `cls`
argument but the argument should be unbound when applying the decorator
to match interpreted Python.

Author: p-sawicki

mypyc/codegen/emit.py

@@ -22,7 +22,13 @@
     TYPE_VAR_PREFIX,
 )
 from mypyc.ir.class_ir import ClassIR, all_concrete_classes
-from mypyc.ir.func_ir import FUNC_STATICMETHOD, FuncDecl, FuncIR, get_text_signature
+from mypyc.ir.func_ir import (
+    FUNC_CLASSMETHOD,
+    FUNC_STATICMETHOD,
+    FuncDecl,
+    FuncIR,
+    get_text_signature,
+)
 from mypyc.ir.ops import (
     NAMESPACE_MODULE,
     NAMESPACE_STATIC,
@@ -1419,13 +1425,17 @@ def emit_cpyfunction_instance(
         cname = f"{PREFIX}{fn.cname(self.names)}"
         wrapper_name = f"{cname}_wrapper"
         cfunc = f"(PyCFunction){cname}"
-        func_flags = "METH_FASTCALL | METH_KEYWORDS"
+        func_flags = ["METH_FASTCALL", "METH_KEYWORDS"]
+        if fn.class_name and fn.decl.kind == FUNC_STATICMETHOD:
+            func_flags.append("METH_STATIC")
+        elif fn.class_name and fn.decl.kind == FUNC_CLASSMETHOD:
+            func_flags.append("METH_CLASS")
         doc = f"PyDoc_STR({native_function_doc_initializer(fn)})"
         has_self_arg = "true" if fn.class_name and fn.decl.kind != FUNC_STATICMETHOD else "false"
 
         code_flags = "CO_COROUTINE"
         self.emit_line(
-            f'PyObject* {wrapper_name} = CPyFunction_New({module}, "{filepath}", "{name}", {cfunc}, {func_flags}, {doc}, {fn.line}, {code_flags}, {has_self_arg});'
+            f'PyObject* {wrapper_name} = CPyFunction_New({module}, "{filepath}", "{name}", {cfunc}, {" | ".join(func_flags)}, {doc}, {fn.line}, {code_flags}, {has_self_arg});'
         )
         self.emit_line(f"if (unlikely(!{wrapper_name}))")
         self.emit_line(error_stmt)

mypyc/irbuild/function.py

@@ -82,7 +82,12 @@
     dict_new_op,
     exact_dict_set_item_op,
 )
-from mypyc.primitives.generic_ops import generic_getattr, generic_setattr, py_setattr_op
+from mypyc.primitives.generic_ops import (
+    generic_getattr,
+    generic_setattr,
+    py_get_item_op,
+    py_setattr_op,
+)
 from mypyc.primitives.misc_ops import register_function
 from mypyc.sametype import is_same_method_signature, is_same_type
 
@@ -486,14 +491,30 @@ def handle_ext_method(builder: IRBuilder, cdef: ClassDef, fdef: FuncDef) -> None
     if is_decorated(builder, fdef):
         # Obtain the function name in order to construct the name of the helper function.
         _, _, name = fdef.fullname.rpartition(".")
-        # Read the PyTypeObject representing the class, get the callable object
-        # representing the non-decorated method
+        # Get the callable representing the non-decorated method directly from the type
+        # dictionary. Attribute access would bind a class method before its decorators are
+        # applied, but the decorators need to receive the unbound function.
         typ = builder.load_native_type_object(cdef.fullname)
-        orig_func = builder.py_get_attr(typ, name, fdef.line)
+        type_dict = builder.py_get_attr(typ, "__dict__", fdef.line)
+        orig_func = builder.primitive_op(
+            py_get_item_op, [type_dict, builder.load_str(name)], fdef.line
+        )
 
         # Decorate the non-decorated method
         decorated_func = load_decorated_func(builder, fdef, orig_func)
 
+        # @classmethod and @staticmethod aren't included in fdefs_to_decorators, since
+        # mypy represents them using the function kind. Reapply the outer descriptor
+        # after the other decorators, matching Python's decorator evaluation order.
+        # TODO: Handle cases where @classmethod/@staticmethod are the inner decorator.
+        # See mypyc#1208 for reference.
+        if func_ir.decl.kind == FUNC_CLASSMETHOD:
+            cls_meth = builder.load_module_attr_by_fullname("builtins.classmethod", fdef.line)
+            decorated_func = builder.py_call(cls_meth, [decorated_func], fdef.line)
+        elif func_ir.decl.kind == FUNC_STATICMETHOD:
+            stat_meth = builder.load_module_attr_by_fullname("builtins.staticmethod", fdef.line)
+            decorated_func = builder.py_call(stat_meth, [decorated_func], fdef.line)
+
         # Set the callable object representing the decorated method as an attribute of the
         # extension class.
         builder.primitive_op(

mypyc/lib-rt/function_wrapper.c

@@ -145,9 +145,19 @@ static PyGetSetDef CPyFunction_getsets[] = {
     {0, 0, 0, 0, 0}
 };
 
-static PyObject* CPy_PyMethod_New(PyObject *func, PyObject *self, PyObject *typ) {
-    (void)typ;
-    if (!self) {
+static PyObject* CPyFunction_descr_get(PyObject *func, PyObject *self, PyObject *typ) {
+    int flags = ((PyCFunctionObject *)func)->m_ml->ml_flags;
+    if (flags & METH_CLASS) {
+        if (typ == NULL) {
+            if (self == NULL) {
+                PyErr_SetString(PyExc_TypeError, "__get__(None, None) is invalid");
+                return NULL;
+            }
+            typ = (PyObject *)Py_TYPE(self);
+        }
+        return PyMethod_New(func, typ);
+    }
+    if (!self || (flags & METH_STATIC)) {
         Py_INCREF(func);
         return func;
     }
@@ -162,7 +172,7 @@ static PyType_Slot CPyFunction_slots[] = {
     {Py_tp_clear, (void *)CPyFunction_clear},
     {Py_tp_members, (void *)CPyFunction_members},
     {Py_tp_getset, (void *)CPyFunction_getsets},
-    {Py_tp_descr_get, (void *)CPy_PyMethod_New},
+    {Py_tp_descr_get, (void *)CPyFunction_descr_get},
     {0, 0},
 };
 

mypyc/test-data/run-async.test

@@ -1554,6 +1554,18 @@ def wrap(fn: F) -> F:
 
     return cast(F, wrapper)
 
+C = Callable[..., Any]
+
+def get_decorator() -> C:
+    def decorator(endpoint: C) -> C:
+        @wraps(endpoint)
+        async def inner(self: Any, val: int, *args: Any) -> Any:
+            return await endpoint(self, val, *args)
+
+        return cast(C, inner)
+
+    return cast(C, decorator)
+
 @wrap
 def wrapped(val: int) -> int:
     return val
@@ -1571,34 +1583,57 @@ async def wrapped2_async(val: int) -> int:
     return val * 2
 
 class T:
-    def returns_one(self) -> int:
+    def returns_one(self, val: int) -> int:
         return 1
 
-    async def returns_one_async(self) -> int:
+    async def returns_one_async(self, val: int) -> int:
         return 1
 
     @wrap
-    def returns_two(self) -> int:
+    def returns_two(self, val: int) -> int:
         return 1
 
     @wrap
-    async def returns_two_async(self) -> int:
+    async def returns_two_async(self, val: int) -> int:
         return 1
 
     @staticmethod
-    def static() -> int:
+    def static(val: int) -> int:
+        return 2
+
+    @staticmethod
+    async def static_async(val: int) -> int:
+        return 2
+
+    @staticmethod
+    @wrap
+    def wrapped_static(val: int) -> int:
         return 2
 
     @staticmethod
-    async def static_async() -> int:
+    @wrap
+    async def wrapped_static_async(val: int) -> int:
         return 2
 
     @classmethod
-    def class_method(cls) -> int:
+    def class_method(cls, val: int) -> int:
         return 3
 
     @classmethod
-    async def class_method_async(cls) -> int:
+    async def class_method_async(cls, val: int) -> int:
+        assert cls is T
+        return 3
+
+    @classmethod
+    @wrap
+    def wrapped_class_method(cls, val: int) -> int:
+        assert cls is T
+        return 3
+
+    @classmethod
+    @get_decorator()
+    async def wrapped_class_method_async(cls, val: int) -> int:
+        assert cls is T
         return 3
 
 def is_coroutine(fn):
@@ -1653,38 +1688,50 @@ def test_method() -> None:
     t = T()
     # Call through variable to make sure the call is through vectorcall and not optimized to a native call.
     f: Any = t.returns_one_async
-    assert asyncio.run(f()) == 1
+    assert asyncio.run(f(1)) == 1
 
     assert not is_coroutine(T.returns_two)
     assert is_coroutine(T.returns_two_async)
-    assert asyncio.run(t.returns_two_async()) == 2
+    assert asyncio.run(t.returns_two_async(2)) == 2
 
     assert not is_coroutine(T.static)
     assert is_coroutine(T.static_async)
-    assert asyncio.run(T.static_async()) == 2
+    assert asyncio.run(T.static_async(3)) == 2
+    assert not is_coroutine(T.wrapped_static)
+    assert is_coroutine(T.wrapped_static_async)
+    assert T.wrapped_static(3) == 4
+    assert t.wrapped_static(3) == 4
+    assert asyncio.run(T.wrapped_static_async(3)) == 4
+    assert asyncio.run(t.wrapped_static_async(3)) == 4
 
     assert not is_coroutine(T.class_method)
     assert is_coroutine(T.class_method_async)
-    assert asyncio.run(T.class_method_async()) == 3
+    assert asyncio.run(T.class_method_async(4)) == 3
+    assert not is_coroutine(T.wrapped_class_method)
+    assert is_coroutine(T.wrapped_class_method_async)
+    assert T.wrapped_class_method(4) == 6
+    assert t.wrapped_class_method(4) == 6
+    assert asyncio.run(T.wrapped_class_method_async(4)) == 3
+    assert asyncio.run(t.wrapped_class_method_async(4)) == 3
 
 def test_nested() -> None:
-    def nested() -> int:
+    def nested(val: int) -> int:
         return 1
 
-    async def nested_async() -> int:
+    async def nested_async(val: int) -> int:
         return 1
 
     @wrap
-    def nested_wrapped() -> int:
+    def nested_wrapped(val: int) -> int:
         return 2
 
     @wrap
-    async def nested_wrapped_async() -> int:
+    async def nested_wrapped_async(val: int) -> int:
         return 2
 
     assert not is_coroutine(nested)
     assert is_coroutine(nested_async)
-    assert asyncio.run(nested_async()) == 1
+    assert asyncio.run(nested_async(1)) == 1
 
     assert getattr(nested_async, "__name__") == "nested_async", getattr(nested_async, "__name__")
     setattr(nested_async, "__name__", "some custom name")
@@ -1696,7 +1743,7 @@ def test_nested() -> None:
 
     assert not is_coroutine(nested_wrapped)
     assert is_coroutine(nested_wrapped_async)
-    assert asyncio.run(nested_wrapped_async()) == 4
+    assert asyncio.run(nested_wrapped_async(2)) == 4
 
 def test_async_function_wrapper_code_refcount() -> None:
     if is_gil_disabled():
@@ -1729,6 +1776,102 @@ def test_nested_async_function_wrapper_code_refcount() -> None:
     assert before == after + 1, (before, after)
     assert after == 1, after
 
+[file driver.py]
+import asyncio
+import sys
+import weakref
+
+import native
+
+def test_function() -> None:
+    native.identity_async.__name__ = "identity_async"
+    native.wrapped_async.__name__ = "wrapped_async"
+
+    assert not native.is_coroutine(native.identity)
+    assert native.is_coroutine(native.identity_async)
+    assert str(native.identity_async).startswith("<function identity_async"), str(native.identity_async)
+    assert asyncio.run(native.identity_async(42)) == 42
+
+    wr = weakref.ref(native.identity_async)
+    f = wr()
+    assert f
+    assert asyncio.run(f(43)) == 43
+
+    assert getattr(native.identity_async, "__name__") == "identity_async"
+    assert getattr(native.identity_async, "__code__") is not None
+    assert getattr(native.identity_async, "__defaults__") is None
+    assert getattr(native.identity_async, "__kwdefaults__") is None
+    assert getattr(native.identity_async, "__annotations__") is None
+
+    assert not native.is_coroutine(native.wrapped)
+    assert native.is_coroutine(native.wrapped_async)
+    assert asyncio.run(native.wrapped_async(22)) == 44
+
+    assert getattr(native.wrapped, "__name__") == "wrapped"
+    assert getattr(native.wrapped2, "__name__") == "wrapped2"
+    assert getattr(native.wrapped_async, "__name__") == "wrapped_async"
+    assert getattr(native.wrapped2_async, "__name__") == "wrapped2_async"
+
+def test_method() -> None:
+    assert not native.is_coroutine(native.T.returns_one)
+    assert native.is_coroutine(native.T.returns_one_async)
+    assert str(native.T.returns_one_async).startswith("<function T.returns_one_async")
+
+    t = native.T()
+    f = t.returns_one_async
+    assert asyncio.run(f(1)) == 1
+
+    assert not native.is_coroutine(native.T.returns_two)
+    assert native.is_coroutine(native.T.returns_two_async)
+    assert asyncio.run(t.returns_two_async(2)) == 2
+
+    assert not native.is_coroutine(native.T.static)
+    assert native.is_coroutine(native.T.static_async)
+    assert asyncio.run(native.T.static_async(3)) == 2
+    assert asyncio.run(t.static_async(3)) == 2
+    assert not native.is_coroutine(native.T.wrapped_static)
+    assert native.is_coroutine(native.T.wrapped_static_async)
+    assert native.T.wrapped_static(3) == 4
+    assert t.wrapped_static(3) == 4
+    assert asyncio.run(native.T.wrapped_static_async(3)) == 4
+    assert asyncio.run(t.wrapped_static_async(3)) == 4
+
+    assert not native.is_coroutine(native.T.class_method)
+    assert native.is_coroutine(native.T.class_method_async)
+    assert asyncio.run(native.T.class_method_async(4)) == 3
+    assert asyncio.run(t.class_method_async(4)) == 3
+    assert not native.is_coroutine(native.T.wrapped_class_method)
+    assert native.is_coroutine(native.T.wrapped_class_method_async)
+    assert native.T.wrapped_class_method(4) == 6
+    assert t.wrapped_class_method(4) == 6
+    assert asyncio.run(native.T.wrapped_class_method_async(4)) == 3
+    assert asyncio.run(t.wrapped_class_method_async(4)) == 3
+
+def test_nested() -> None:
+    def nested(val: int) -> int:
+        return 1
+
+    async def nested_async(val: int) -> int:
+        return 1
+
+    nested_wrapped = native.wrap(nested)
+    nested_wrapped_async = native.wrap(nested_async)
+
+    assert not native.is_coroutine(nested_wrapped)
+    assert native.is_coroutine(nested_wrapped_async)
+    assert nested_wrapped(1) == 2
+    assert asyncio.run(nested_wrapped_async(2)) == 2
+
+native.test_function()
+native.test_method()
+native.test_nested()
+native.test_async_function_wrapper_code_refcount()
+native.test_nested_async_function_wrapper_code_refcount()
+
+test_function()
+test_nested()
+test_method()
+
 [file asyncio/__init__.pyi]
 def run(x: object) -> object: ...