[mypyc] Fix class decorators that don't return a class

mypyc/codegen/emitmodule.py

@@ -23,7 +23,7 @@
 )
 from mypy.errors import CompileError
 from mypy.fscache import FileSystemCache
-from mypy.nodes import MypyFile, TypeInfo
+from mypy.nodes import ClassDef, MypyFile, TypeInfo
 from mypy.options import Options
 from mypy.plugin import Plugin, ReportConfigContext
 from mypy.util import hash_digest, json_dumps
@@ -82,6 +82,7 @@
 from mypyc.irbuild.main import build_ir
 from mypyc.irbuild.mapper import Mapper
 from mypyc.irbuild.prepare import load_type_map
+from mypyc.irbuild.util import is_decorated_class
 from mypyc.namegen import NameGenerator, exported_name
 from mypyc.options import CompilerOptions
 from mypyc.transform.borrow_generator_attrs import borrow_generator_attrs
@@ -139,6 +140,10 @@ class MypycPlugin(Plugin):
       * If the IR metadata is missing or stale or any of the generated
         C source files associated missing or stale, then we need to
         recompile the module so we mark it as stale.
+      * Mypy ignores class decorators, but a decorator that can replace
+        a class changes how other modules refer to the class (see
+        ClassIR.is_decorated), so we make such decorated classes part
+        of the module's interface.
     """
 
     def __init__(
@@ -154,8 +159,11 @@ def __init__(
         self.compiler_options = compiler_options
         self.metastore = create_metastore(options, parallel_worker=False)
 
-    def report_config_data(self, ctx: ReportConfigContext) -> tuple[str | None, list[str]] | None:
-        # The config data we report is the group map entry for the module.
+    def report_config_data(
+        self, ctx: ReportConfigContext
+    ) -> tuple[str | None, list[str], list[str]] | None:
+        # The config data we report is the group map entry for the module,
+        # followed by the names of its decorated classes.
         # If the data is being used to check validity, we do additional checks
         # that the IR cache exists and matches the metadata cache and all
         # output source files exist and are up to date.
@@ -167,7 +175,8 @@ def report_config_data(self, ctx: ReportConfigContext) -> tuple[str | None, list
 
         # If we aren't doing validity checks, just return the cache data
         if not is_check:
-            return self.group_map[id]
+            assert self._modules is not None
+            return (*self.group_map[id], decorated_class_names(self._modules[id]))
 
         # Load the metadata and IR cache
         meta_path, _, _ = get_cache_names(id, path, self.options)
@@ -201,7 +210,10 @@ def report_config_data(self, ctx: ReportConfigContext) -> tuple[str | None, list
             if hash != real_hash:
                 return None
 
-        return self.group_map[id]
+        # The module hasn't been parsed yet, so take the decorated classes
+        # from the cached IR. It was built from the source the cache is for.
+        decorated = sorted(c["name"] for c in ir_data["ir"]["classes"] if c["is_decorated"])
+        return (*self.group_map[id], decorated)
 
     def get_additional_deps(self, file: MypyFile) -> list[tuple[int, str, int]]:
         # Report dependency on modules in the module's group
@@ -230,6 +242,13 @@ def get_additional_indirect_deps(self, file: MypyFile) -> set[str]:
         return mods
 
 
+def decorated_class_names(tree: MypyFile) -> list[str]:
+    """Names of the classes that build_type_map sets ClassIR.is_decorated for."""
+    return sorted(
+        node.name for node in tree.defs if isinstance(node, ClassDef) and is_decorated_class(node)
+    )
+
+
 def parse_and_typecheck(
     sources: list[BuildSource],
     options: Options,

mypyc/ir/class_ir.py

@@ -105,6 +105,10 @@ def __init__(
         # An augmented class has additional methods separate from what mypyc generates.
         # Right now the only one is dataclasses.
         self.is_augmented = False
+        # Does this (non-extension) class have class decorators other than the ones native
+        # classes support? These can replace the class with an arbitrary object, such as a
+        # functools.cache wrapper, so the name of the class may not refer to the class.
+        self.is_decorated = False
         # Does this inherit from a Python class?
         self.inherits_python = False
         # Do instances of this class have __dict__?
@@ -480,6 +484,7 @@ def serialize(self) -> JsonDict:
             "is_abstract": self.is_abstract,
             "is_generated": self.is_generated,
             "is_augmented": self.is_augmented,
+            "is_decorated": self.is_decorated,
             "is_final_class": self.is_final_class,
             "inherits_python": self.inherits_python,
             "has_dict": self.has_dict,
@@ -550,6 +555,7 @@ def deserialize(cls, data: JsonDict, ctx: DeserMaps) -> ClassIR:
         ir.is_abstract = data["is_abstract"]
         ir.is_ext_class = data["is_ext_class"]
         ir.is_augmented = data["is_augmented"]
+        ir.is_decorated = data["is_decorated"]
         ir.is_final_class = data["is_final_class"]
         ir.inherits_python = data["inherits_python"]
         ir.has_dict = data["has_dict"]

mypyc/irbuild/builder.py

@@ -1250,6 +1250,14 @@ def is_synthetic_type(self, typ: TypeInfo) -> bool:
         """Is a type something other than just a class we've created?"""
         return typ.is_named_tuple or typ.is_newtype or typ.typeddict_type is not None
 
