[mypyc] Transform object.__new__ inside __new__

Author: p-sawicki

mypyc/irbuild/expression.py

@@ -57,7 +57,6 @@
 from mypyc.ir.ops import (
     Assign,
     BasicBlock,
-    Call,
     ComparisonOp,
     Integer,
     LoadAddress,
@@ -98,7 +97,11 @@
     join_formatted_strings,
     tokenizer_printf_style,
 )
-from mypyc.irbuild.specialize import apply_function_specialization, apply_method_specialization
+from mypyc.irbuild.specialize import (
+    apply_function_specialization,
+    apply_method_specialization,
+    translate_object_new,
+)
 from mypyc.primitives.bytes_ops import bytes_slice_op
 from mypyc.primitives.dict_ops import dict_get_item_op, dict_new_op, exact_dict_set_item_op
 from mypyc.primitives.generic_ops import iter_op, name_op
@@ -473,35 +476,15 @@ def translate_super_method_call(builder: IRBuilder, expr: CallExpr, callee: Supe
         if callee.name in base.method_decls:
             break
     else:
+        if callee.name == "__new__":
+            result = translate_object_new(builder, expr, MemberExpr(callee.call, "__new__"))
+            if result:
+                return result
         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)

mypyc/irbuild/specialize.py

@@ -36,6 +36,7 @@
 from mypy.types import AnyType, TypeOfAny
 from mypyc.ir.ops import (
     BasicBlock,
+    Call,
     Extend,
     Integer,
     RaiseStandardError,
@@ -68,6 +69,7 @@
     is_list_rprimitive,
     is_uint8_rprimitive,
     list_rprimitive,
+    object_rprimitive,
     set_rprimitive,
     str_rprimitive,
     uint8_rprimitive,
@@ -1002,3 +1004,38 @@ def translate_ord(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value
     if isinstance(arg, (StrExpr, BytesExpr)) and len(arg.value) == 1:
         return Integer(ord(arg.value))
     return None
+
+
+@specialize_function("__new__", object_rprimitive)
+def translate_object_new(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    fn = builder.fn_info
+    if fn.name != "__new__":
+        return None
+
+    ir = None
+    for cls in builder.classes:
+        if cls.name == fn.class_name:
+            ir = cls
+    if ir is None:
+        return None
+
+    call = "object.__new__()"
+    if not ir.is_ext_class:
+        builder.error(f"{call} not supported for non-extension classes", expr.line)
+        return None
+    if ir.inherits_python:
+        builder.error(
+            f"{call} not supported for classes inheriting from non-native classes", expr.line
+        )
+        return None
+
+    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 = 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))
+
+    return None

mypyc/test-data/irbuild-classes.test

@@ -1662,6 +1662,7 @@ L0:
 
 [case testDunderNew]
 from __future__ import annotations
+from typing import Any
 
 class Test:
     val: int
@@ -1686,6 +1687,14 @@ class NewClassMethod:
 def fn2() -> NewClassMethod:
     return NewClassMethod.__new__(42)
 
+class NotTransformed:
+    def __new__(cls, val: int) -> Any:
+        return super().__new__(str)
+    
+    @classmethod
+    def factory(cls, val: int) -> NotTransformed:
+        return super().__new__(cls)
+
 [out]
 def Test.__new__(cls, val):
     cls :: object
@@ -1721,6 +1730,233 @@ L0:
     r0 = __main__.NewClassMethod :: type
     r1 = NewClassMethod.__new__(r0, 84)
     return r1
+def NotTransformed.__new__(cls, val):
+    cls :: object
+    val :: int
+    r0 :: object
+    r1 :: str
+    r2, r3 :: object
+    r4 :: object[2]
+    r5 :: object_ptr
+    r6 :: object
+    r7 :: str
+    r8, r9 :: object
+    r10 :: object[1]
+    r11 :: object_ptr
+    r12 :: object
+    r13 :: str
+L0:
+    r0 = builtins :: module
+    r1 = 'super'
+    r2 = CPyObject_GetAttr(r0, r1)
+    r3 = __main__.NotTransformed :: type
+    r4 = [r3, cls]
+    r5 = load_address r4
+    r6 = PyObject_Vectorcall(r2, r5, 2, 0)
+    keep_alive r3, cls
+    r7 = '__new__'
+    r8 = CPyObject_GetAttr(r6, r7)
+    r9 = load_address PyUnicode_Type
+    r10 = [r9]
+    r11 = load_address r10
+    r12 = PyObject_Vectorcall(r8, r11, 1, 0)
+    keep_alive r9
+    r13 = cast(str, r12)
+    return r13
+def NotTransformed.factory(cls, val):
+    cls :: object
+    val :: int
+    r0 :: object
+    r1 :: str
+    r2, r3 :: object
+    r4 :: object[2]
+    r5 :: object_ptr
+    r6 :: object
+    r7 :: str
+    r8 :: object
+    r9 :: object[1]
+    r10 :: object_ptr
+    r11 :: object
+    r12 :: __main__.NotTransformed
+L0:
+    r0 = builtins :: module
+    r1 = 'super'
+    r2 = CPyObject_GetAttr(r0, r1)
+    r3 = __main__.NotTransformed :: type
+    r4 = [r3, cls]
+    r5 = load_address r4
+    r6 = PyObject_Vectorcall(r2, r5, 2, 0)
+    keep_alive r3, cls
+    r7 = '__new__'
+    r8 = CPyObject_GetAttr(r6, r7)
+    r9 = [cls]
+    r10 = load_address r9
+    r11 = PyObject_Vectorcall(r8, r10, 1, 0)
+    keep_alive cls
+    r12 = cast(__main__.NotTransformed, r11)
+    return r12
+
+[case testObjectDunderNew]
+from __future__ import annotations
+from typing import Any
+
+class Test:
+    val: int
+
+    def __new__(cls, val: int) -> Test:
+        obj = object.__new__(cls)
+        obj.val = val
+        return obj
+
+def fn() -> Test:
+    return Test.__new__(Test, 42)
+
+class NewClassMethod:
+    val: int
+
+    @classmethod
+    def __new__(cls, val: int) -> NewClassMethod:
+        obj = object.__new__(cls)
+        obj.val = val
+        return obj
+
+def fn2() -> NewClassMethod:
+    return NewClassMethod.__new__(42)
+
+class NotTransformed:
+    def __new__(cls, val: int) -> Any:
+        return object.__new__(str)
+    
+    @classmethod
+    def factory(cls, val: int) -> NotTransformed:
+        return object.__new__(cls)
+
+class ObjectNewOutsideDunderNew:
+    def __init__(self) -> None:
+        object.__new__(ObjectNewOutsideDunderNew)
+
+def object_new_outside_class() -> None:
+    object.__new__(Test)
+
+[out]
+def Test.__new__(cls, val):
+    cls :: object
+    val :: int
+    r0, obj :: __main__.Test
+    r1 :: bool
+L0:
+    r0 = __mypyc__Test_setup(cls)
+    obj = r0
+    obj.val = val; r1 = is_error
+    return obj
+def fn():
+    r0 :: object
+    r1 :: __main__.Test
+L0:
+    r0 = __main__.Test :: type
+    r1 = Test.__new__(r0, 84)
+    return r1
+def NewClassMethod.__new__(cls, val):
+    cls :: object
+    val :: int
+    r0, obj :: __main__.NewClassMethod
+    r1 :: bool
+L0:
+    r0 = __mypyc__NewClassMethod_setup(cls)
+    obj = r0
+    obj.val = val; r1 = is_error
+    return obj
+def fn2():
+    r0 :: object
+    r1 :: __main__.NewClassMethod
+L0:
+    r0 = __main__.NewClassMethod :: type
+    r1 = NewClassMethod.__new__(r0, 84)
+    return r1
+def NotTransformed.__new__(cls, val):
+    cls :: object
+    val :: int
+    r0 :: object
+    r1 :: str
+    r2, r3 :: object
+    r4 :: str
+    r5 :: object[2]
+    r6 :: object_ptr
+    r7 :: object
+    r8 :: str
+L0:
+    r0 = builtins :: module
+    r1 = 'object'
+    r2 = CPyObject_GetAttr(r0, r1)
+    r3 = load_address PyUnicode_Type
+    r4 = '__new__'
+    r5 = [r2, r3]
+    r6 = load_address r5
+    r7 = PyObject_VectorcallMethod(r4, r6, 9223372036854775810, 0)
+    keep_alive r2, r3
+    r8 = cast(str, r7)
+    return r8
+def NotTransformed.factory(cls, val):
+    cls :: object
+    val :: int
+    r0 :: object
+    r1 :: str
+    r2 :: object
+    r3 :: str
+    r4 :: object[2]
+    r5 :: object_ptr
+    r6 :: object
+    r7 :: __main__.NotTransformed
+L0:
+    r0 = builtins :: module
+    r1 = 'object'
+    r2 = CPyObject_GetAttr(r0, r1)
+    r3 = '__new__'
+    r4 = [r2, cls]
+    r5 = load_address r4
+    r6 = PyObject_VectorcallMethod(r3, r5, 9223372036854775810, 0)
+    keep_alive r2, cls
+    r7 = cast(__main__.NotTransformed, r6)
+    return r7
+def ObjectNewOutsideDunderNew.__init__(self):
+    self :: __main__.ObjectNewOutsideDunderNew
+    r0 :: object
+    r1 :: str
+    r2, r3 :: object
+    r4 :: str
+    r5 :: object[2]
+    r6 :: object_ptr
+    r7 :: object
+L0:
+    r0 = builtins :: module
+    r1 = 'object'
+    r2 = CPyObject_GetAttr(r0, r1)
+    r3 = __main__.ObjectNewOutsideDunderNew :: type
+    r4 = '__new__'
+    r5 = [r2, r3]
+    r6 = load_address r5
+    r7 = PyObject_VectorcallMethod(r4, r6, 9223372036854775810, 0)
+    keep_alive r2, r3
+    return 1
+def object_new_outside_class():
+    r0 :: object
+    r1 :: str
+    r2, r3 :: object
+    r4 :: str
+    r5 :: object[2]
+    r6 :: object_ptr
+    r7 :: object
+L0:
+    r0 = builtins :: module
+    r1 = 'object'
+    r2 = CPyObject_GetAttr(r0, r1)
+    r3 = __main__.Test :: type
+    r4 = '__new__'
+    r5 = [r2, r3]
+    r6 = load_address r5
+    r7 = PyObject_VectorcallMethod(r4, r6, 9223372036854775810, 0)
+    keep_alive r2, r3
+    return 1
 
 [case testUnsupportedDunderNew]
 from __future__ import annotations
@@ -1729,11 +1965,20 @@ from mypy_extensions import mypyc_attr
 @mypyc_attr(native_class=False)
 class NonNative:
     def __new__(cls) -> NonNative:
-        return super().__new__(cls)  # E: super().__new__() not supported for non-extension classes
+        return super().__new__(cls)  # E: object.__new__() not supported for non-extension classes
 
 class InheritsPython(dict):
     def __new__(cls) -> InheritsPython:
-        return super().__new__(cls)  # E: super().__new__() not supported for classes inheriting from non-native classes
+        return super().__new__(cls)  # E: object.__new__() not supported for classes inheriting from non-native classes
+
+@mypyc_attr(native_class=False)
+class NonNativeObjectNew:
+    def __new__(cls) -> NonNativeObjectNew:
+        return object.__new__(cls)  # E: object.__new__() not supported for non-extension classes
+
+class InheritsPythonObjectNew(dict):
+    def __new__(cls) -> InheritsPythonObjectNew:
+        return object.__new__(cls)  # E: object.__new__() not supported for classes inheriting from non-native classes
 
 [case testClassWithFreeList]
 from mypy_extensions import mypyc_attr, trait