Make sure Pydantic dataclasses with slots and validate_assignment are unpickable

Viicos · Viicos · commit 63c9727a2d25 · 2025-04-17T14:53:45.000+02:00
diff --git a/pydantic/_internal/_dataclasses.py b/pydantic/_internal/_dataclasses.py
@@ -85,6 +85,7 @@ def complete_dataclass(
     cls: type[Any],
     config_wrapper: _config.ConfigWrapper,
     *,
+    slots: bool,
     raise_errors: bool = True,
     ns_resolver: NsResolver | None = None,
     _force_build: bool = False,
@@ -98,6 +99,7 @@ def complete_dataclass(
     Args:
         cls: The class.
         config_wrapper: The config wrapper instance.
+        slots: Whether slots was set on the class.
         raise_errors: Whether to raise errors, defaults to `True`.
         ns_resolver: The namespace resolver instance to use when collecting dataclass fields
             and during schema building.
@@ -194,6 +196,27 @@ def validated_setattr(instance: Any, field: str, value: str, /) -> None:
 
         cls.__setattr__ = validated_setattr.__get__(None, cls)  # type: ignore
 
+        if slots and not hasattr(cls, '__setstate__'):
+            # If slots is set, `pickle` (relied on by `copy.copy()`) will use
+            # `__setattr__()` to reconstruct the dataclass. However, the custom
+            # `__setattr__()` set above relies on `validate_assignment()`, which
+            # in turn excepts all the field values to be already present on the
+            # instance, resulting in attribute errors.
+            # As such, we make use of `object.__setattr__()` instead.
+            # Note that we do so only if `__setstate__()` isn't already set (this is the
+            # case if on top of `slots`, `frozen` is used).
+
+            # Taken from `dataclasses._dataclass_get/setstate()`:
+            def _dataclass_getstate(self: Any) -> list[Any]:
+                return [getattr(self, f.name) for f in dataclasses.fields(self)]
+
+            def _dataclass_setstate(self: Any, state: list[Any]) -> None:
+                for field, value in zip(dataclasses.fields(self), state):
+                    object.__setattr__(self, field.name, value)
+
+            cls.__getstate__ = _dataclass_getstate  # pyright: ignore[reportAttributeAccessIssue]
+            cls.__setstate__ = _dataclass_setstate
+
     cls.__pydantic_complete__ = True
     return True
 
diff --git a/pydantic/dataclasses.py b/pydantic/dataclasses.py
@@ -276,7 +276,7 @@ def create_dataclass(cls: type[Any]) -> type[PydanticDataclass]:
         # TODO `parent_namespace` is currently None, but we could do the same thing as Pydantic models:
         # fetch the parent ns using `parent_frame_namespace` (if the dataclass was defined in a function),
         # and possibly cache it (see the `__pydantic_parent_namespace__` logic for models).
-        _pydantic_dataclasses.complete_dataclass(cls, config_wrapper, raise_errors=False)
+        _pydantic_dataclasses.complete_dataclass(cls, config_wrapper, slots=slots, raise_errors=False)
         return cls
 
     return create_dataclass if _cls is None else create_dataclass(_cls)
diff --git a/tests/test_dataclasses.py b/tests/test_dataclasses.py
@@ -2430,6 +2430,23 @@ class Model:
     assert dc.b == 'bar'
 
 
+# Must be defined at the module level to be pickable:
+@pydantic.dataclasses.dataclass(slots=True, config={'validate_assignment': True})
+class DataclassSlotsValidateAssignment:
+    a: int
+
+
+@pytest.mark.skipif(sys.version_info < (3, 10), reason='slots are only supported for dataclasses in Python >= 3.10')
+def test_dataclass_slots_validate_assignment():
+    """https://github.com/pydantic/pydantic/issues/11768"""
+
+    m = DataclassSlotsValidateAssignment(1)
+    m_pickle = pickle.loads(pickle.dumps(m))
+    assert m_pickle.a == 1
+    with pytest.raises(ValidationError):
+        m.a = 'not_an_int'
+
+
 @pytest.mark.parametrize(
     'dataclass_decorator',
     [
