fix dataclass validation & serialization

Author: samuelcolvin

pydantic_core/core_schema.py

@@ -3123,7 +3123,7 @@ def dataclass_schema(
     another schema, not as the root type.
 
     Args:
-        cls: The dataclass type, used to to perform subclass checks
+        cls: The dataclass type, used to perform subclass checks
         schema: The schema to use for the dataclass fields
         cls_name: The name to use in error locs, etc; this is useful for generics (default: `cls.__name__`)
         post_init: Whether to call `__post_init__` after validation

src/serializers/infer.rs

@@ -21,7 +21,7 @@ use super::errors::{py_err_se_err, PydanticSerializationError};
 use super::extra::{Extra, SerMode};
 use super::filter::AnyFilter;
 use super::ob_type::ObType;
-use super::shared::object_to_dict;
+use super::shared::dataclass_to_dict;
 
 pub(crate) fn infer_to_python(
     value: &PyAny,
@@ -97,29 +97,23 @@ pub(crate) fn infer_to_python_known(
         Ok::<PyObject, PyErr>(new_dict.into_py(py))
     };
 
-    let serialize_with_serializer = |value: &PyAny, is_model: bool| {
-        if let Ok(py_serializer) = value.getattr(intern!(py, "__pydantic_serializer__")) {
-            if let Ok(serializer) = py_serializer.extract::<SchemaSerializer>() {
-                let extra = serializer.build_extra(
-                    py,
-                    extra.mode,
-                    extra.by_alias,
-                    extra.warnings,
-                    extra.exclude_unset,
-                    extra.exclude_defaults,
-                    extra.exclude_none,
-                    extra.round_trip,
-                    extra.rec_guard,
-                    extra.serialize_unknown,
-                    extra.fallback,
-                );
-                return serializer.serializer.to_python(value, include, exclude, &extra);
-            }
-        }
-        // Fallback to dict serialization if `__pydantic_serializer__` is not set.
-        // This currently only affects non-pydantic dataclasses.
-        let dict = object_to_dict(value, is_model, extra)?;
-        serialize_dict(dict)
+    let serialize_with_serializer = || {
+        let py_serializer = value.getattr(intern!(py, "__pydantic_serializer__"))?;
+        let serializer: SchemaSerializer = py_serializer.extract()?;
+        let extra = serializer.build_extra(
+            py,
+            extra.mode,
+            extra.by_alias,
+            extra.warnings,
+            extra.exclude_unset,
+            extra.exclude_defaults,
+            extra.exclude_none,
+            extra.round_trip,
+            extra.rec_guard,
+            extra.serialize_unknown,
+            extra.fallback,
+        );
+        serializer.serializer.to_python(value, include, exclude, &extra)
     };
 
     let value = match extra.mode {
@@ -191,8 +185,8 @@ pub(crate) fn infer_to_python_known(
                 let py_url: PyMultiHostUrl = value.extract()?;
                 py_url.__str__().into_py(py)
             }
-            ObType::PydanticSerializable => serialize_with_serializer(value, true)?,
-            ObType::Dataclass => serialize_with_serializer(value, false)?,
+            ObType::PydanticSerializable => serialize_with_serializer()?,
+            ObType::Dataclass => serialize_dict(dataclass_to_dict(value)?)?,
             ObType::Enum => {
                 let v = value.getattr(intern!(py, "value"))?;
                 infer_to_python(v, include, exclude, extra)?.into_py(py)
@@ -257,8 +251,8 @@ pub(crate) fn infer_to_python_known(
                 }
                 new_dict.into_py(py)
             }
-            ObType::PydanticSerializable => serialize_with_serializer(value, true)?,
-            ObType::Dataclass => serialize_with_serializer(value, false)?,
+            ObType::PydanticSerializable => serialize_with_serializer()?,
+            ObType::Dataclass => serialize_dict(dataclass_to_dict(value)?)?,
             ObType::Generator => {
                 let iter = super::type_serializers::generator::SerializationIterator::new(
                     value.downcast()?,
@@ -406,36 +400,6 @@ pub(crate) fn infer_serialize_known<S: Serializer>(
         }};
     }
 
-    macro_rules! serialize_with_serializer {
-        ($py_serializable:expr, $is_model:expr) => {{
-            let py = $py_serializable.py();
-            if let Ok(py_serializer) = value.getattr(intern!(py, "__pydantic_serializer__")) {
-                if let Ok(extracted_serializer) = py_serializer.extract::<SchemaSerializer>() {
-                    let extra = extracted_serializer.build_extra(
-                        py,
-                        extra.mode,
-                        extra.by_alias,
-                        extra.warnings,
-                        extra.exclude_unset,
-                        extra.exclude_defaults,
-                        extra.exclude_none,
-                        extra.round_trip,
-                        extra.rec_guard,
-                        extra.serialize_unknown,
-                        extra.fallback,
-                    );
-                    let pydantic_serializer =
-                        PydanticSerializer::new(value, &extracted_serializer.serializer, include, exclude, &extra);
-                    return pydantic_serializer.serialize(serializer);
-                }
-            }
-            // Fallback to dict serialization if `__pydantic_serializer__` is not set.
-            // This currently only affects non-pydantic dataclasses.
-            let dict = object_to_dict(value, $is_model, extra).map_err(py_err_se_err)?;
-            serialize_dict!(dict)
-        }};
-    }
-
     let ser_result = match ob_type {
         ObType::None => serializer.serialize_none(),
         ObType::Int | ObType::IntSubclass => serialize!(i64),
@@ -490,8 +454,30 @@ pub(crate) fn infer_serialize_known<S: Serializer>(
             let py_url: PyMultiHostUrl = value.extract().map_err(py_err_se_err)?;
             serializer.serialize_str(&py_url.__str__())
         }
-        ObType::Dataclass => serialize_with_serializer!(value, false),
-        ObType::PydanticSerializable => serialize_with_serializer!(value, true),
+        ObType::PydanticSerializable => {
+            let py = value.py();
+            let py_serializer = value
+                .getattr(intern!(py, "__pydantic_serializer__"))
+                .map_err(py_err_se_err)?;
+            let extracted_serializer: SchemaSerializer = py_serializer.extract().map_err(py_err_se_err)?;
+            let extra = extracted_serializer.build_extra(
+                py,
+                extra.mode,
+                extra.by_alias,
+                extra.warnings,
+                extra.exclude_unset,
+                extra.exclude_defaults,
+                extra.exclude_none,
+                extra.round_trip,
+                extra.rec_guard,
+                extra.serialize_unknown,
+                extra.fallback,
+            );
+            let pydantic_serializer =
+                PydanticSerializer::new(value, &extracted_serializer.serializer, include, exclude, &extra);
+            pydantic_serializer.serialize(serializer)
+        }
+        ObType::Dataclass => serialize_dict!(dataclass_to_dict(value).map_err(py_err_se_err)?),
         ObType::Enum => {
             let v = value.getattr(intern!(value.py(), "value")).map_err(py_err_se_err)?;
             infer_serialize(v, serializer, include, exclude, extra)

src/serializers/mod.rs

@@ -11,7 +11,7 @@ use config::SerializationConfig;
 pub use errors::{PydanticSerializationError, PydanticSerializationUnexpectedValue};
 use extra::{CollectWarnings, SerRecursionGuard};
 pub(crate) use extra::{Extra, SerMode, SerializationState};
-pub(crate) use shared::slots_dc_dict;
+pub(crate) use shared::dataclass_to_dict;
 pub use shared::CombinedSerializer;
 use shared::{to_json_bytes, BuildSerializer, TypeSerializer};
 

src/serializers/shared.rs

@@ -4,7 +4,7 @@ use std::fmt::Debug;
 use pyo3::exceptions::PyTypeError;
 use pyo3::once_cell::GILOnceCell;
 use pyo3::prelude::*;
-use pyo3::types::{PyDict, PySet, PyString};
+use pyo3::types::{PyDict, PyString};
 use pyo3::{intern, PyTraverseError, PyVisit};
 
 use enum_dispatch::enum_dispatch;
@@ -96,7 +96,6 @@ combined_serializer! {
         super::type_serializers::other::CallableBuilder;
         super::type_serializers::definitions::DefinitionsSerializerBuilder;
         super::type_serializers::dataclass::DataclassArgsBuilder;
-        super::type_serializers::dataclass::DataclassBuilder;
         super::type_serializers::function::FunctionBeforeSerializerBuilder;
         super::type_serializers::function::FunctionAfterSerializerBuilder;
         super::type_serializers::function::FunctionPlainSerializerBuilder;
@@ -124,6 +123,7 @@ combined_serializer! {
         Generator: super::type_serializers::generator::GeneratorSerializer;
         Dict: super::type_serializers::dict::DictSerializer;
         Model: super::type_serializers::model::ModelSerializer;
+        Dataclass: super::type_serializers::dataclass::DataclassSerializer;
         Url: super::type_serializers::url::UrlSerializer;
         MultiHostUrl: super::type_serializers::url::MultiHostUrlSerializer;
         Any: super::type_serializers::any::AnySerializer;
@@ -328,42 +328,23 @@ pub(crate) fn to_json_bytes(
     Ok(bytes)
 }
 
-pub(super) fn object_to_dict<'py>(value: &'py PyAny, is_model: bool, extra: &Extra) -> PyResult<&'py PyDict> {
-    let py = value.py();
-    let attrs: &PyDict = match value.getattr(intern!(py, "__dict__")) {
-        Ok(attr) => attr.downcast()?,
-        Err(_) => return slots_dc_dict(value),
-    };
-
-    if is_model && extra.exclude_unset {
-        let fields_set: &PySet = value.getattr(intern!(py, "__pydantic_fields_set__"))?.downcast()?;
-
-        let new_attrs = attrs.copy()?;
-        for key in new_attrs.keys() {
-            if !fields_set.contains(key)? {
-                new_attrs.del_item(key)?;
-            }
-        }
-        Ok(new_attrs)
-    } else {
-        Ok(attrs)
-    }
-}
-
 static DC_FIELD_MARKER: GILOnceCell<PyObject> = GILOnceCell::new();
 
-pub(crate) fn slots_dc_dict(dc: &PyAny) -> PyResult<&PyDict> {
-    let py = dc.py();
-    let dc_fields: &PyDict = dc.getattr(intern!(py, "__dataclass_fields__"))?.downcast()?;
-    let dict = PyDict::new(py);
-
-    // need to match the logic from dataclasses.fields `tuple(f for f in fields.values() if f._field_type is _FIELD)`
+/// needed to match the logic from dataclasses.fields `tuple(f for f in fields.values() if f._field_type is _FIELD)`
+pub(super) fn get_field_marker(py: Python<'_>) -> PyResult<&PyAny> {
     let field_type_marker_obj = DC_FIELD_MARKER.get_or_try_init(py, || {
         let field_ = py.import("dataclasses")?.getattr("_FIELD")?;
         Ok::<PyObject, PyErr>(field_.into_py(py))
     })?;
-    let field_type_marker = field_type_marker_obj.as_ref(py);
+    Ok(field_type_marker_obj.as_ref(py))
+}
+
+pub(crate) fn dataclass_to_dict(dc: &PyAny) -> PyResult<&PyDict> {
+    let py = dc.py();
+    let dc_fields: &PyDict = dc.getattr(intern!(py, "__dataclass_fields__"))?.downcast()?;
+    let dict = PyDict::new(py);
 
+    let field_type_marker = get_field_marker(py)?;
     for (field_name, field) in dc_fields.iter() {
         let field_type = field.getattr(intern!(py, "_field_type"))?;
         if field_type.is(field_type_marker) {