add dataclass.fields to schema

Author: samuelcolvin

pydantic_core/core_schema.py

@@ -3093,6 +3093,7 @@ class DataclassSchema(TypedDict, total=False):
     type: Required[Literal['dataclass']]
     cls: Required[Type[Any]]
     schema: Required[CoreSchema]
+    fields: Required[List[str]]
     cls_name: str
     post_init: bool  # default: False
     revalidate_instances: Literal['always', 'never', 'subclass-instances']  # default: 'never'
@@ -3101,12 +3102,13 @@ class DataclassSchema(TypedDict, total=False):
     ref: str
     metadata: Any
     serialization: SerSchema
-    slots: List[str]
+    slots: bool
 
 
 def dataclass_schema(
     cls: Type[Any],
     schema: CoreSchema,
+    fields: List[str],
     *,
     cls_name: str | None = None,
     post_init: bool | None = None,
@@ -3116,7 +3118,7 @@ def dataclass_schema(
     metadata: Any = None,
     serialization: SerSchema | None = None,
     frozen: bool | None = None,
-    slots: List[str] | None = None,
+    slots: bool | None = None,
 ) -> DataclassSchema:
     """
     Returns a schema for a dataclass. As with `ModelSchema`, this schema can only be used as a field within
@@ -3125,6 +3127,8 @@ def dataclass_schema(
     Args:
         cls: The dataclass type, used to perform subclass checks
         schema: The schema to use for the dataclass fields
+        fields: Fields of the dataclass, this is used in serialization and in validation during re-validation
+            and while validating assignment
         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
         revalidate_instances: whether instances of models and dataclasses (including subclass instances)
@@ -3134,11 +3138,13 @@ def dataclass_schema(
         metadata: Any other information you want to include with the schema, not used by pydantic-core
         serialization: Custom serialization schema
         frozen: Whether the dataclass is frozen
-        slots: The slots to use for the dataclass, set only if `slots=True` on the dataclass
+        slots: Whether `slots=True` on the dataclass, means each field is assigned independently, rather than
+            simply setting `__dict__`, default false
     """
     return dict_not_none(
         type='dataclass',
         cls=cls,
+        fields=fields,
         cls_name=cls_name,
         schema=schema,
         post_init=post_init,

src/serializers/mod.rs

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

src/serializers/shared.rs

@@ -339,7 +339,7 @@ pub(super) fn get_field_marker(py: Python<'_>) -> PyResult<&PyAny> {
     Ok(field_type_marker_obj.as_ref(py))
 }
 
-pub(crate) fn dataclass_to_dict(dc: &PyAny) -> PyResult<&PyDict> {
+pub(super) 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);

src/serializers/type_serializers/dataclass.rs

@@ -9,9 +9,9 @@ use crate::build_tools::{py_error_type, ExtraBehavior, SchemaDict};
 use crate::definitions::DefinitionsBuilder;
 
 use super::{
-    get_field_marker, infer_json_key, infer_json_key_known, infer_serialize, infer_to_python, py_err_se_err,
-    BuildSerializer, CombinedSerializer, ComputedFields, Extra, FieldsMode, GeneralFieldsSerializer, ObType, SerCheck,
-    SerField, TypeSerializer,
+    infer_json_key, infer_json_key_known, infer_serialize, infer_to_python, py_err_se_err, BuildSerializer,
+    CombinedSerializer, ComputedFields, Extra, FieldsMode, GeneralFieldsSerializer, ObType, SerCheck, SerField,
+    TypeSerializer,
 };
 
 pub struct DataclassArgsBuilder;
@@ -83,17 +83,11 @@ impl BuildSerializer for DataclassSerializer {
         let sub_schema: &PyDict = schema.get_as_req(intern!(py, "schema"))?;
         let serializer = Box::new(CombinedSerializer::build(sub_schema, config, definitions)?);
 
-        let dc_fields: &PyDict = class.getattr(intern!(py, "__dataclass_fields__"))?.downcast()?;
-        let mut fields = Vec::with_capacity(dc_fields.len());
-
-        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) {
-                let field_name: &PyString = field_name.downcast()?;
-                fields.push(field_name.into_py(py));
-            }
-        }
+        let fields = schema
+            .get_as_req::<&PyList>(intern!(py, "fields"))?
+            .iter()
+            .map(|s| Ok(s.downcast::<PyString>()?.into_py(py)))
+            .collect::<PyResult<Vec<_>>>()?;
 
         Ok(Self {
             class: class.into(),

src/serializers/type_serializers/mod.rs

@@ -35,6 +35,4 @@ pub(self) use super::infer::{
     infer_to_python_known,
 };
 pub(self) use super::ob_type::{IsType, ObType};
-pub(self) use super::shared::{
-    get_field_marker, to_json_bytes, BuildSerializer, CombinedSerializer, PydanticSerializer, TypeSerializer,
-};
+pub(self) use super::shared::{to_json_bytes, BuildSerializer, CombinedSerializer, PydanticSerializer, TypeSerializer};

src/validators/dataclass.rs

@@ -10,7 +10,6 @@ use crate::errors::{ErrorType, ValError, ValLineError, ValResult};
 use crate::input::{GenericArguments, Input};
 use crate::lookup_key::LookupKey;
 use crate::recursion_guard::RecursionGuard;
-use crate::serializers::dataclass_to_dict;
 use crate::validators::function::convert_err;
 
 use super::arguments::{json_get, json_slice, py_get, py_slice};
@@ -411,11 +410,12 @@ pub struct DataclassValidator {
     strict: bool,
     validator: Box<CombinedValidator>,
     class: Py<PyType>,
+    fields: Vec<Py<PyString>>,
     post_init: Option<Py<PyString>>,
     revalidate: Revalidate,
     name: String,
     frozen: bool,
-    slots: Option<Vec<Py<PyString>>>,
+    slots: bool,
 }
 
 impl BuildValidator for DataclassValidator {
@@ -442,21 +442,17 @@ impl BuildValidator for DataclassValidator {
             None
         };
 
-        let slots = match schema.get_as::<&PyList>(intern!(py, "slots"))? {
-            Some(slots) => {
-                let slots = slots
-                    .iter()
-                    .map(|s| Ok(s.downcast::<PyString>()?.into_py(py)))
-                    .collect::<PyResult<Vec<_>>>()?;
-                Some(slots)
-            }
-            None => None,
-        };
+        let fields = schema
+            .get_as_req::<&PyList>(intern!(py, "fields"))?
+            .iter()
+            .map(|s| Ok(s.downcast::<PyString>()?.into_py(py)))
+            .collect::<PyResult<Vec<_>>>()?;
 
         Ok(Self {
             strict: is_strict(schema, config)?,
             validator: Box::new(validator),
             class: class.into(),
+            fields,
             post_init,
             revalidate: Revalidate::from_str(schema_or_config_same(
                 schema,
@@ -465,7 +461,7 @@ impl BuildValidator for DataclassValidator {
             )?)?,
             name,
             frozen: schema.get_as(intern!(py, "frozen"))?.unwrap_or(false),
-            slots,
+            slots: schema.get_as(intern!(py, "slots"))?.unwrap_or(false),
         }
         .into())
     }
@@ -489,7 +485,7 @@ impl Validator for DataclassValidator {
         let class = self.class.as_ref(py);
         if let Some(py_input) = input.input_is_instance(class) {
             if self.revalidate.should_revalidate(py_input, class) {
-                let input_dict: &PyAny = dataclass_to_dict(py_input)?;
+                let input_dict: &PyAny = self.dataclass_to_dict(py, py_input)?;
                 let val_output = self
                     .validator
                     .validate(py, input_dict, extra, definitions, recursion_guard)?;
@@ -530,17 +526,7 @@ impl Validator for DataclassValidator {
             return Err(ValError::new(ErrorType::FrozenInstance, field_value));
         }
 
-        let new_dict = if let Some(ref slots) = self.slots {
-            let slots_dict = PyDict::new(py);
-            for slot in slots {
-                let slot = slot.as_ref(py);
-                slots_dict.set_item(slot, obj.getattr(slot)?)?;
-            }
-            slots_dict
-        } else {
-            let dunder_dict: &PyDict = obj.getattr(intern!(py, "__dict__"))?.downcast()?;
-            dunder_dict.copy()?
-        };
+        let new_dict = self.dataclass_to_dict(py, obj)?;
 
         new_dict.set_item(field_name, field_value)?;
 
@@ -556,8 +542,10 @@ impl Validator for DataclassValidator {
 
         let (dc_dict, _): (&PyDict, PyObject) = val_assignment_result.extract(py)?;
 
-        if self.slots.is_some() {
-            let value = dc_dict.get_item(field_name).unwrap();
+        if self.slots {
+            let value = dc_dict
+                .get_item(field_name)
+                .ok_or_else(|| PyKeyError::new_err(field_name.to_string()))?;
             force_setattr(py, obj, field_name, value)?;
         } else {
             force_setattr(py, obj, intern!(py, "__dict__"), dc_dict)?;
@@ -613,6 +601,16 @@ impl DataclassValidator {
         Ok(self_instance.into_py(py))
     }
 
+    fn dataclass_to_dict<'py>(&self, py: Python<'py>, dc: &'py PyAny) -> PyResult<&'py PyDict> {
+        let dict = PyDict::new(py);
+
+        for field_name in &self.fields {
+            let field_name = field_name.as_ref(py);
+            dict.set_item(field_name, dc.getattr(field_name)?)?;
+        }
+        Ok(dict)
+    }
+
     fn set_dict_call<'s, 'data>(
         &'s self,
         py: Python<'data>,
@@ -621,7 +619,7 @@ impl DataclassValidator {
         input: &'data impl Input<'data>,
     ) -> ValResult<'data, ()> {
         let (dc_dict, post_init_kwargs): (&PyAny, &PyAny) = val_output.extract(py)?;
-        if self.slots.is_some() {
+        if self.slots {
             let dc_dict: &PyDict = dc_dict.downcast()?;
             for (key, value) in dc_dict.iter() {
                 force_setattr(py, dc, key, value)?;

tests/serializers/test_any.py

@@ -439,6 +439,7 @@ class Foo:
         core_schema.dataclass_args_schema(
             'Foo', [core_schema.dataclass_field(name='a', schema=core_schema.str_schema())]
         ),
+        ['a'],
     )
     Foo.__pydantic_serializer__ = SchemaSerializer(schema)
 
@@ -467,6 +468,7 @@ class Foo:
         core_schema.dataclass_args_schema(
             'Foo', [core_schema.dataclass_field(name='a', schema=core_schema.str_schema())]
         ),
+        ['a'],
     )
     Foo.__pydantic_validator__ = SchemaValidator(schema)
     Foo.__pydantic_serializer__ = SchemaSerializer(schema)

tests/serializers/test_dataclasses.py

@@ -32,6 +32,7 @@ def test_dataclass():
                 core_schema.dataclass_field(name='b', schema=core_schema.bytes_schema()),
             ],
         ),
+        ['a', 'b'],
     )
     s = SchemaSerializer(schema)
     assert s.to_python(Foo(a='hello', b=b'more')) == IsStrictDict(a='hello', b=b'more')
@@ -57,6 +58,7 @@ def test_serialization_exclude():
                 core_schema.dataclass_field(name='b', schema=core_schema.bytes_schema(), serialization_exclude=True),
             ],
         ),
+        ['a', 'b'],
     )
     s = SchemaSerializer(schema)
     assert s.to_python(Foo(a='hello', b=b'more')) == {'a': 'hello'}
@@ -79,6 +81,7 @@ def test_serialization_alias():
                 core_schema.dataclass_field(name='b', schema=core_schema.bytes_schema(), serialization_alias='BAR'),
             ],
         ),
+        ['a', 'b'],
     )
     s = SchemaSerializer(schema)
     assert s.to_python(Foo(a='hello', b=b'more')) == IsStrictDict(a='hello', BAR=b'more')
@@ -111,6 +114,7 @@ def c(self) -> str:
             ],
             computed_fields=[core_schema.computed_field('c', core_schema.str_schema())],
         ),
+        ['a', 'b'],
     )
     s = SchemaSerializer(schema)
     assert s.to_python(FooProp(a='hello', b=b'more')) == IsStrictDict(a='hello', b=b'more', c='hello more')
