[mypyc] Fix vtable pointer with inherited dunder new (#20302)

Fixes an issue where a subclass would have its vtable pointer set to the
base class' vtable when there is a `__new__` method defined in the base
class. This resulted in the subclass constructor calling the setup
function of the base class because mypyc transforms `object.__new__`
into the setup function.

The fix is to store the pointers to the setup functions in `tp_methods`
of type objects and look them up dynamically when instantiating new
objects.

---------

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

Author: p-sawicki

mypyc/codegen/emitclass.py

@@ -359,7 +359,7 @@ def emit_line() -> None:
     if cl.is_trait:
         generate_new_for_trait(cl, new_name, emitter)
 
-    generate_methods_table(cl, methods_name, emitter)
+    generate_methods_table(cl, methods_name, setup_name if generate_full else None, emitter)
     emit_line()
 
     flags = ["Py_TPFLAGS_DEFAULT", "Py_TPFLAGS_HEAPTYPE", "Py_TPFLAGS_BASETYPE"]
@@ -960,8 +960,17 @@ def generate_finalize_for_class(
     emitter.emit_line("}")
 
 
-def generate_methods_table(cl: ClassIR, name: str, emitter: Emitter) -> None:
+def generate_methods_table(
+    cl: ClassIR, name: str, setup_name: str | None, emitter: Emitter
+) -> None:
     emitter.emit_line(f"static PyMethodDef {name}[] = {{")
+    if setup_name:
+        # Store pointer to the setup function so it can be resolved dynamically
+        # in case of instance creation in __new__.
+        # CPy_SetupObject expects this method to be the first one in tp_methods.
+        emitter.emit_line(
+            f'{{"__internal_mypyc_setup", (PyCFunction){setup_name}, METH_O, NULL}},'
+        )
     for fn in cl.methods.values():
         if fn.decl.is_prop_setter or fn.decl.is_prop_getter or fn.internal:
             continue

mypyc/irbuild/specialize.py

@@ -99,7 +99,7 @@
     isinstance_dict,
 )
 from mypyc.primitives.float_ops import isinstance_float
-from mypyc.primitives.generic_ops import generic_setattr
+from mypyc.primitives.generic_ops import generic_setattr, setup_object
 from mypyc.primitives.int_ops import isinstance_int
 from mypyc.primitives.list_ops import isinstance_list, new_list_set_item_op
 from mypyc.primitives.misc_ops import isinstance_bool
@@ -1103,7 +1103,14 @@ def translate_object_new(builder: IRBuilder, expr: CallExpr, callee: RefExpr) ->
     method_args = fn.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))
+        subs = ir.subclasses()
+        if subs is not None and len(subs) == 0:
+            return builder.add(Call(ir.setup, [subtype], expr.line))
+        # Call a function that dynamically resolves the setup function of extension classes from the type object.
+        # This is necessary because the setup involves default attribute initialization and setting up
+        # the vtable which are specific to a given type and will not work if a subtype is created using
+        # the setup function of its base.
+        return builder.call_c(setup_object, [subtype], expr.line)
 
     return None
 

mypyc/lib-rt/CPy.h

@@ -958,6 +958,8 @@ static inline int CPyObject_GenericSetAttr(PyObject *self, PyObject *name, PyObj
     return _PyObject_GenericSetAttrWithDict(self, name, value, NULL);
 }
 
+PyObject *CPy_SetupObject(PyObject *type);
+
 #if CPY_3_11_FEATURES
 PyObject *CPy_GetName(PyObject *obj);
 #endif

mypyc/lib-rt/generic_ops.c

@@ -62,3 +62,23 @@ PyObject *CPyObject_GetSlice(PyObject *obj, CPyTagged start, CPyTagged end) {
     Py_DECREF(slice);
     return result;
 }
+
+typedef PyObject *(*SetupFunction)(PyObject *);
+
+PyObject *CPy_SetupObject(PyObject *type) {
+    PyTypeObject *tp = (PyTypeObject *)type;
+    PyMethodDef *def = NULL;
+    for(; tp; tp = tp->tp_base) {
+        def = tp->tp_methods;
+        if (!def || !def->ml_name) {
+            continue;
+        }
+
+        if (!strcmp(def->ml_name, "__internal_mypyc_setup")) {
+            return ((SetupFunction)(void(*)(void))def->ml_meth)(type);
+        }
+    }
+
+    PyErr_SetString(PyExc_RuntimeError, "Internal mypyc error: Unable to find object setup function");
+    return NULL;
+}

mypyc/primitives/generic_ops.py

@@ -417,3 +417,10 @@
     c_function_name="CPyObject_GenericSetAttr",
     error_kind=ERR_NEG_INT,
 )
+
+setup_object = custom_op(
+    arg_types=[object_rprimitive],
+    return_type=object_rprimitive,
+    c_function_name="CPy_SetupObject",
+    error_kind=ERR_MAGIC,
+)

mypyc/test-data/irbuild-classes.test

@@ -1685,6 +1685,13 @@ class Test:
         obj.val = val
         return obj
 
+class Test2:
+    def __new__(cls) -> Test2:
+        return super().__new__(cls)
+
+class Sub(Test2):
+    pass
+
 def fn() -> Test:
     return Test.__new__(Test, 42)
 
@@ -1719,6 +1726,13 @@ L0:
     obj = r0
     obj.val = val; r1 = is_error
     return obj
+def Test2.__new__(cls):
+    cls, r0 :: object
+    r1 :: __main__.Test2
+L0:
+    r0 = CPy_SetupObject(cls)
+    r1 = cast(__main__.Test2, r0)
+    return r1
 def fn():
     r0 :: object
     r1 :: __main__.Test
@@ -1822,6 +1836,13 @@ class Test:
         obj.val = val
         return obj
 
