Address review comments

Author: p-sawicki

mypyc/irbuild/builder.py

@@ -1437,6 +1437,10 @@ def add_function(self, func_ir: FuncIR, line: int) -> None:
         self.function_names.add(name)
         self.functions.append(func_ir)
 
+    def get_current_class_ir(self) -> ClassIR | None:
+        type_info = self.fn_info.fitem.info
+        return self.mapper.type_to_ir.get(type_info)
+
 
 def gen_arg_defaults(builder: IRBuilder) -> None:
     """Generate blocks for arguments that have default values.

mypyc/irbuild/specialize.py

@@ -1012,10 +1012,7 @@ def translate_object_new(builder: IRBuilder, expr: CallExpr, callee: RefExpr) ->
     if fn.name != "__new__":
         return None
 
-    ir = None
-    for cls in builder.classes:
-        if cls.name == fn.class_name:
-            ir = cls
+    ir = builder.get_current_class_ir()
     if ir is None:
         return None
 

mypyc/irbuild/statement.py

@@ -33,6 +33,7 @@
     ListExpr,
     Lvalue,
     MatchStmt,
+    NameExpr,
     OperatorAssignmentStmt,
     RaiseStmt,
     ReturnStmt,
@@ -170,10 +171,38 @@ def transform_return_stmt(builder: IRBuilder, stmt: ReturnStmt) -> None:
     builder.nonlocal_control[-1].gen_return(builder, retval, stmt.line)
 
 
+def check_unsupported_cls_assignment(builder: IRBuilder, stmt: AssignmentStmt) -> None:
+    fn = builder.fn_info
+    method_args = fn.fitem.arg_names
+    if fn.name != "__new__" or len(method_args) == 0:
+        return
+
+    ir = builder.get_current_class_ir()
+    if ir is None or ir.inherits_python or not ir.is_ext_class:
+        return
+
+    cls_arg = method_args[0]
+    lvalues: list[Expression] = []
+    for lvalue in stmt.lvalues:
+        if isinstance(lvalue, (TupleExpr, ListExpr)):
+            lvalues += lvalue.items
+        else:
+            lvalues.append(lvalue)
+
+    for lvalue in lvalues:
+        if isinstance(lvalue, NameExpr) and lvalue.name == cls_arg:
+            # Disallowed because it could break the transformation of object.__new__ calls
+            # inside __new__ methods.
+            builder.error(
+                f"Assignment to argument {cls_arg} in __new__ method unsupported", stmt.line
+            )
+
+
 def transform_assignment_stmt(builder: IRBuilder, stmt: AssignmentStmt) -> None:
     lvalues = stmt.lvalues
     assert lvalues
     builder.disallow_class_assignments(lvalues, stmt.line)
+    check_unsupported_cls_assignment(builder, stmt)
     first_lvalue = lvalues[0]
     if stmt.type and isinstance(stmt.rvalue, TempNode):
         # This is actually a variable annotation without initializer. Don't generate

mypyc/test-data/irbuild-classes.test

@@ -1691,8 +1691,8 @@ class NotTransformed:
     def __new__(cls, val: int) -> Any:
         return super().__new__(str)
 
-    @classmethod
-    def factory(cls, val: int) -> NotTransformed:
+    def factory(cls: Any, val: int) -> Any:
+        cls = str
         return super().__new__(cls)
 
 [out]
@@ -1766,38 +1766,39 @@ L0:
 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
+    r0, r1 :: object
+    r2 :: str
+    r3, r4 :: object
+    r5 :: object[2]
+    r6 :: object_ptr
+    r7 :: object
+    r8 :: str
+    r9 :: object
+    r10 :: object[1]
+    r11 :: object_ptr
+    r12 :: object
 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)
+    r0 = load_address PyUnicode_Type
+    cls = r0
+    r1 = builtins :: module
+    r2 = 'super'
+    r3 = CPyObject_GetAttr(r1, r2)
+    r4 = __main__.NotTransformed :: type
+    r5 = [r4, cls]
+    r6 = load_address r5
+    r7 = PyObject_Vectorcall(r3, r6, 2, 0)
+    keep_alive r4, cls
+    r8 = '__new__'
+    r9 = CPyObject_GetAttr(r7, r8)
+    r10 = [cls]
+    r11 = load_address r10
+    r12 = PyObject_Vectorcall(r9, r11, 1, 0)
     keep_alive cls
-    r12 = cast(__main__.NotTransformed, r11)
     return r12
 
-[case testObjectDunderNew]
+[case testObjectDunderNew_64bit]
 from __future__ import annotations
+from mypy_extensions import mypyc_attr
 from typing import Any
 
 class Test:
@@ -1827,10 +1828,21 @@ class NotTransformed:
     def __new__(cls, val: int) -> Any:
         return object.__new__(str)
 
-    @classmethod
-    def factory(cls, val: int) -> NotTransformed:
+    def factory(cls: Any, val: int) -> Any:
+        cls = str
         return object.__new__(cls)
 
+@mypyc_attr(native_class=False)
+class NonNative:
+    def __new__(cls: Any) -> Any:
+        cls = str
+        return cls("str")
+
+class InheritsPython(dict):
+    def __new__(cls: Any) -> Any:
+        cls = dict
+        return cls({})
+
 class ObjectNewOutsideDunderNew:
     def __init__(self) -> None:
         object.__new__(ObjectNewOutsideDunderNew)
@@ -1899,25 +1911,69 @@ L0:
 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
+    r0, r1 :: object
+    r2 :: str
+    r3 :: object
+    r4 :: str
+    r5 :: object[2]
+    r6 :: object_ptr
+    r7 :: object
 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)
+    r0 = load_address PyUnicode_Type
+    cls = r0
+    r1 = builtins :: module
+    r2 = 'object'
+    r3 = CPyObject_GetAttr(r1, r2)
+    r4 = '__new__'
+    r5 = [r3, cls]
+    r6 = load_address r5
+    r7 = PyObject_VectorcallMethod(r4, r6, 9223372036854775810, 0)
+    keep_alive r3, cls
     return r7
+def __new___NonNative_obj.__get__(__mypyc_self__, instance, owner):
+    __mypyc_self__, instance, owner, r0 :: object
+    r1 :: bit
+    r2 :: object
+L0:
+    r0 = load_address _Py_NoneStruct
+    r1 = instance == r0
+    if r1 goto L1 else goto L2 :: bool
+L1:
+    return __mypyc_self__
+L2:
+    r2 = PyMethod_New(__mypyc_self__, instance)
+    return r2
+def __new___NonNative_obj.__call__(__mypyc_self__, cls):
+    __mypyc_self__ :: __main__.__new___NonNative_obj
+    cls, r0 :: object
+    r1 :: str
+    r2 :: object[1]
+    r3 :: object_ptr
+    r4 :: object
+L0:
+    r0 = load_address PyUnicode_Type
+    cls = r0
+    r1 = 'str'
+    r2 = [r1]
+    r3 = load_address r2
+    r4 = PyObject_Vectorcall(cls, r3, 1, 0)
+    keep_alive r1
+    return r4
+def InheritsPython.__new__(cls):
+    cls, r0 :: object
+    r1 :: dict
+    r2 :: object[1]
+    r3 :: object_ptr
+    r4 :: object
+L0:
+    r0 = load_address PyDict_Type
+    cls = r0
+    r1 = PyDict_New()
+    r2 = [r1]
+    r3 = load_address r2
+    r4 = PyObject_Vectorcall(cls, r3, 1, 0)
+    keep_alive r1
+    return r4
 def ObjectNewOutsideDunderNew.__init__(self):
     self :: __main__.ObjectNewOutsideDunderNew
     r0 :: object
@@ -1961,6 +2017,7 @@ L0:
 [case testUnsupportedDunderNew]
 from __future__ import annotations
 from mypy_extensions import mypyc_attr
+from typing import Any
 
 @mypyc_attr(native_class=False)
 class NonNative:
@@ -1980,6 +2037,21 @@ class InheritsPythonObjectNew(dict):
     def __new__(cls) -> InheritsPythonObjectNew:
         return object.__new__(cls)  # E: object.__new__() not supported for classes inheriting from non-native classes
 
+class ClsAssignment:
+    def __new__(cls: Any) -> Any:
+        cls = str  # E: Assignment to argument cls in __new__ method unsupported
+        return super().__new__(cls)
+
+class ClsTupleAssignment:
+    def __new__(class_i_want: Any, val: int) -> Any:
+        class_i_want, val = dict, 1  # E: Assignment to argument class_i_want in __new__ method unsupported
+        return object.__new__(class_i_want)
+
+class ClsListAssignment:
+    def __new__(cls: Any, val: str) -> Any:
+        [cls, val] = [object, "object"]  # E: Assignment to argument cls in __new__ method unsupported
+        return object.__new__(cls)
+
 [case testClassWithFreeList]
 from mypy_extensions import mypyc_attr, trait