@@ -151,7 +155,8 @@ class SubModel(Model):
                 core_schema.dataclass_field(name='y2', init_only=True, schema=core_schema.str_schema()),
             ],
         ),
-        slots=['x'],
+        ['x', 'x2'],
+        slots=True,
     )
     dc = SubModel(x=1, y='a', x2=2, y2='b')
     assert dataclasses.asdict(dc) == {'x': 1, 'x2': 2}

tests/test_schema_functions.py

@@ -255,13 +255,13 @@ def args(*args, **kwargs):
     ),
     (
         core_schema.dataclass_schema,
-        args(MyDataclass, {'type': 'int'}),
-        {'type': 'dataclass', 'schema': {'type': 'int'}, 'cls': MyDataclass},
+        args(MyDataclass, {'type': 'int'}, ['foobar']),
+        {'type': 'dataclass', 'schema': {'type': 'int'}, 'fields': ['foobar'], 'cls': MyDataclass},
     ),
     (
         core_schema.dataclass_schema,
-        args(MyDataclass, {'type': 'int'}, slots=['a']),
-        {'type': 'dataclass', 'schema': {'type': 'int'}, 'cls': MyDataclass, 'slots': ['a']},
+        args(MyDataclass, {'type': 'int'}, ['foobar'], slots=True),
+        {'type': 'dataclass', 'schema': {'type': 'int'}, 'fields': ['foobar'], 'cls': MyDataclass, 'slots': True},
     ),
 ]