+class Test2:
+    def __new__(cls) -> Test2:
+        return object.__new__(cls)
+
+class Sub(Test2):
+    pass
+
 def fn() -> Test:
     return Test.__new__(Test, 42)
 
@@ -1874,6 +1895,13 @@ L0:
     obj = r0
     obj.val = val; r1 = is_error
     return obj
+def Test2.__new__(cls):
+    cls, r0 :: object
+    r1 :: __main__.Test2
+L0:
+    r0 = CPy_SetupObject(cls)
+    r1 = cast(__main__.Test2, r0)
+    return r1
 def fn():
     r0 :: object
     r1 :: __main__.Test

mypyc/test-data/run-classes.test

@@ -3859,6 +3859,7 @@ Add(1, 0)=1
 [case testInheritedDunderNew]
 from __future__ import annotations
 from mypy_extensions import mypyc_attr
+from testutil import assertRaises
 from typing_extensions import Self
 
 from m import interpreted_subclass
@@ -3875,19 +3876,32 @@ class Base:
     def __init__(self, val: int) -> None:
         self.init_val = val
 
+    def method(self) -> int:
+        raise NotImplementedError
+
 class Sub(Base):
+
     def __new__(cls, val: int) -> Self:
         return super().__new__(cls, val + 1)
 
     def __init__(self, val: int) -> None:
         super().__init__(val)
         self.init_val = self.init_val * 2
 
+    def method(self) -> int:
+        return 0
+
 class SubWithoutNew(Base):
+    sub_only_str = ""
+    sub_only_int: int
+
     def __init__(self, val: int) -> None:
         super().__init__(val)
         self.init_val = self.init_val * 2
 
+    def method(self) -> int:
+        return 1
+
 class BaseWithoutInterpretedSubclasses:
     val: int
 
@@ -3899,6 +3913,9 @@ class BaseWithoutInterpretedSubclasses:
     def __init__(self, val: int) -> None:
         self.init_val = val
 
+    def method(self) -> int:
+        raise NotImplementedError
+
 class SubNoInterpreted(BaseWithoutInterpretedSubclasses):
     def __new__(cls, val: int) -> Self:
         return super().__new__(cls, val + 1)
@@ -3907,55 +3924,77 @@ class SubNoInterpreted(BaseWithoutInterpretedSubclasses):
         super().__init__(val)
         self.init_val = self.init_val * 2
 
+    def method(self) -> int:
+        return 0
+
 class SubNoInterpretedWithoutNew(BaseWithoutInterpretedSubclasses):
     def __init__(self, val: int) -> None:
         super().__init__(val)
         self.init_val = self.init_val * 2
 
+    def method(self) -> int:
+        return 1
+
 def test_inherited_dunder_new() -> None:
     b = Base(42)
     assert type(b) == Base
     assert b.val == 43
     assert b.init_val == 42
+    with assertRaises(NotImplementedError):
+        b.method()
 
     s = Sub(42)
     assert type(s) == Sub
     assert s.val == 44
     assert s.init_val == 84
+    assert s.method() == 0
 
     s2 = SubWithoutNew(42)
     assert type(s2) == SubWithoutNew
     assert s2.val == 43
     assert s2.init_val == 84
+    assert s2.method() == 1
+    assert s2.sub_only_str == ""
+    with assertRaises(AttributeError):
+        s2.sub_only_int
+    s2.sub_only_int = 11
+    assert s2.sub_only_int == 11
 
 def test_inherited_dunder_new_without_interpreted_subclasses() -> None:
     b = BaseWithoutInterpretedSubclasses(42)
     assert type(b) == BaseWithoutInterpretedSubclasses
     assert b.val == 43
     assert b.init_val == 42
+    with assertRaises(NotImplementedError):
+        b.method()
 
     s = SubNoInterpreted(42)
     assert type(s) == SubNoInterpreted
     assert s.val == 44
     assert s.init_val == 84
+    assert s.method() == 0
 
     s2 = SubNoInterpretedWithoutNew(42)
     assert type(s2) == SubNoInterpretedWithoutNew
     assert s2.val == 43
     assert s2.init_val == 84
+    assert s2.method() == 1
 
 def test_interpreted_subclass() -> None:
     interpreted_subclass(Base)
 
 [file m.py]
 from __future__ import annotations
+from testutil import assertRaises
 from typing_extensions import Self
 
 def interpreted_subclass(base) -> None:
     b = base(42)
     assert type(b) == base
     assert b.val == 43
     assert b.init_val == 42
+    with assertRaises(NotImplementedError):
+        b.method()
 
     class InterpretedSub(base):
         def __new__(cls, val: int) -> Self:
@@ -3965,20 +4004,36 @@ def interpreted_subclass(base) -> None:
             super().__init__(val)
             self.init_val : int = self.init_val * 2
 
+        def method(self) -> int:
+            return 3
+
     s = InterpretedSub(42)
     assert type(s) == InterpretedSub
     assert s.val == 44
     assert s.init_val == 84
+    assert s.method() == 3
 
     class InterpretedSubWithoutNew(base):
+        sub_only_str = ""
+        sub_only_int: int
+
         def __init__(self, val: int) -> None:
             super().__init__(val)
             self.init_val : int = self.init_val * 2
 
+        def method(self) -> int:
+            return 4
+
     s2 = InterpretedSubWithoutNew(42)
     assert type(s2) == InterpretedSubWithoutNew
     assert s2.val == 43
     assert s2.init_val == 84
+    assert s2.method() == 4
+    assert s2.sub_only_str == ""
+    with assertRaises(AttributeError):
+        s2.sub_only_int
+    s2.sub_only_int = 11
+    assert s2.sub_only_int == 11
 
 [typing fixtures/typing-full.pyi]