[mypyc] Fix decorated __new__ in native classes

mypyc/analysis/attrdefined.py

@@ -139,6 +139,7 @@ def analyze_always_defined_attrs_in_class(cl: ClassIR, seen: set[ClassIR]) -> No
         or cl.children is None
         or cl.is_serializable()
         or cl.has_method("__new__")
+        or cl.has_non_native_new()
     ):
         # Give up -- we can't enforce that attributes are always defined.
         return

mypyc/codegen/emitclass.py

@@ -314,6 +314,17 @@ def emit_line() -> None:
     # Fill out slots in the type object from dunder methods.
     fields.update(generate_slots(cl, SLOT_DEFS, emitter))
 
+    if cl.has_non_native_new():
+        # Leave tp_new unset. When a decorated __new__ is set as an attribute of the
+        # class that defines it, Python sets tp_new to a function that calls __new__
+        # from the type dict, and subclasses inherit that. A tp_new of our own would
+        # add a default __new__ to the type dict, which would hide the decorated one.
+        fields.pop("tp_new", None)
+        if init_fn is not None and "tp_init" not in fields:
+            # An inherited tp_init may do nothing, since tp_new of a class without
+            # __new__ calls __init__.
+            fields["tp_init"] = generate_init_for_class(cl, init_fn, emitter)
+
     # Fill out dunder methods that live in tables hanging off the side.
     for table_name, type, slot_defs in SIDE_TABLES:
         slots = generate_slots(cl, slot_defs, emitter)
