Checking dataclass type before copying attributes.

Author: advait-dixit

mypyc/irbuild/classdef.py

@@ -377,9 +377,10 @@ def finalize(self, ir: ClassIR) -> None:
         dec = self.builder.accept(
             next(d for d in self.cdef.decorators if is_dataclass_decorator(d))
         )
+        dataclass_type_val = self.builder.load_str(dataclass_type(self.cdef) or "unknown")
         self.builder.call_c(
             dataclass_sleight_of_hand,
-            [dec, self.type_obj, self.non_ext.dict, self.non_ext.anns],
+            [dec, self.type_obj, self.non_ext.dict, self.non_ext.anns, dataclass_type_val],
             self.cdef.line,
         )
 

mypyc/irbuild/util.py

@@ -73,6 +73,8 @@ def is_dataclass(cdef: ClassDef) -> bool:
     return any(is_dataclass_decorator(d) for d in cdef.decorators)
 
 
+# The string values returned by this function are inspected in
+# mypyc/lib-rt/misc_ops.c:CPyDataclass_SleightOfHand(...).
 def dataclass_type(cdef: ClassDef) -> str | None:
     for d in cdef.decorators:
         typ = dataclass_decorator_type(d)

mypyc/lib-rt/CPy.h

@@ -860,7 +860,8 @@ PyObject *CPyType_FromTemplateWrapper(PyObject *template_,
                                       PyObject *orig_bases,
                                       PyObject *modname);
 int CPyDataclass_SleightOfHand(PyObject *dataclass_dec, PyObject *tp,
-                               PyObject *dict, PyObject *annotations);
+                               PyObject *dict, PyObject *annotations,
+                               PyObject *dataclass_type);
 PyObject *CPyPickle_SetState(PyObject *obj, PyObject *state);
 PyObject *CPyPickle_GetState(PyObject *obj);
 CPyTagged CPyTagged_Id(PyObject *o);

mypyc/lib-rt/misc_ops.c

@@ -347,10 +347,12 @@ static int _CPy_UpdateObjFromDict(PyObject *obj, PyObject *dict)
  *   tp: The class we are making a dataclass
  *   dict: The dictionary containing values that dataclasses needs
  *   annotations: The type annotation dictionary
+ *   dataclass_type: A str object with the return value of util.py:dataclass_type()
  */
 int
 CPyDataclass_SleightOfHand(PyObject *dataclass_dec, PyObject *tp,
-                           PyObject *dict, PyObject *annotations) {
+                           PyObject *dict, PyObject *annotations,
+                           PyObject *dataclass_type) {
     PyTypeObject *ttp = (PyTypeObject *)tp;
     Py_ssize_t pos;
     PyObject *res = NULL;
@@ -381,16 +383,23 @@ CPyDataclass_SleightOfHand(PyObject *dataclass_dec, PyObject *tp,
     if (!res) {
         goto fail;
     }
-    // These attributes are added or modified by @attr.s(slots=True).
-    const char * const keys[] = {"__attrs_attrs__", "__attrs_own_setattr__", "__init__", ""};
-    for (const char * const *key_iter = keys; **key_iter != '\0'; key_iter++) {
-        PyObject *value = NULL;
-        int rv = PyObject_GetOptionalAttrString(res, *key_iter, &value);
-        if (rv == 1) {
-            PyObject_SetAttrString(tp, *key_iter, value);
-            Py_DECREF(value);
-        } else if (rv == -1) {
-          goto fail;
+    const char *dataclass_type_ptr = PyUnicode_AsUTF8(dataclass_type);
+    if (dataclass_type_ptr == NULL) {
+        goto fail;
+    }
+    if (strcmp(dataclass_type_ptr, "attr") == 0 ||
+        strcmp(dataclass_type_ptr, "attr-auto") == 0) {
+        // These attributes are added or modified by @attr.s(slots=True).
+        const char * const keys[] = {"__attrs_attrs__", "__attrs_own_setattr__", "__init__", ""};
+        for (const char * const *key_iter = keys; **key_iter != '\0'; key_iter++) {
+            PyObject *value = NULL;
+            int rv = PyObject_GetOptionalAttrString(res, *key_iter, &value);
+            if (rv == 1) {
+                PyObject_SetAttrString(tp, *key_iter, value);
+                Py_DECREF(value);
+            } else if (rv == -1) {
+              goto fail;
+            }
         }
     }
 

mypyc/primitives/misc_ops.py

@@ -224,7 +224,13 @@
 # Create a dataclass from an extension class. See
 # CPyDataclass_SleightOfHand for more docs.
 dataclass_sleight_of_hand = custom_op(
-    arg_types=[object_rprimitive, object_rprimitive, dict_rprimitive, dict_rprimitive],
+    arg_types=[
+        object_rprimitive,
+        object_rprimitive,
+        dict_rprimitive,
+        dict_rprimitive,
+        str_rprimitive,
+    ],
     return_type=bit_rprimitive,
     c_function_name="CPyDataclass_SleightOfHand",
     error_kind=ERR_FALSE,