+    def is_decorated_class(self, typ: TypeInfo) -> bool:
+        """Might a class decorator have replaced the class with another object?
+
+        The name of the class then refers to the decorated object, not to the class.
+        """
+        ir = self.mapper.type_to_ir.get(typ)
+        return ir is not None and ir.is_decorated
+
     def get_final_ref(self, expr: MemberExpr) -> tuple[str, Var, bool] | None:
         """Check if `expr` is a final class, module or instance attribute.
 
@@ -1755,6 +1763,7 @@ def load_global(self, expr: NameExpr) -> Value:
             self.is_native_module_ref_expr(expr)
             and isinstance(expr.node, TypeInfo)
             and not self.is_synthetic_type(expr.node)
+            and not self.is_decorated_class(expr.node)
         ):
             assert expr.fullname
             return self.load_native_type_object(expr.fullname)

mypyc/irbuild/classdef.py

@@ -76,6 +76,7 @@
 from mypyc.primitives.misc_ops import (
     dataclass_sleight_of_hand,
     import_op,
+    isinstance_type,
     not_implemented_op,
     py_calc_meta_op,
     py_init_subclass_op,
@@ -287,23 +288,30 @@ def add_attr(self, lvalue: NameExpr, stmt: AssignmentStmt) -> None:
     def finalize(self, ir: ClassIR) -> None:
         # Dynamically create the class via the type constructor
         non_ext_class = load_non_ext_class(self.builder, ir, self.non_ext, self.cdef.line)
-        non_ext_class = load_decorated_class(self.builder, self.cdef, non_ext_class)
+        decorated_class = load_decorated_class(self.builder, self.cdef, non_ext_class)
+        type_obj = decorated_class
+        if ir.is_decorated:
+            type_obj = load_decorated_class_type(
+                self.builder, non_ext_class, decorated_class, self.cdef.line
+            )
 
         # Try to avoid contention when using free threading.
-        self.builder.set_immortal_if_free_threaded(non_ext_class, self.cdef.line)
+        self.builder.set_immortal_if_free_threaded(type_obj, self.cdef.line)
+        if type_obj is not decorated_class:
+            self.builder.set_immortal_if_free_threaded(decorated_class, self.cdef.line)
 
-        # Save the decorated class
+        # Save the class object
         self.builder.add(
-            InitStatic(non_ext_class, self.cdef.name, self.builder.module_name, NAMESPACE_TYPE)
+            InitStatic(type_obj, self.cdef.name, self.builder.module_name, NAMESPACE_TYPE)
         )
 
-        # Add the non-extension class to the dict
+        # Add the decorated class to the dict
         self.builder.call_c(
             exact_dict_set_item_op,
             [
                 self.builder.load_globals_dict(),
                 self.builder.load_str(self.cdef.name),
-                non_ext_class,
+                decorated_class,
             ],
             self.cdef.line,
         )
@@ -933,6 +941,27 @@ def load_decorated_class(builder: IRBuilder, cdef: ClassDef, type_obj: Value) ->
     return dec_class
 
 
+def load_decorated_class_type(
+    builder: IRBuilder, undecorated: Value, decorated: Value, line: int
+) -> Value:
+    """Find the class object to use in type checks of a decorated non-extension class.
+
+    A class decorator can replace the class with an object that isn't a class, such as
+    a functools.cache wrapper. Instances still have the undecorated class as their type
+    then, so use it. If the decorator returned a class (perhaps a new one, as
+    six.add_metaclass does), use that.
+    """
+    result = Register(object_rprimitive)
+    builder.assign(result, decorated, line)
+    is_class = builder.primitive_op(isinstance_type, [decorated], line)
+    not_class, done = BasicBlock(), BasicBlock()
+    builder.add(Branch(is_class, done, not_class, Branch.BOOL))
+    builder.activate_block(not_class)
+    builder.assign(result, undecorated, line)
+    builder.goto_and_activate(done)
+    return result
+
+
 def cache_class_attrs(
     builder: IRBuilder, attrs_to_cache: list[tuple[Lvalue, RType]], cdef: ClassDef
 ) -> None:

mypyc/irbuild/prepare.py

@@ -78,6 +78,7 @@
     get_mypyc_attrs,
     is_class_body_final,
     is_dataclass,
+    is_decorated_class,
     is_extension_class,
     is_trait,
 )
@@ -110,6 +111,7 @@ def build_type_map(
             is_final_class=cdef.info.is_final,
         )
         class_ir.is_ext_class = is_extension_class(module.path, cdef, errors)
+        class_ir.is_decorated = is_decorated_class(cdef)
         if class_ir.is_ext_class:
             class_ir.deletable = cdef.info.deletable_attributes.copy()
         # If global optimizations are disabled, turn of tracking of class children

mypyc/irbuild/util.py

@@ -304,21 +304,30 @@ def get_explicit_native_class(path: str, cdef: ClassDef, errors: Errors) -> bool
     return None
 
 
+def is_decorated_class(cdef: ClassDef) -> bool:
+    """Does a class have decorators other than the ones native classes support?
+
+    Only non-extension classes can have them. Unlike the supported decorators, which
+    return a class, these can replace the class with an arbitrary object.
+    """
+    return any(
+        not is_trait_decorator(d)
+        and not is_dataclass_decorator(d)
+        and not get_mypyc_attr_call(d)
+        and not is_final_decorator(d)
+        for d in cdef.decorators
+    )
+
+
 def is_implicit_extension_class(cdef: ClassDef) -> tuple[bool, str]:
     """Check if class can be extension class and return a user-friendly reason it can't be one."""
 