@@ -676,6 +687,10 @@ def generate_setup_for_class(
     struct_name = cl.struct_name(emitter.names)
     emitter.emit_line(f"{struct_name} *self;")
 
+    if cl.has_non_native_new() and not cl.allow_interpreted_subclasses:
+        # tp_new normally checks this, but this class doesn't have its own.
+        emit_interpreted_subclass_check(cl, emitter)
+
     prefix = cl.name_prefix(emitter.names)
     if cl.reuse_freed_instance:
         # Attempt to use a per-type free list first (a free "list" with up to one object only).
@@ -781,6 +796,31 @@ def emit_null_check() -> None:
     emitter.emit_line("    return self;")
 
 
+def emit_interpreted_subclass_check(cl: ClassIR, emitter: Emitter) -> None:
+    """Emit C code to raise an exception unless the C variable "type" is the type object of cl."""
+    emitter.emit_line(f"if (type != {emitter.type_struct_name(cl)}) {{")
+    emitter.emit_line(
+        'PyErr_SetString(PyExc_TypeError, "interpreted classes cannot inherit from compiled");'
+    )
+    emitter.emit_line("return NULL;")
+    emitter.emit_line("}")
+
+
+def emit_type_object_call(cl: ClassIR, fn: FuncDecl, emitter: Emitter) -> None:
+    """Emit C code to return the result of calling the type object of cl.
+
+    The arguments are the arguments of fn, which is the constructor of cl.
+    """
+    type_arg = "(PyObject *)" + emitter.type_struct_name(cl)
+    if fn.sig.args:
+        # The constructor takes *args, **kwargs.
+        assert [arg.kind for arg in fn.sig.args] == [ARG_STAR, ARG_STAR2], fn.sig
+        args = ", ".join(REG_PREFIX + arg.name for arg in fn.sig.args)
+        emitter.emit_line(f"return PyObject_Call({type_arg}, {args});")
+    else:
+        emitter.emit_line(f"return PyObject_CallNoArgs({type_arg});")
+
+
 def generate_constructor_for_class(
     cl: ClassIR,
     fn: FuncDecl,
@@ -793,6 +833,12 @@ def generate_constructor_for_class(
     emitter.emit_line(f"{native_function_header(fn, emitter)}")
     emitter.emit_line("{")
 
+    if cl.has_non_native_new():
+        # A decorated __new__ isn't called natively, so call the type object.
+        emit_type_object_call(cl, fn, emitter)
+        emitter.emit_line("}")
+        return
+
     fn_args = [REG_PREFIX + arg.name for arg in fn.sig.args]
     type_arg = "(PyObject *)" + emitter.type_struct_name(cl)
     new_args = ", ".join(fn_args)
@@ -844,7 +890,12 @@ def generate_init_for_class(cl: ClassIR, init_fn: FuncIR, emitter: Emitter) -> s
     emitter.emit_line("static int")
     emitter.emit_line(f"{func_name}(PyObject *self, PyObject *args, PyObject *kwds)")
     emitter.emit_line("{")
-    if cl.allow_interpreted_subclasses or cl.builtin_base or cl.has_method("__new__"):
+    if (
+        cl.allow_interpreted_subclasses
+        or cl.builtin_base
+        or cl.has_method("__new__")
+        or cl.has_non_native_new()
+    ):
         emitter.emit_line(
             f"return {emitter.wrapper_function_call(init_fn.decl)}"
             "(self, args, kwds) != NULL ? 0 : -1;"
@@ -864,17 +915,15 @@ def generate_new_for_class(
     init_fn: FuncIR | None,
     emitter: Emitter,
 ) -> None:
+    if cl.has_non_native_new():
+        # The class doesn't have its own tp_new.
+        return
     emitter.emit_line("static PyObject *")
     emitter.emit_line(f"{func_name}(PyTypeObject *type, PyObject *args, PyObject *kwds)")
     emitter.emit_line("{")
     # TODO: Check and unbox arguments
     if not cl.allow_interpreted_subclasses:
-        emitter.emit_line(f"if (type != {emitter.type_struct_name(cl)}) {{")
-        emitter.emit_line(
-            'PyErr_SetString(PyExc_TypeError, "interpreted classes cannot inherit from compiled");'
-        )
-        emitter.emit_line("return NULL;")
-        emitter.emit_line("}")
+        emit_interpreted_subclass_check(cl, emitter)
 
     type_arg = "(PyObject*)type"
     new_args = "args, kwds"
@@ -1117,7 +1166,13 @@ def generate_methods_table(
         emitter.emit_line(f'{{"{fn.name}",')
         emitter.emit_line(f" (PyCFunction){PREFIX}{fn.cname(emitter.names)},")
         flags = ["METH_FASTCALL", "METH_KEYWORDS"]
-        if fn.decl.kind == FUNC_STATICMETHOD:
+        if fn.name == "__new__" and "__new__" not in cl.method_decls:
+            # A decorated __new__. Its decorators get the callable that is in the type
+            # dict. The wrapper takes the class as self, like a class method, and a
+            # class method descriptor that is called with a class as the first argument
+            # binds itself to that class.
+            flags = ["METH_VARARGS", "METH_KEYWORDS", "METH_CLASS"]
+        elif fn.decl.kind == FUNC_STATICMETHOD:
             flags.append("METH_STATIC")
         elif fn.decl.kind == FUNC_CLASSMETHOD:
             flags.append("METH_CLASS")

mypyc/ir/class_ir.py

@@ -438,6 +438,17 @@ def get_method(self, name: str, *, prefer_method: bool = False) -> FuncIR | None
     def has_method_decl(self, name: str) -> bool:
         return any(name in ir.method_decls for ir in self.mro)
 
+    def has_non_native_new(self) -> bool:
+        """Are instances created by a __new__ method that isn't called natively?
+
+        A decorated __new__ is one. Like any decorated method it has no method
+        declaration, since only the decorated callable in the type dict can be called.
+        Instances are then created like instances of a Python class: tp_new looks up
+        __new__ in the type dict, and tp_init calls __init__.
+        """
+        res = self.get_method_and_class("__new__")
+        return res is not None and "__new__" not in res[1].method_decls
+
     def has_no_subclasses(self) -> bool:
         return self.children == [] and not self.allow_interpreted_subclasses
 

mypyc/irbuild/function.py

@@ -489,6 +489,10 @@ def handle_ext_method(builder: IRBuilder, cdef: ClassDef, fdef: FuncDef) -> None
     builder.functions.append(func_ir)
 
     if is_decorated(builder, fdef):
+        if name == "__new__" and class_ir.is_trait:
+            # A decorated __new__ is called by an inherited tp_new, and a class only
+            # inherits tp_new from its first base class.
+            builder.error('Decorated "__new__" not supported in traits', fdef.line)
         # Obtain the function name in order to construct the name of the helper function.
         _, _, name = fdef.fullname.rpartition(".")
         # Get the callable representing the non-decorated method directly from the type

mypyc/irbuild/prepare.py

@@ -772,6 +772,28 @@ def check_matching_args(init_sig: FuncSignature, new_sig: FuncSignature) -> bool
     return True
 
 
+def has_non_native_new(cdef: ClassDef, mapper: Mapper) -> bool:
+    """Are instances of a class created by a __new__ method that isn't called natively?
+
+    This is ClassIR.has_non_native_new() for use before the methods have been built.
+    """
+    for info in cdef.info.mro:
+        class_ir = mapper.type_to_ir.get(info)
+        symbol = info.names.get("__new__")
+        if class_ir is None or not class_ir.is_ext_class or symbol is None:
+            continue
+        node = symbol.node
+        if isinstance(node, OverloadedFuncDef):
+            node = node.impl
+        if isinstance(node, Decorator):
+            # A decorated method has no method declaration. Other methods only lack one
+            # while their class isn't prepared yet, which can happen in an import cycle.
+            return "__new__" not in class_ir.method_decls
+        if isinstance(node, FuncDef):
+            return False
+    return False
+
+
 def prepare_init_method(cdef: ClassDef, ir: ClassIR, module_name: str, mapper: Mapper) -> None:
     # Set up a constructor decl
     init_node = cdef.info["__init__"].node
@@ -792,6 +814,9 @@ def prepare_init_method(cdef: ClassDef, ir: ClassIR, module_name: str, mapper: M
         if isinstance(new_node, FuncDef):
             new_sig = mapper.fdef_to_sig(new_node, True)
             args_match = check_matching_args(init_sig, new_sig)
+        if has_non_native_new(cdef, mapper):
+            # The constructor calls the type object with *args, **kwargs.
+            args_match = False
 
         defining_ir = mapper.type_to_ir.get(init_node.info)
         # If there is a nontrivial __init__ that wasn't defined in an

mypyc/test-data/irbuild-classes.test

@@ -2212,9 +2212,12 @@ L0:
 
 [case testUnsupportedDunderNew]
 from __future__ import annotations
-from mypy_extensions import mypyc_attr
+from mypy_extensions import mypyc_attr, trait
 from typing import Any
 
+def decorator(f: Any) -> Any:
+    return f
+
 @mypyc_attr(native_class=False)
 class NonNative:
     def __new__(cls) -> NonNative:
@@ -2261,6 +2264,95 @@ class WrongNumberOfArgsObjectNew:
     def __new__(cls):
         return object.__new__(cls, 1)  # E: "object.__new__()" supported only with 1 argument, got 2
 
+@trait
+class DecoratedInTrait:
+    @decorator
+    def __new__(cls) -> Any:  # E: Decorated "__new__" not supported in traits
+        return super().__new__(cls)
+
+[case testDecoratedDunderNew]
+from typing import Any, Callable, TypeVar
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+def deco(f: F) -> F:
+    return f
+
+class A:
+    @deco
+    def __new__(cls, x: int) -> Any:
+        return super().__new__(cls)
+
+    def __init__(self, x: int) -> None:
+        pass
+
+class B(A):
+    pass
+
+class C:
+    def __new__(cls, x: int) -> Any:
+        return super().__new__(cls)
+
+    def __init__(self, x: int) -> None:
+        pass
+
+def f() -> None:
+    A(1)
+    B(2)
+    C(3)
+[out]
+def deco(f):
+    f :: object
+L0:
+    return f
+def A.__new__(cls, x):
+    cls :: object
+    x :: int
+    r0 :: object
+    r1 :: __main__.A
+L0:
+    r0 = CPy_SetupObject(cls)
+    r1 = cast(__main__.A, r0)
+    return r1
+def A.__init__(self, x):
+    self :: __main__.A
+    x :: int
+L0:
+    return 1
+def C.__new__(cls, x):
+    cls :: object
+    x :: int
+    r0 :: __main__.C
+L0:
+    r0 = __mypyc__C_setup(cls)
+    return r0
+def C.__init__(self, x):
+    self :: __main__.C
+    x :: int
+L0:
+    return 1
+def f():
+    r0 :: object
+    r1 :: tuple
+    r2 :: dict
+    r3 :: __main__.A
+    r4 :: object
+    r5 :: tuple
+    r6 :: dict
+    r7 :: __main__.B
+    r8 :: __main__.C
+L0:
+    r0 = object 1
+    r1 = PyTuple_Pack(1, r0)
+    r2 = PyDict_New()
+    r3 = A(r1, r2)
+    r4 = object 2
+    r5 = PyTuple_Pack(1, r4)
+    r6 = PyDict_New()
+    r7 = B(r5, r6)
+    r8 = C(6)
+    return 1
+
 [case testClassWithFreeList]
 from mypy_extensions import mypyc_attr, trait