[mypyc] Call __setattr__ for assignments that come before its definition

mypyc/irbuild/builder.py

@@ -929,11 +929,12 @@ def assign(self, target: Register | AssignmentTarget, rvalue_reg: Value, line: i
             self.add(Assign(target.register, rvalue_reg, line))
         elif isinstance(target, AssignmentTargetAttr):
             if isinstance(target.obj_type, RInstance):
-                setattr = target.obj_type.class_ir.get_method("__setattr__")
-                if setattr:
+                # Look at the method declarations, since the method bodies are added as
+                # they are compiled, and "__setattr__" may be compiled after this.
+                if target.obj_type.class_ir.has_method("__setattr__"):
                     key = self.load_str(target.attr, line)
                     boxed_reg = self.builder.box(rvalue_reg)
-                    call = MethodCall(target.obj, setattr.name, [key, boxed_reg], line)
+                    call = MethodCall(target.obj, "__setattr__", [key, boxed_reg], line)
                     self.add(call)
                 else:
                     rvalue_reg = self.coerce_rvalue(rvalue_reg, target.type, line)

mypyc/test-data/irbuild-classes.test

@@ -2764,6 +2764,64 @@ L0:
     r2 = r1 >= 0 :: signed
     return 1
 
+[case testSetAttrDefinedAfterAssignment]
+def assign(i: "SetAttr") -> None:
+    i.regular_attr = 1
+
+class SetAttr:
+    def __init__(self) -> None:
+        self.regular_attr = 0
+
+    def __setattr__(self, key: str, val: object) -> None:
+        pass
+
+[out]
+def assign(i):
+    i :: __main__.SetAttr
+    r0 :: str
+    r1 :: object
+    r2 :: None
+L0:
+    r0 = 'regular_attr'
+    r1 = object 1
+    r2 = i.__setattr__(r0, r1)
+    return 1
+def SetAttr.__init__(self):
+    self :: __main__.SetAttr
+    r0 :: str
+    r1 :: object
+    r2 :: None
+L0:
+    r0 = 'regular_attr'
+    r1 = object 0
+    r2 = self.__setattr__(r0, r1)
+    return 1
+def SetAttr.__setattr__(self, key, val):
+    self :: __main__.SetAttr
+    key :: str
+    val :: object
+L0:
+    return 1
+def SetAttr.__setattr____wrapper(__mypyc_self__, attr, value):
+    __mypyc_self__ :: __main__.SetAttr
+    attr, value :: object
+    r0 :: bit
+    r1 :: i32
+    r2 :: bit
+    r3 :: str
+    r4 :: None
+L0:
+    r0 = value == 0
+    if r0 goto L1 else goto L2 :: bool
+L1:
+    r1 = CPyObject_GenericSetAttr(__mypyc_self__, attr, 0)
+    r2 = r1 >= 0 :: signed
+    return 0
+L2:
+    r3 = cast(str, attr)
+    r4 = __mypyc_self__.__setattr__(r3, value)
+    return 0
+
 [case testUntransformedSetAttr_64bit]
 from mypy_extensions import mypyc_attr
 

mypyc/test-data/run-classes.test

@@ -6096,6 +6096,68 @@ test_no_setattr_nonnative()
 
 [typing fixtures/typing-full.pyi]
 
+[case testDunderSetAttrDefinedAfterAssignment]
+setattr_calls: list[str] = []
+
+def assign_in_function(obj: "LateSetAttr") -> None:
+    obj.attr = 1
+
+class Other:
+    def assign_in_other_class(self, obj: "LateSetAttr") -> None:
+        obj.attr = 2
+
+class LateSetAttr:
+    def __init__(self) -> None:
+        self.attr = 0
+
+    def assign_in_method(self) -> None:
+        self.attr = 3
+        self.attr += 1
+
+    def __setattr__(self, key: str, val: object) -> None:
+        setattr_calls.append(f"{key}={val}")
+        super().__setattr__(key, val)
+
+class AssignInSetAttr:
+    attr: int
+    assigned: bool
+
+    def __init__(self) -> None:
+        object.__setattr__(self, "attr", 0)
+        object.__setattr__(self, "assigned", False)
+
+    def __setattr__(self, key: str, val: object) -> None:
+        setattr_calls.append(f"{key}={val}")
+        if key == "attr":
+            # This calls __setattr__ again.
+            self.assigned = True
+        super().__setattr__(key, val)
+
+def test_setattr_defined_after_assignment() -> None:
+    setattr_calls.clear()
+    obj = LateSetAttr()
+    assert setattr_calls == ["attr=0"]
+
+    assign_in_function(obj)
+    assert setattr_calls == ["attr=0", "attr=1"]
+
+    Other().assign_in_other_class(obj)
+    assert setattr_calls == ["attr=0", "attr=1", "attr=2"]
+
+    obj.assign_in_method()
+    assert setattr_calls == ["attr=0", "attr=1", "attr=2", "attr=3", "attr=4"]
+    assert obj.attr == 4
+
+def test_assignment_in_setattr() -> None:
+    setattr_calls.clear()
+    obj = AssignInSetAttr()
+    assert setattr_calls == []
+
+    obj.attr = 1
+    assert setattr_calls == ["attr=1", "assigned=True"]
+    assert obj.attr == 1
+    assert obj.assigned
+
 [case testDelAttrWithDeletableAttr]
 from testutil import assertRaises