[mypyc] Use defined __new__ method in tp_new and constructor

mypyc/analysis/attrdefined.py

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

mypyc/codegen/emitclass.py

@@ -5,6 +5,7 @@
 from collections.abc import Mapping
 from typing import Callable
 
+from mypy.nodes import ARG_STAR, ARG_STAR2
 from mypyc.codegen.cstring import c_string_initializer
 from mypyc.codegen.emit import Emitter, HeaderDeclaration, ReturnHandler
 from mypyc.codegen.emitfunc import native_function_doc_initializer, native_function_header
@@ -224,7 +225,7 @@ def generate_class(cl: ClassIR, module: str, emitter: Emitter) -> None:
     name = cl.name
     name_prefix = cl.name_prefix(emitter.names)
 
-    setup_name = f"{name_prefix}_setup"
+    setup_name = emitter.native_function_name(cl.setup)
     new_name = f"{name_prefix}_new"
     finalize_name = f"{name_prefix}_finalize"
     members_name = f"{name_prefix}_members"
@@ -317,10 +318,8 @@ def emit_line() -> None:
         fields["tp_basicsize"] = base_size
 
     if generate_full:
-        # Declare setup method that allocates and initializes an object. type is the
-        # type of the class being initialized, which could be another class if there
-        # is an interpreted subclass.
-        emitter.emit_line(f"static PyObject *{setup_name}(PyTypeObject *type);")
+        assert cl.setup is not None
+        emitter.emit_line(native_function_header(cl.setup, emitter) + ";")
         assert cl.ctor is not None
         emitter.emit_line(native_function_header(cl.ctor, emitter) + ";")
 
@@ -390,9 +389,7 @@ def emit_line() -> None:
 
     emitter.emit_line()
     if generate_full:
-        generate_setup_for_class(
-            cl, setup_name, defaults_fn, vtable_name, shadow_vtable_name, emitter
-        )
+        generate_setup_for_class(cl, defaults_fn, vtable_name, shadow_vtable_name, emitter)
         emitter.emit_line()
         generate_constructor_for_class(cl, cl.ctor, init_fn, setup_name, vtable_name, emitter)
         emitter.emit_line()
