[mypyc] Generate __setattr__ wrapper

Author: p-sawicki

mypyc/codegen/emitclass.py

@@ -62,6 +62,7 @@ def dunder_attr_slot(cl: ClassIR, fn: FuncIR, emitter: Emitter) -> str:
     "__hash__": ("tp_hash", generate_hash_wrapper),
     "__get__": ("tp_descr_get", generate_get_wrapper),
     "__getattr__": ("tp_getattro", dunder_attr_slot),
+    "__setattr__": ("tp_setattro", dunder_attr_slot),
 }
 
 AS_MAPPING_SLOT_DEFS: SlotTable = {

mypyc/irbuild/builder.py

@@ -76,6 +76,7 @@
     Integer,
     IntOp,
     LoadStatic,
+    MethodCall,
     Op,
     PrimitiveDescription,
     RaiseStandardError,
@@ -735,8 +736,15 @@ def assign(self, target: Register | AssignmentTarget, rvalue_reg: Value, line: i
             self.add(Assign(target.register, rvalue_reg))
         elif isinstance(target, AssignmentTargetAttr):
             if isinstance(target.obj_type, RInstance):
-                rvalue_reg = self.coerce_rvalue(rvalue_reg, target.type, line)
-                self.add(SetAttr(target.obj, target.attr, rvalue_reg, line))
+                setattr = target.obj_type.class_ir.get_method("__setattr__")
+                if setattr:
+                    key = self.load_str(target.attr)
+                    boxed_reg = self.builder.box(rvalue_reg)
+                    call = MethodCall(target.obj, setattr.name, [key, boxed_reg], line)
+                    self.add(call)
+                else:
+                    rvalue_reg = self.coerce_rvalue(rvalue_reg, target.type, line)
+                    self.add(SetAttr(target.obj, target.attr, rvalue_reg, line))
             else:
                 key = self.load_str(target.attr)
                 boxed_reg = self.builder.box(rvalue_reg)

mypyc/irbuild/expression.py

@@ -101,6 +101,7 @@
     apply_function_specialization,
     apply_method_specialization,
     translate_object_new,
+    translate_object_setattr,
 )
 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
@@ -485,6 +486,12 @@ def translate_super_method_call(builder: IRBuilder, expr: CallExpr, callee: Supe
                 # Call translates to object.__init__(self), which is a
                 # no-op, so omit the call.
                 return builder.none()
+            elif callee.name == "__setattr__":
+                result = translate_object_setattr(
+                    builder, expr, MemberExpr(callee.call, "__setattr__")
+                )
+                if result:
+                    return result
         return translate_call(builder, expr, callee)
 
     decl = base.method_decl(callee.name)

mypyc/irbuild/function.py

@@ -57,6 +57,7 @@
 from mypyc.ir.rtypes import (
     RInstance,
     bool_rprimitive,
+    c_int_rprimitive,
     dict_rprimitive,
     int_rprimitive,
     object_rprimitive,
@@ -415,6 +416,34 @@ def generate_getattr_wrapper(builder: IRBuilder, cdef: ClassDef, getattr: FuncDe
         builder.add(Return(getattr_result, line))
 
 
+def generate_setattr_wrapper(builder: IRBuilder, cdef: ClassDef, setattr: FuncDef) -> None:
+    """
+    Generate a wrapper function for __setattr__ that can be put into the tp_setattro slot.
+    The wrapper takes two arguments besides self - attribute name and the new value.
+    Returns 0 on success and -1 on failure. Restrictions are similar to the __getattr__
+    wrapper above.
+
+    This one is simpler because to match interpreted python semantics it's enough to always
+    call the user-provided function, including for names matching regular attributes.
+    """
+    name = setattr.name + "__wrapper"
+    ir = builder.mapper.type_to_ir[cdef.info]
+    line = setattr.line
+
+    error_base = f'"__setattr__" not supported in class "{cdef.name}" because '
+    if ir.allow_interpreted_subclasses:
+        builder.error(error_base + "it allows interpreted subclasses", line)
+    if ir.inherits_python:
+        builder.error(error_base + "it inherits from a non-native class", line)
+
+    with builder.enter_method(ir, name, c_int_rprimitive, internal=True):
+        attr_arg = builder.add_argument("attr", object_rprimitive)
+        value_arg = builder.add_argument("value", object_rprimitive)
+
+        builder.gen_method_call(builder.self(), setattr.name, [attr_arg, value_arg], None, line)
+        builder.add(Return(Integer(0, c_int_rprimitive), line))
+
+
 def handle_ext_method(builder: IRBuilder, cdef: ClassDef, fdef: FuncDef) -> None:
     # Perform the function of visit_method for methods inside extension classes.
     name = fdef.name
@@ -483,6 +512,8 @@ def handle_ext_method(builder: IRBuilder, cdef: ClassDef, fdef: FuncDef) -> None
 
     if fdef.name == "__getattr__":
         generate_getattr_wrapper(builder, cdef, fdef)
+    elif fdef.name == "__setattr__":
+        generate_setattr_wrapper(builder, cdef, fdef)
 
 
 def handle_non_ext_method(

mypyc/irbuild/specialize.py

@@ -42,6 +42,7 @@
     Integer,
     RaiseStandardError,
     Register,
+    SetAttr,
     Truncate,
     Unreachable,
     Value,
@@ -97,6 +98,7 @@
     isinstance_dict,
 )
 from mypyc.primitives.float_ops import isinstance_float
+from mypyc.primitives.generic_ops import generic_setattr
 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
@@ -1007,19 +1009,28 @@ def translate_ord(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> 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
-
-    is_super_new = isinstance(expr.callee, SuperExpr)
-    is_object_new = (
+def is_object(callee: RefExpr) -> bool:
+    """
+    Returns True for object.<name> calls
+    """
+    return (
         isinstance(callee, MemberExpr)
         and isinstance(callee.expr, NameExpr)
         and callee.expr.fullname == "builtins.object"
     )
-    if not (is_super_new or is_object_new):
+
+
+def is_super_or_object(expr: CallExpr, callee: RefExpr) -> bool:
+    """
+    Returns True for super().<name> or object.<name> calls.
+    """
+    return isinstance(expr.callee, SuperExpr) or is_object(callee)
+
+
+@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__" or not is_super_or_object(expr, callee):
         return None
 
     ir = builder.get_current_class_ir()
@@ -1046,3 +1057,28 @@ def translate_object_new(builder: IRBuilder, expr: CallExpr, callee: RefExpr) ->
         return builder.add(Call(ir.setup, [subtype], expr.line))
 
     return None
+
+
+@specialize_function("__setattr__", object_rprimitive)
+def translate_object_setattr(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    is_super = isinstance(expr.callee, SuperExpr)
+    is_object_callee = is_object(callee)
+    if not ((is_super and len(expr.args) >= 2) or (is_object_callee and len(expr.args) >= 3)):
+        return None
+
+    self_reg = builder.accept(expr.args[0]) if is_object_callee else builder.self()
+    ir = builder.get_current_class_ir()
+    # Need to offset by 1 for super().__setattr__ calls because there is no self arg in this case.
+    name_idx = 0 if is_super else 1
+    value_idx = 1 if is_super else 2
+    attr_name = expr.args[name_idx]
+    attr_value = expr.args[value_idx]
+    value = builder.accept(attr_value)
+
+    if isinstance(attr_name, StrExpr) and ir and ir.has_attr(attr_name.value):
+        name = attr_name.value
+        value = builder.coerce(value, ir.attributes[name], expr.line)
+        return builder.add(SetAttr(self_reg, name, value, expr.line))
+
+    name_reg = builder.accept(attr_name)
+    return builder.call_c(generic_setattr, [self_reg, name_reg, value], expr.line)

mypyc/lib-rt/CPy.h

@@ -952,6 +952,9 @@ void CPyTrace_LogEvent(const char *location, const char *line, const char *op, c
 static inline PyObject *CPyObject_GenericGetAttr(PyObject *self, PyObject *name) {
     return _PyObject_GenericGetAttrWithDict(self, name, NULL, 1);
 }
+static inline int CPyObject_GenericSetAttr(PyObject *self, PyObject *name, PyObject *value) {
+    return _PyObject_GenericSetAttrWithDict(self, name, value, NULL);
+}
 
 #if CPY_3_11_FEATURES
 PyObject *CPy_GetName(PyObject *obj);

mypyc/primitives/generic_ops.py

@@ -410,3 +410,10 @@
     error_kind=ERR_NEVER,
     returns_null=True,
 )
+
+generic_setattr = custom_op(
+    arg_types=[object_rprimitive, object_rprimitive, object_rprimitive],
+    return_type=c_int_rprimitive,
+    c_function_name="CPyObject_GenericSetAttr",
+    error_kind=ERR_NEG_INT,
+)

mypyc/test-data/fixtures/ir.py

@@ -45,6 +45,7 @@ def __init__(self) -> None: pass
     def __eq__(self, x: object) -> bool: pass
     def __ne__(self, x: object) -> bool: pass
     def __str__(self) -> str: pass
+    def __setattr__(self, k: str, v: object) -> None: pass
 
 class type:
     def __init__(self, o: object) -> None: ...