Implement frozen and extra_behavior for dataclasses (#505)

Author: adriangb

pydantic_core/core_schema.py

@@ -19,6 +19,9 @@ def dict_not_none(**kwargs: Any) -> Any:
     return {k: v for k, v in kwargs.items() if v is not None}
 
 
+ExtraBehavior = Literal['allow', 'forbid', 'ignore']
+
+
 class CoreConfig(TypedDict, total=False):
     title: str
     strict: bool
@@ -27,7 +30,7 @@ class CoreConfig(TypedDict, total=False):
     # if configs are merged, which should take precedence, default 0, default means child takes precedence
     config_merge_priority: int
     # settings related to typed_dicts only
-    typed_dict_extra_behavior: Literal['allow', 'forbid', 'ignore']
+    extra_fields_behavior: ExtraBehavior
     typed_dict_total: bool  # default: True
     # used on typed-dicts and tagged union keys
     from_attributes: bool
@@ -2494,7 +2497,7 @@ class TypedDictSchema(TypedDict, total=False):
     extra_validator: CoreSchema
     return_fields_set: bool
     # all these values can be set via config, equivalent fields have `typed_dict_` prefix
-    extra_behavior: Literal['allow', 'forbid', 'ignore']
+    extra_behavior: ExtraBehavior
     total: bool  # default: True
     populate_by_name: bool  # replaces `allow_population_by_field_name` in pydantic v1
     from_attributes: bool
@@ -2509,7 +2512,7 @@ def typed_dict_schema(
     strict: bool | None = None,
     extra_validator: CoreSchema | None = None,
     return_fields_set: bool | None = None,
-    extra_behavior: Literal['allow', 'forbid', 'ignore'] | None = None,
+    extra_behavior: ExtraBehavior | None = None,
     total: bool | None = None,
     populate_by_name: bool | None = None,
     from_attributes: bool | None = None,
@@ -2645,6 +2648,7 @@ class DataclassField(TypedDict, total=False):
     schema: Required[CoreSchema]
     kw_only: bool  # default: True
     init_only: bool  # default: False
+    frozen: bool  # default: False
     validation_alias: Union[str, List[Union[str, int]], List[List[Union[str, int]]]]
     serialization_alias: str
     serialization_exclude: bool  # default: False
@@ -2661,6 +2665,7 @@ def dataclass_field(
     serialization_alias: str | None = None,
     serialization_exclude: bool | None = None,
     metadata: Any = None,
+    frozen: bool | None = None,
 ) -> DataclassField:
     """
     Returns a schema for a dataclass field, e.g.:
@@ -2696,6 +2701,7 @@ def dataclass_field(
         serialization_alias=serialization_alias,
         serialization_exclude=serialization_exclude,
         metadata=metadata,
+        frozen=frozen,
     )
 
 
@@ -2708,6 +2714,7 @@ class DataclassArgsSchema(TypedDict, total=False):
     ref: str
     metadata: Any
     serialization: SerSchema
+    extra_behavior: ExtraBehavior
 
 
 def dataclass_args_schema(
@@ -2718,6 +2725,7 @@ def dataclass_args_schema(
     ref: str | None = None,
     metadata: Any = None,
     serialization: SerSchema | None = None,
+    extra_behavior: ExtraBehavior | None = None,
 ) -> DataclassArgsSchema:
     """
     Returns a schema for validating dataclass arguments, e.g.:
@@ -2754,6 +2762,7 @@ def dataclass_args_schema(
         ref=ref,
         metadata=metadata,
         serialization=serialization,
+        extra_behavior=extra_behavior,
     )
 
 
@@ -2764,6 +2773,7 @@ class DataclassSchema(TypedDict, total=False):
     post_init: bool  # default: False
     revalidate_instances: Literal['always', 'never', 'subclass-instances']  # default: 'never'
     strict: bool  # default: False
+    frozen: bool  # default False
     ref: str
     metadata: Any
     serialization: SerSchema
@@ -2779,6 +2789,7 @@ def dataclass_schema(
     ref: str | None = None,
     metadata: Any = None,
     serialization: SerSchema | None = None,
+    frozen: bool | None = None,
 ) -> DataclassSchema:
     """
     Returns a schema for a dataclass. As with `ModelSchema`, this schema can only be used as a field within
@@ -2805,6 +2816,7 @@ def dataclass_schema(
         ref=ref,
         metadata=metadata,
         serialization=serialization,
+        frozen=frozen,
     )
 
 

src/build_tools.rs

@@ -248,3 +248,35 @@ pub fn safe_repr(v: &PyAny) -> Cow<str> {
         "<unprintable object>".into()
     }
 }
+
+#[derive(Debug, Clone)]
+pub(crate) enum ExtraBehavior {
+    Allow,
+    Forbid,
+    Ignore,
+}
+
+impl ExtraBehavior {
+    pub fn from_schema_or_config(
+        py: Python,
+        schema: &PyDict,
+        config: Option<&PyDict>,
+        default: Self,
+    ) -> PyResult<Self> {
+        let extra_behavior = schema_or_config::<Option<&str>>(
+            schema,
+            config,
+            intern!(py, "extra_behavior"),
+            intern!(py, "extra_fields_behavior"),
+        )?
+        .flatten();
+        let res = match extra_behavior {
+            Some("allow") => Self::Allow,
+            Some("ignore") => Self::Ignore,
+            Some("forbid") => Self::Forbid,
+            Some(v) => return py_err!("Invalid extra_behavior: `{}`", v),
+            None => default,
+        };
+        Ok(res)
+    }
+}

src/serializers/type_serializers/typed_dict.rs

@@ -8,7 +8,7 @@ use ahash::{AHashMap, AHashSet};
 use serde::ser::SerializeMap;
 
 use crate::build_context::BuildContext;
-use crate::build_tools::{py_error_type, schema_or_config, SchemaDict};
+use crate::build_tools::{py_error_type, schema_or_config, ExtraBehavior, SchemaDict};
 use crate::PydanticSerializationUnexpectedValue;
 
 use super::{
@@ -80,16 +80,13 @@ impl BuildSerializer for TypedDictSerializer {
     ) -> PyResult<CombinedSerializer> {
         let py = schema.py();
 
-        let extra_behavior = schema_or_config::<&str>(
-            schema,
-            config,
-            intern!(py, "extra_behavior"),
-            intern!(py, "typed_dict_extra_behavior"),
-        )?;
         let total =
             schema_or_config(schema, config, intern!(py, "total"), intern!(py, "typed_dict_total"))?.unwrap_or(true);
 
-        let include_extra = extra_behavior == Some("allow");
+        let include_extra = matches!(
+            ExtraBehavior::from_schema_or_config(py, schema, config, ExtraBehavior::Ignore)?,
+            ExtraBehavior::Allow
+        );
 
         let fields_dict: &PyDict = schema.get_as_req(intern!(py, "fields"))?;
         let mut fields: AHashMap<String, TypedDictField> = AHashMap::with_capacity(fields_dict.len());

src/validators/dataclass.rs

@@ -5,7 +5,7 @@ use pyo3::types::{PyDict, PyList, PyString, PyTuple, PyType};
 
 use ahash::AHashSet;
 
-use crate::build_tools::{is_strict, py_err, schema_or_config_same, SchemaDict};
+use crate::build_tools::{is_strict, py_err, schema_or_config_same, ExtraBehavior, SchemaDict};
 use crate::errors::{ErrorType, ValError, ValLineError, ValResult};
 use crate::input::{GenericArguments, Input};
 use crate::lookup_key::LookupKey;
@@ -24,6 +24,7 @@ struct Field {
     init_only: bool,
     lookup_key: LookupKey,
     validator: CombinedValidator,
+    frozen: bool,
 }
 
 #[derive(Debug, Clone)]
@@ -33,6 +34,7 @@ pub struct DataclassArgsValidator {
     init_only_count: Option<usize>,
     dataclass_name: String,
     validator_name: String,
+    extra_behavior: ExtraBehavior,
 }
 
 impl BuildValidator for DataclassArgsValidator {
@@ -47,6 +49,8 @@ impl BuildValidator for DataclassArgsValidator {
 
         let populate_by_name = schema_or_config_same(schema, config, intern!(py, "populate_by_name"))?.unwrap_or(false);
 
+        let extra_behavior = ExtraBehavior::from_schema_or_config(py, schema, config, ExtraBehavior::Ignore)?;
+
         let fields_schema: &PyList = schema.get_as_req(intern!(py, "fields"))?;
         let mut fields: Vec<Field> = Vec::with_capacity(fields_schema.len());
 
@@ -91,6 +95,7 @@ impl BuildValidator for DataclassArgsValidator {
                 lookup_key,
                 validator,
                 init_only: field.get_as(intern!(py, "init_only"))?.unwrap_or(false),
+                frozen: field.get_as::<bool>(intern!(py, "frozen"))?.unwrap_or(false),
             });
         }
 
@@ -108,6 +113,7 @@ impl BuildValidator for DataclassArgsValidator {
             init_only_count,
             dataclass_name,
             validator_name,
+            extra_behavior,
         }
         .into())
     }
@@ -254,11 +260,20 @@ impl Validator for DataclassArgsValidator {
                             match raw_key.strict_str() {
                                 Ok(either_str) => {
                                     if !used_keys.contains(either_str.as_cow()?.as_ref()) {
-                                        errors.push(ValLineError::new_with_loc(
-                                            ErrorType::UnexpectedKeywordArgument,
-                                            value,
-                                            raw_key.as_loc_item(),
-                                        ));
+                                        // Unknown / extra field
+                                        match self.extra_behavior {
+                                            ExtraBehavior::Forbid => {
+                                                errors.push(ValLineError::new_with_loc(
+                                                    ErrorType::UnexpectedKeywordArgument,
+                                                    value,
+                                                    raw_key.as_loc_item(),
+                                                ));
+                                            }
+                                            ExtraBehavior::Ignore => {}
+                                            ExtraBehavior::Allow => {
+                                                output_dict.set_item(either_str.as_py_string(py), value)?
+                                            }
+                                        }
                                     }
                                 }
                                 Err(ValError::LineErrors(line_errors)) => {
@@ -303,7 +318,19 @@ impl Validator for DataclassArgsValidator {
     ) -> ValResult<'data, PyObject> {
         let dict: &PyDict = obj.downcast()?;
 
+        let ok = |output: PyObject| {
+            dict.set_item(field_name, output)?;
+            Ok(dict.to_object(py))
+        };
+
         if let Some(field) = self.fields.iter().find(|f| f.name == field_name) {
+            if field.frozen {
+                return Err(ValError::new_with_loc(
+                    ErrorType::FrozenField,
+                    field_value,
+                    field.name.to_string(),
+                ));
+            }
             // by using dict but removing the field in question, we match V1 behaviour
             let data_dict = dict.copy()?;
             if let Err(err) = data_dict.del_item(field_name) {
@@ -321,10 +348,7 @@ impl Validator for DataclassArgsValidator {
                 .validator
                 .validate(py, field_value, &next_extra, slots, recursion_guard)
             {
-                Ok(output) => {
-                    dict.set_item(field_name, output)?;
-                    Ok(dict.to_object(py))
-                }
+                Ok(output) => ok(output),
                 Err(ValError::LineErrors(line_errors)) => {
                     let errors = line_errors
                         .into_iter()
@@ -335,13 +359,21 @@ impl Validator for DataclassArgsValidator {
                 Err(err) => Err(err),
             }
         } else {
-            Err(ValError::new_with_loc(
-                ErrorType::NoSuchAttribute {
-                    attribute: field_name.to_string(),
-                },
-                field_value,
-                field_name.to_string(),
-            ))
+            // Handle extra (unknown) field
+            // We partially use the extra_behavior for initialization / validation
+            // to determine how to handle assignment
+            match self.extra_behavior {
+                // For dataclasses we allow assigning unknown fields
+                // to match stdlib dataclass behavior
+                ExtraBehavior::Allow => ok(field_value.to_object(py)),
+                _ => Err(ValError::new_with_loc(
+                    ErrorType::NoSuchAttribute {
+                        attribute: field_name.to_string(),
+                    },
+                    field_value,
+                    field_name.to_string(),
+                )),
+            }
         }
     }
 
@@ -364,6 +396,7 @@ pub struct DataclassValidator {
     post_init: Option<Py<PyString>>,
     revalidate: Revalidate,
     name: String,
+    frozen: bool,
 }
 
 impl BuildValidator for DataclassValidator {
@@ -399,6 +432,7 @@ impl BuildValidator for DataclassValidator {
             // as with model, get the class's `__name__`, not using `class.name()` since it uses `__qualname__`
             // which is not what we want here
             name: class.getattr(intern!(py, "__name__"))?.extract()?,
+            frozen: schema.get_as(intern!(py, "frozen"))?.unwrap_or(false),
         }
         .into())
     }
@@ -455,6 +489,9 @@ impl Validator for DataclassValidator {
         slots: &'data [CombinedValidator],
         recursion_guard: &'s mut RecursionGuard,
     ) -> ValResult<'data, PyObject> {
+        if self.frozen {
+            return Err(ValError::new(ErrorType::FrozenInstance, field_value));
+        }
         let dict_py_str = intern!(py, "__dict__");
         let dict: &PyDict = obj.getattr(dict_py_str)?.downcast()?;