[mypyc] Fix pickling and copying frozen dataclasses

mypyc/lib-rt/misc_ops.c

@@ -475,8 +475,15 @@ CPyDataclass_SleightOfHand(PyObject *dataclass_dec, PyObject *tp,
 PyObject *
 CPyPickle_SetState(PyObject *obj, PyObject *state)
 {
-    if (_CPy_UpdateObjFromDict(obj, state) != 0) {
-        return NULL;
+    // Set the attributes like object.__setattr__ does, without calling any __setattr__ of the
+    // class, since pickle and copy restore the state of a Python object by updating its
+    // __dict__ directly. The __setattr__ of a frozen dataclass always raises, for example.
+    Py_ssize_t pos = 0;
+    PyObject *key, *value;
+    while (PyDict_Next(state, &pos, &key, &value)) {
+        if (PyObject_GenericSetAttr(obj, key, value) != 0) {
+            return NULL;
+        }
     }
     Py_RETURN_NONE;
 }

mypyc/test-data/run-classes.test

@@ -1731,6 +1731,74 @@ assert e.__getstate__() == {'y': 20, '__dict__': {'x': 10}}
 e2 = pickle.loads(pickle.dumps(e))
 assert e is not e2 and e.x == e2.x and e.y == e2.y
 
+[case testPickleFrozenDataclass]
+import copy
+import pickle
+from dataclasses import dataclass
+from typing import Any, Dict, List
+
+from mypy_extensions import mypyc_attr
+
+@dataclass(frozen=True)
+class Config:
+    py_version: str = "3"
+
+@dataclass(frozen=True)
+class Frozen:
+    x: int
+    y: str
+
+@dataclass(frozen=True)
+class FrozenSub(Frozen):
+    z: List[int]
+
+@dataclass(frozen=True, slots=True)
+class FrozenSlots:
+    a: int
+    b: Dict[int, int]
+
+setattr_calls: List[str] = []
+
+@mypyc_attr(serializable=True)
+class WithSetattr:
+    def __init__(self, x: int) -> None:
+        self.x = x
+
+    def __setattr__(self, name: str, value: object) -> None:
+        setattr_calls.append(name)
+        super().__setattr__(name, value)
+
+def round_trips(obj: Any) -> List[Any]:
+    # Protocols 0 and 1 can't rebuild any native class
+    out = [pickle.loads(pickle.dumps(obj, protocol))
+           for protocol in range(2, pickle.HIGHEST_PROTOCOL + 1)]
+    return out + [copy.copy(obj), copy.deepcopy(obj)]
+
+def test_frozen_dataclass() -> None:
+    for obj in Config(), Frozen(1, 'a'), FrozenSub(2, 'b', [3]), FrozenSlots(4, {5: 6}):
+        for copied in round_trips(obj):
+            assert copied == obj
+            assert copied is not obj
+            assert type(copied) is type(obj)
+
+def test_copy_is_still_frozen() -> None:
+    f: Any = copy.copy(Frozen(1, 'a'))
+    try:
+        f.x = 2
+    except AttributeError as e:
+        assert type(e).__name__ == 'FrozenInstanceError'
+    else:
+        assert False
+    assert f.x == 1
+
+def test_setattr_not_called() -> None:
+    # Python restores the state without calling __setattr__ too, by updating __dict__
+    obj = WithSetattr(1)
+    calls = len(setattr_calls)
+    for copied in round_trips(obj):
+        assert copied.x == 1
+    assert len(setattr_calls) == calls
+
 
 [case testInterpretedParentInit]
 from interp import C