@@ -579,16 +576,16 @@ def generate_vtable(
 
 def generate_setup_for_class(
     cl: ClassIR,
-    func_name: str,
     defaults_fn: FuncIR | None,
     vtable_name: str,
     shadow_vtable_name: str | None,
     emitter: Emitter,
 ) -> None:
     """Generate a native function that allocates an instance of a class."""
-    emitter.emit_line("static PyObject *")
-    emitter.emit_line(f"{func_name}(PyTypeObject *type)")
+    emitter.emit_line(native_function_header(cl.setup, emitter))
     emitter.emit_line("{")
+    type_arg_name = REG_PREFIX + cl.setup.sig.args[0].name
+    emitter.emit_line(f"PyTypeObject *type = (PyTypeObject*){type_arg_name};")
     struct_name = cl.struct_name(emitter.names)
     emitter.emit_line(f"{struct_name} *self;")
 
@@ -663,6 +660,35 @@ def emit_attr_defaults_func_call(defaults_fn: FuncIR, self_name: str, emitter: E
     )
 
 
+def emit_setup_or_dunder_new_call(
+    cl: ClassIR,
+    setup_name: str,
+    type_arg: str,
+    native_prefix: bool,
+    new_args: str,
+    emitter: Emitter,
+) -> None:
+    def emit_null_check() -> None:
+        emitter.emit_line("if (self == NULL)")
+        emitter.emit_line("    return NULL;")
+
+    new_fn = cl.get_method("__new__")
+    if not new_fn:
+        emitter.emit_line(f"PyObject *self = {setup_name}({type_arg});")
+        emit_null_check()
+        return
+    prefix = emitter.get_group_prefix(new_fn.decl) + NATIVE_PREFIX if native_prefix else PREFIX
+    all_args = type_arg
+    if new_args != "":
+        all_args += ", " + new_args
+    emitter.emit_line(f"PyObject *self = {prefix}{new_fn.cname(emitter.names)}({all_args});")
+    emit_null_check()
+
+    # skip __init__ if __new__ returns some other type
+    emitter.emit_line(f"if (Py_TYPE(self) != {emitter.type_struct_name(cl)})")
+    emitter.emit_line("    return self;")
+
+
 def generate_constructor_for_class(
     cl: ClassIR,
     fn: FuncDecl,
@@ -674,17 +700,30 @@ def generate_constructor_for_class(
     """Generate a native function that allocates and initializes an instance of a class."""
     emitter.emit_line(f"{native_function_header(fn, emitter)}")
     emitter.emit_line("{")
-    emitter.emit_line(f"PyObject *self = {setup_name}({emitter.type_struct_name(cl)});")
-    emitter.emit_line("if (self == NULL)")
-    emitter.emit_line("    return NULL;")
-    args = ", ".join(["self"] + [REG_PREFIX + arg.name for arg in fn.sig.args])
+
+    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)
+
+    use_wrapper = (
+        cl.has_method("__new__")
+        and len(fn.sig.args) == 2
+        and fn.sig.args[0].kind == ARG_STAR
+        and fn.sig.args[1].kind == ARG_STAR2
+    )
+    emit_setup_or_dunder_new_call(cl, setup_name, type_arg, not use_wrapper, new_args, emitter)
+
+    args = ", ".join(["self"] + fn_args)
     if init_fn is not None:
+        prefix = PREFIX if use_wrapper else NATIVE_PREFIX
+        cast = "!= NULL ? 0 : -1" if use_wrapper else ""
         emitter.emit_line(
-            "char res = {}{}{}({});".format(
+            "char res = {}{}{}({}){};".format(
                 emitter.get_group_prefix(init_fn.decl),
-                NATIVE_PREFIX,
+                prefix,
                 init_fn.cname(emitter.names),
                 args,
+                cast,
             )
         )
         emitter.emit_line("if (res == 2) {")
@@ -717,7 +756,7 @@ 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:
+    if cl.allow_interpreted_subclasses or cl.builtin_base or cl.has_method("__new__"):
         emitter.emit_line(
             "return {}{}(self, args, kwds) != NULL ? 0 : -1;".format(
                 PREFIX, init_fn.cname(emitter.names)
@@ -750,15 +789,22 @@ def generate_new_for_class(
         emitter.emit_line("return NULL;")
         emitter.emit_line("}")
 
-    if not init_fn or cl.allow_interpreted_subclasses or cl.builtin_base or cl.is_serializable():
+    type_arg = "(PyObject*)type"
+    new_args = "args, kwds"
+    emit_setup_or_dunder_new_call(cl, setup_name, type_arg, False, new_args, emitter)
+    if (
+        not init_fn
+        or cl.allow_interpreted_subclasses
+        or cl.builtin_base
+        or cl.is_serializable()
+        or cl.has_method("__new__")
+    ):
         # Match Python semantics -- __new__ doesn't call __init__.
-        emitter.emit_line(f"return {setup_name}(type);")
+        emitter.emit_line("return self;")
     else:
         # __new__ of a native class implicitly calls __init__ so that we
         # can enforce that instances are always properly initialized. This
         # is needed to support always defined attributes.
-        emitter.emit_line(f"PyObject *self = {setup_name}(type);")
-        emitter.emit_lines("if (self == NULL)", "    return NULL;")
         emitter.emit_line(
             f"PyObject *ret = {PREFIX}{init_fn.cname(emitter.names)}(self, args, kwds);"
         )

mypyc/codegen/emitmodule.py

@@ -1274,8 +1274,8 @@ def is_fastcall_supported(fn: FuncIR, capi_version: tuple[int, int]) -> bool:
         if fn.name == "__call__":
             # We can use vectorcalls (PEP 590) when supported
             return True
-        # TODO: Support fastcall for __init__.
-        return fn.name != "__init__"
+        # TODO: Support fastcall for __init__ and __new__.
+        return fn.name != "__init__" and fn.name != "__new__"
     return True
 
 

mypyc/codegen/emitwrapper.py

@@ -238,7 +238,7 @@ def generate_legacy_wrapper_function(
     real_args = list(fn.args)
     if fn.sig.num_bitmap_args:
         real_args = real_args[: -fn.sig.num_bitmap_args]
-    if fn.class_name and fn.decl.kind != FUNC_STATICMETHOD:
+    if fn.class_name and (fn.decl.name == "__new__" or fn.decl.kind != FUNC_STATICMETHOD):
         arg = real_args.pop(0)
         emitter.emit_line(f"PyObject *obj_{arg.name} = self;")
 

mypyc/ir/class_ir.py

@@ -5,9 +5,9 @@
 from typing import NamedTuple
 
 from mypyc.common import PROPSET_PREFIX, JsonDict
-from mypyc.ir.func_ir import FuncDecl, FuncIR, FuncSignature
+from mypyc.ir.func_ir import FuncDecl, FuncIR, FuncSignature, RuntimeArg
 from mypyc.ir.ops import DeserMaps, Value
-from mypyc.ir.rtypes import RInstance, RType, deserialize_type
+from mypyc.ir.rtypes import RInstance, RType, deserialize_type, object_rprimitive
 from mypyc.namegen import NameGenerator, exported_name
 
 # Some notes on the vtable layout: Each concrete class has a vtable
@@ -133,6 +133,16 @@ def __init__(
         self.builtin_base: str | None = None
         # Default empty constructor
         self.ctor = FuncDecl(name, None, module_name, FuncSignature([], RInstance(self)))
+        # Declare setup method that allocates and initializes an object. type is the
+        # type of the class being initialized, which could be another class if there
+        # is an interpreted subclass.
+        # TODO: Make it a regular method and generate its body in IR
+        self.setup = FuncDecl(
+            "__mypyc__" + name + "_setup",
+            None,
+            module_name,
+            FuncSignature([RuntimeArg("type", object_rprimitive)], RInstance(self)),
+        )
         # Attributes defined in the class (not inherited)
         self.attributes: dict[str, RType] = {}
         # Deletable attributes

mypyc/irbuild/expression.py

@@ -57,6 +57,7 @@
 from mypyc.ir.ops import (
     Assign,
     BasicBlock,
+    Call,
     ComparisonOp,
     Integer,
     LoadAddress,
@@ -472,23 +473,42 @@ def translate_super_method_call(builder: IRBuilder, expr: CallExpr, callee: Supe
         if callee.name in base.method_decls:
             break
     else:
-        if (
-            ir.is_ext_class
-            and ir.builtin_base is None
-            and not ir.inherits_python
-            and callee.name == "__init__"
-            and len(expr.args) == 0
-        ):
-            # Call translates to object.__init__(self), which is a
-            # no-op, so omit the call.
-            return builder.none()
+        if ir.is_ext_class and ir.builtin_base is None and not ir.inherits_python:
+            if callee.name == "__init__" and len(expr.args) == 0:
+                # Call translates to object.__init__(self), which is a
+                # no-op, so omit the call.
+                return builder.none()
+            elif callee.name == "__new__":
+                # object.__new__(cls)
+                assert (
+                    len(expr.args) == 1
+                ), f"Expected object.__new__() call to have exactly 1 argument, got {len(expr.args)}"
+                typ_arg = expr.args[0]
+                method_args = builder.fn_info.fitem.arg_names
+                if (
+                    isinstance(typ_arg, NameExpr)
+                    and len(method_args) > 0
+                    and method_args[0] == typ_arg.name
+                ):
+                    subtype = builder.accept(expr.args[0])
+                    return builder.add(Call(ir.setup, [subtype], expr.line))
+
+        if callee.name == "__new__":
+            call = "super().__new__()"
+            if not ir.is_ext_class:
+                builder.error(f"{call} not supported for non-extension classes", expr.line)
+            if ir.inherits_python:
+                builder.error(
+                    f"{call} not supported for classes inheriting from non-native classes",
+                    expr.line,
+                )
         return translate_call(builder, expr, callee)
 
     decl = base.method_decl(callee.name)
     arg_values = [builder.accept(arg) for arg in expr.args]
     arg_kinds, arg_names = expr.arg_kinds.copy(), expr.arg_names.copy()
 
-    if decl.kind != FUNC_STATICMETHOD:
+    if decl.kind != FUNC_STATICMETHOD and decl.name != "__new__":
         # Grab first argument
         vself: Value = builder.self()
         if decl.kind == FUNC_CLASSMETHOD:

mypyc/irbuild/prepare.py

@@ -193,9 +193,9 @@ def prepare_func_def(
     create_generator_class_if_needed(module_name, class_name, fdef, mapper)
 
     kind = (
-        FUNC_STATICMETHOD
-        if fdef.is_static
-        else (FUNC_CLASSMETHOD if fdef.is_class else FUNC_NORMAL)
+        FUNC_CLASSMETHOD
+        if fdef.is_class
+        else (FUNC_STATICMETHOD if fdef.is_static else FUNC_NORMAL)
     )
     sig = mapper.fdef_to_sig(fdef, options.strict_dunders_typing)
     decl = FuncDecl(fdef.name, class_name, module_name, sig, kind)
@@ -555,21 +555,57 @@ def add_setter_declaration(
     ir.method_decls[setter_name] = decl
 
 
+def check_matching_args(init_sig: FuncSignature, new_sig: FuncSignature) -> bool:
+    num_init_args = len(init_sig.args) - init_sig.num_bitmap_args
+    num_new_args = len(new_sig.args) - new_sig.num_bitmap_args
+    if num_init_args != num_new_args:
+        return False
+
+    for idx in range(1, num_init_args):
+        init_arg = init_sig.args[idx]
+        new_arg = new_sig.args[idx]
+        if init_arg.type != new_arg.type:
+            return False
+
+        if init_arg.kind != new_arg.kind:
+            return False
+
+    return True
+
+
 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
+
+    new_node: SymbolNode | None = None
+    new_symbol = cdef.info.get("__new__")
+    # We are only interested in __new__ method defined in a user-defined class,
+    # so we ignore it if it comes from a builtin type. It's usually builtins.object
+    # but could also be builtins.type for metaclasses so we detect the prefix which
+    # matches both.
+    if new_symbol and new_symbol.fullname and not new_symbol.fullname.startswith("builtins."):
+        new_node = new_symbol.node
+    if isinstance(new_node, (Decorator, OverloadedFuncDef)):
+        new_node = get_func_def(new_node)
     if not ir.is_trait and not ir.builtin_base and isinstance(init_node, FuncDef):
         init_sig = mapper.fdef_to_sig(init_node, True)
+        args_match = True
+        if isinstance(new_node, FuncDef):
+            new_sig = mapper.fdef_to_sig(new_node, True)
+            args_match = check_matching_args(init_sig, new_sig)
 
         defining_ir = mapper.type_to_ir.get(init_node.info)
         # If there is a nontrivial __init__ that wasn't defined in an
         # extension class, we need to make the constructor take *args,
         # **kwargs so it can call tp_init.
         if (
-            defining_ir is None
-            or not defining_ir.is_ext_class
-            or cdef.info["__init__"].plugin_generated
-        ) and init_node.info.fullname != "builtins.object":
+            (
+                defining_ir is None
+                or not defining_ir.is_ext_class
+                or cdef.info["__init__"].plugin_generated
+            )
+            and init_node.info.fullname != "builtins.object"
+        ) or not args_match:
             init_sig = FuncSignature(
                 [
                     init_sig.args[0],

mypyc/lib-rt/misc_ops.c

@@ -227,6 +227,17 @@ PyObject *CPyType_FromTemplate(PyObject *template,
     if (!name)
         goto error;
 
+    if (template_->tp_doc) {
+        // cpython expects tp_doc to be heap-allocated so convert it here to
+        // avoid segfaults on deallocation.
+        Py_ssize_t size = strlen(template_->tp_doc) + 1;
+        char *doc = (char *)PyMem_Malloc(size);
+        if (!doc)
+            goto error;
+        memcpy(doc, template_->tp_doc, size);
+        template_->tp_doc = doc;
+    }
+
     // Allocate the type and then copy the main stuff in.
     t = (PyHeapTypeObject*)PyType_GenericAlloc(&PyType_Type, 0);
     if (!t)