[mypyc] Support decorated methods that override native methods

mypyc/codegen/emitclass.py

@@ -37,6 +37,7 @@
     PREFIX,
     REG_PREFIX,
     RUNNING_FIELD,
+    UNDECORATED_PREFIX,
     short_id_from_name,
 )
 from mypyc.ir.class_ir import ClassIR, VTableEntries
@@ -1114,7 +1115,7 @@ def generate_methods_table(
     for fn in cl.methods.values():
         if fn.decl.is_prop_setter or fn.decl.is_prop_getter or fn.internal:
             continue
-        emitter.emit_line(f'{{"{fn.name}",')
+        emitter.emit_line(f'{{"{fn.name.removeprefix(UNDECORATED_PREFIX)}",')
         emitter.emit_line(f" (PyCFunction){PREFIX}{fn.cname(emitter.names)},")
         flags = ["METH_FASTCALL", "METH_KEYWORDS"]
         if fn.decl.kind == FUNC_STATICMETHOD:

mypyc/common.py

@@ -27,6 +27,8 @@
 TEMP_ATTR_NAME: Final = "__mypyc_temp__"
 LAMBDA_NAME: Final = "__mypyc_lambda__"
 PROPSET_PREFIX: Final = "__mypyc_setter__"
+# Key in ClassIR.methods for the undecorated version of a decorated method override
+UNDECORATED_PREFIX: Final = "__mypyc_undecorated__"
 SELF_NAME: Final = "__mypyc_self__"
 MYPYC_DEFAULTS_SETUP: Final = "__mypyc_defaults_setup"
 GENERATOR_ATTRIBUTE_PREFIX: Final = "__mypyc_generator_attribute__"

mypyc/irbuild/expression.py

@@ -617,6 +617,10 @@ def translate_super_method_call(builder: IRBuilder, expr: CallExpr, callee: Supe
         return translate_call(builder, expr, callee)
 
     decl = base.method_decl(callee.name)
+    if decl.internal:
+        # A decorated override, which must be called through the Python API so that
+        # the lookup starts from the right class
+        return translate_call(builder, expr, callee)
     arg_values = [builder.accept(arg) for arg in expr.args]
     arg_kinds, arg_names = expr.arg_kinds.copy(), expr.arg_names.copy()
 

mypyc/irbuild/function.py

@@ -29,7 +29,7 @@
     Var,
 )
 from mypy.types import CallableType, Type, UnboundType, get_proper_type
-from mypyc.common import FAST_PREFIX, LAMBDA_NAME, PROPSET_PREFIX, SELF_NAME
+from mypyc.common import FAST_PREFIX, LAMBDA_NAME, PROPSET_PREFIX, SELF_NAME, UNDECORATED_PREFIX
 from mypyc.ir.class_ir import ClassIR, NonExtClassInfo
 from mypyc.ir.func_ir import (
     FUNC_CLASSMETHOD,
@@ -535,6 +535,24 @@ def handle_ext_method(builder: IRBuilder, cdef: ClassDef, fdef: FuncDef) -> None
 
     class_ir.methods[func_ir.decl.name] = func_ir
 
+    # A decorated method that overrides a native method isn't a native method itself,
+    # so native calls to it go through a glue method that calls the decorated attribute
+    # using the Python API. The undecorated method stays in the type dict, since the
+    # decorators are applied to it at runtime.
+    decorated_override = (
+        is_decorated(builder, fdef)
+        and not fdef.is_property
+        and fdef not in builder.prop_setters
+        and name in class_ir.method_decls
+    )
+    if decorated_override:
+        glue = gen_glue_method(builder, func_ir.sig, func_ir, class_ir, class_ir, fdef.line, True)
+        # Native calls are compiled against the method's declaration, so the glue method
+        # takes it over from the undecorated method, which is renamed
+        glue = FuncIR(class_ir.method_decls[name], glue.arg_regs, glue.blocks, fdef.line)
+        builder.functions.append(glue)
+        class_ir.methods[name] = glue
+
     # If this overrides a parent class method with a different type, we need
     # to generate a glue method to mediate between them.
     for base in class_ir.mro[1:]:
@@ -549,7 +567,15 @@ def handle_ext_method(builder: IRBuilder, cdef: ClassDef, fdef: FuncDef) -> None
             # property setters. Need to make a special glue method for handling this,
             # similar to gen_glue_property.
 
-            f = gen_glue(builder, base.method_decls[name].sig, func_ir, class_ir, base, fdef)
+            f = gen_glue(
+                builder,
+                base.method_decls[name].sig,
+                func_ir,
+                class_ir,
+                base,
+                fdef,
+                do_py_ops=decorated_override,
+            )
             class_ir.glue_methods[(base, name)] = f
             builder.functions.append(f)
 
@@ -566,6 +592,12 @@ def handle_ext_method(builder: IRBuilder, cdef: ClassDef, fdef: FuncDef) -> None
         class_ir.glue_methods[(class_ir, func_ir.decl.name)] = f
         builder.functions.append(f)
 
+    # Rename the undecorated method only after all glue methods calling it by name exist
+    if decorated_override:
+        func_ir.decl.name = UNDECORATED_PREFIX + name
+        func_ir.decl.internal = False
+        class_ir.methods[func_ir.decl.name] = func_ir
+
     if fdef.name == "__getattr__":
         generate_getattr_wrapper(builder, cdef, fdef)
     elif fdef.name == "__setattr__":

mypyc/irbuild/prepare.py

@@ -40,6 +40,7 @@
 from mypy.semanal import refers_to_fullname
 from mypy.traverser import TraverserVisitor
 from mypy.types import Instance, Type, get_proper_type
+from mypy.util import is_dunder
 from mypyc.common import (
     FAST_PREFIX,
     GENERATOR_HELPER_NAME,
@@ -322,13 +323,41 @@ def prepare_method_def(
             # Making the argument implicitly positional-only avoids unnecessary glue methods
             decl.sig.args[1].pos_only = True
             ir.method_decls[PROPSET_PREFIX + node.name] = decl
+        elif (
+            decl.kind == FUNC_NORMAL
+            and not node.func.is_property
+            and not node.func.is_generator
+            and not node.func.is_coroutine
+            and overrides_native_method(cdef, node.name, mapper)
+        ):
+            # Native calls to a decorated override are dispatched through the vtable
+            # to a glue method that calls the decorated attribute. It's internal since
+            # the decorated attribute replaces it in the type dict.
+            decl.internal = True
+            ir.method_decls[node.name] = decl
 
         if node.func.is_property:
             assert node.func.type, f"Expected return type annotation for property '{node.name}'"
             decl.is_prop_getter = True
             ir.property_types[node.name] = decl.sig.ret_type
 
 
+def overrides_native_method(cdef: ClassDef, name: str, mapper: Mapper) -> bool:
+    # Dunders are excluded since they are called through type slots
+    if is_dunder(name):
+        return False
+    for base in cdef.info.mro[1:]:
+        sym = base.names.get(name)
+        if sym is not None:
+            base_ir = mapper.type_to_ir.get(base)
+            return (
+                base_ir is not None
+                and base_ir.is_ext_class
+                and isinstance(sym.node, (FuncDef, Decorator, OverloadedFuncDef))
+            )
+    return False
+
+
 def prepare_fast_path(
     ir: ClassIR,
     module_name: str,

mypyc/irbuild/vtable.py

@@ -2,8 +2,6 @@
 
 from __future__ import annotations
 
-import itertools
-
 from mypyc.ir.class_ir import ClassIR, VTableEntries, VTableMethod
 from mypyc.sametype import is_same_method_signature
 
@@ -38,14 +36,14 @@ def compute_vtable(cls: ClassIR) -> None:
     all_traits = [t for t in cls.mro if t.is_trait]
 
     for t in [cls] + [t for t in all_traits if t is not cls]:
-        for fn in itertools.chain(t.methods.values()):
+        for name, fn in t.methods.items():
             # TODO: don't generate a new entry when we overload without changing the type
-            if fn == cls.get_method(fn.name, prefer_method=True):
-                cls.vtable[fn.name] = len(entries)
+            if fn == cls.get_method(name, prefer_method=True):
+                cls.vtable[name] = len(entries)
                 # If the class contains a glue method referring to itself, that is a
                 # shadow glue method to support interpreted subclasses.
-                shadow = cls.glue_methods.get((cls, fn.name))
-                entries.append(VTableMethod(t, fn.name, fn, shadow))
+                shadow = cls.glue_methods.get((cls, name))
+                entries.append(VTableMethod(t, name, fn, shadow))
 
     # Compute vtables for all of the traits that the class implements
     if not cls.is_trait:

mypyc/test-data/run-classes.test

@@ -2385,6 +2385,114 @@ import other_interpreted
 
 [out]
 
+[case testDecoratedMethodOverridesNativeMethod]
+import functools
+from typing import Any, Callable, TypeVar
+from mypy_extensions import mypyc_attr, trait
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+def add_ten(f: F) -> F:
+    @functools.wraps(f)
+    def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+        return f(self, *args, **kwargs) + 10
+    return wrapper  # type: ignore[return-value]
+
+def tag(f: F) -> F:
+    def wrapper(self: Any) -> str:
+        return "<" + f(self) + ">"
+    return wrapper  # type: ignore[return-value]
+
+class Base:
+    def method(self, x: int) -> int:
+        return x
+
+    def call(self, x: int) -> int:
+        return self.method(x)
+
+class Derived(Base):
+    @add_ten
+    def method(self, x: int) -> int:
+        return super().method(x) * 2
+
+class Sub(Derived):
+    def method(self, x: int) -> int:
+        return super().method(x) + 100
+
+class ObjBase:
+    def m(self) -> object:
+        return 0
+
+class NarrowerReturn(ObjBase):
+    @add_ten
+    def m(self) -> int:
+        return 1
+
+@trait
+class T:
+    def m(self) -> int:
+        return 0
+
+class TraitImpl(T):
+    @add_ten
+    def m(self) -> int:
+        return 2
+
+class Repr:
+    @tag
+    def __repr__(self) -> str:
+        return "Repr"
+
+@mypyc_attr(allow_interpreted_subclasses=True)
+class OpenBase:
+    def m(self) -> int:
+        return 0
+
+@mypyc_attr(allow_interpreted_subclasses=True)
+class Open(OpenBase):
+    @add_ten
+    def m(self) -> int:
+        return 3
+
+def via_base(b: Base, x: int) -> int:
+    return b.method(x)
+
+def via_derived(d: Derived, x: int) -> int:
+    return d.method(x)
+
+def via_obj_base(b: ObjBase) -> object:
+    return b.m()
+
+def via_trait(t: T) -> int:
+    return t.m()
+
+def via_open_base(b: OpenBase) -> int:
+    return b.m()
+
+def test_decorated_override() -> None:
+    d = Derived()
+    assert via_base(d, 1) == 12
+    assert via_derived(d, 1) == 12
+    assert d.call(1) == 12
+    assert via_base(Base(), 1) == 1
+    assert via_base(Sub(), 1) == 112
+    assert via_obj_base(NarrowerReturn()) == 11
+    assert via_trait(TraitImpl()) == 12
+    assert repr(Repr()) == "<Repr>"
+
+[typing fixtures/typing-full.pyi]
+[file driver.py]
+from native import Derived, Open, via_open_base, test_decorated_override
+test_decorated_override()
+
+assert Derived().method(1) == 12
+assert Derived.method.__name__ == "method"
+
+class Inherits(Open):
+    pass
+
+assert via_open_base(Inherits()) == 13
+
 [case testPropertySetterWiderThanGetter]
 from typing import Union
 from mypy_extensions import trait

mypyc/test-data/run-multimodule.test

@@ -530,6 +530,38 @@ globals()['A'] = None
 [file driver.py]
 import other_main
 
+[case testMultiModuleDecoratedMethodOverride]
+from typing import Any, Callable, TypeVar
+from other_base import Base, call_method
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+def add_ten(f: F) -> F:
+    def wrapper(self: Any, x: int) -> int:
+        return f(self, x) + 10
+    return wrapper  # type: ignore[return-value]
+
+class Derived(Base):
+    @add_ten
+    def method(self, x: int) -> int:
+        return x * 2
+
+def test_decorated_override() -> None:
+    d = Derived()
+    assert d.method(1) == 12
+    assert call_method(d, 1) == 12
+    assert call_method(Base(), 1) == 1
+
+[file other_base.py]
+class Base:
+    def method(self, x: int) -> int:
+        return x
+
+def call_method(b: Base, x: int) -> int:
+    return b.method(x)
+
+[typing fixtures/typing-full.pyi]
+
 [case testNonNativeImportInPackageFile]
 # The import is really non-native only in separate compilation mode where __init__.py and
 # other_cache.py are in different libraries and the import uses the standard Python procedure.