[mypyc] Use defined __new__ method in tp_new and constructor (#19739)

Fixes #16012 

mypyc ignored custom implementations of `__new__` because, even though a
C function representing the method was generated, it was called neither
in the type constructor nor in the method assigned to the `tp_new`
pointer.

Now if there's a `__new__` method defined for a type, the corresponding
function is called in place of the setup function which is responsible
for allocating memory for new objects and initializing their attributes
to default values.

The setup function is still called when creating instances of the type
as calls resolving to `object.__new__()` are transformed to call the
setup function. This way, `__new__` can return instances of other types
and instances of the type of the class where `__new__` is defined are
setup correctly.

There are a couple of limitations:
- Programs with `super().__new__()` calls in `__new__` methods of
non-native classes are rejected because it's more difficult to resolve
the setup function for non-native classes but this could probably be
supported in the future.
- Similarly, programs are rejected when a class inherits from a
non-compiled class. In this case calling the `tp_new` method of the
parent type results in an error because cpython expects the sub type to
use a wrapper for `tp_new` which compiled classes don't. Allowing this
would require compiled types to be initialized more closely to the way
cpython does it which might need a lot of work.
- Lastly, when `__new__` is annotated with `@classmethod`, calling it
without the type parameter works in compiled code but raises an error in
interpreted. I'm not sure of the reason and it's difficult to make it a
compiler error because it's outside of what mypyc sees.

---------

Co-authored-by: Brian Schubert <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

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)