-    for d in cdef.decorators:
-        if (
-            not is_trait_decorator(d)
-            and not is_dataclass_decorator(d)
-            and not get_mypyc_attr_call(d)
-            and not is_final_decorator(d)
-        ):
-            return (
-                False,
-                "Classes that have decorators other than supported decorators"
-                " can't be native classes.",
-            )
+    if is_decorated_class(cdef):
+        return (
+            False,
+            "Classes that have decorators other than supported decorators"
+            " can't be native classes.",
+        )
 
     if cdef.info.typeddict_type:
         return False, "TypedDict classes can't be native classes."

mypyc/test-data/irbuild-classes.test

@@ -1498,6 +1498,51 @@ class M(type):  # E: Inheriting from most builtin types is unimplemented \
 class A(metaclass=M):  # E: Class is marked as native_class=True but it can't be a native class. Classes with a metaclass other than ABCMeta, TypingMeta or GenericMeta can't be native classes.
     pass
 
+[case testDecoratedNonExtClassReference]
+from typing import Any
+
+def decorator(cls: Any) -> Any:
+    return cls
+
+@decorator
+class C:
+    pass
+
+def f() -> C:
+    return C()
+
+def g(x: object) -> bool:
+    return isinstance(x, C)
+[out]
+def decorator(cls):
+    cls :: object
+L0:
+    return cls
+def f():
+    r0 :: dict
+    r1 :: str
+    r2, r3 :: object
+    r4 :: __main__.C
+L0:
+    r0 = __main__.globals :: static
+    r1 = 'C'
+    r2 = CPyDict_GetItem(r0, r1)
+    r3 = PyObject_Vectorcall(r2, 0, 0, 0)
+    r4 = cast(__main__.C, r3)
+    return r4
+def g(x):
+    x, r0 :: object
+    r1 :: ptr
+    r2 :: object
+    r3 :: bit
+L0:
+    r0 = __main__.C :: type
+    r1 = get_element_ptr x ob_type :: PyObject
+    r2 = borrow load_mem r1 :: builtins.object*
+    keep_alive x
+    r3 = r2 == r0
+    return r3
+
 [case testReservedName]
 from typing import Any, overload