improve field handling

Author: samuelcolvin

pydantic_core/core_schema.py

@@ -2841,6 +2841,7 @@ class ModelSchema(TypedDict, total=False):
     revalidate_instances: Literal['always', 'never', 'subclass-instances']  # default: 'never'
     strict: bool
     frozen: bool
+    extra_behavior: ExtraBehavior
     config: CoreConfig
     ref: str
     metadata: Any
@@ -2855,6 +2856,7 @@ def model_schema(
     revalidate_instances: Literal['always', 'never', 'subclass-instances'] | None = None,
     strict: bool | None = None,
     frozen: bool | None = None,
+    extra_behavior: ExtraBehavior | None = None,
     config: CoreConfig | None = None,
     ref: str | None = None,
     metadata: Any = None,
@@ -2894,6 +2896,7 @@ class MyModel:
           should re-validate defaults to config.revalidate_instances, else 'never'
         strict: Whether the model is strict
         frozen: Whether the model is frozen
+        extra_behavior: The extra behavior to use for the model, used in serialization
         config: The config to use for the model
         ref: optional unique identifier of the schema, used to reference the schema in other places
         metadata: Any other information you want to include with the schema, not used by pydantic-core
@@ -2907,6 +2910,7 @@ class MyModel:
         revalidate_instances=revalidate_instances,
         strict=strict,
         frozen=frozen,
+        extra_behavior=extra_behavior,
         config=config,
         ref=ref,
         metadata=metadata,

src/serializers/fields.rs

@@ -62,24 +62,6 @@ impl SerField {
         }
         Cow::Borrowed(key_str)
     }
-
-    pub fn to_python(
-        &self,
-        output_dict: &PyDict,
-        value: &PyAny,
-        next_include: Option<&PyAny>,
-        next_exclude: Option<&PyAny>,
-        extra: &Extra,
-    ) -> PyResult<()> {
-        if let Some(ref serializer) = self.serializer {
-            if !exclude_default(value, extra, serializer)? {
-                let value = serializer.to_python(value, next_include, next_exclude, extra)?;
-                let output_key = self.get_key_py(output_dict.py(), extra);
-                output_dict.set_item(output_key, value)?;
-            }
-        }
-        Ok(())
-    }
 }
 
 fn exclude_default(value: &PyAny, extra: &Extra, serializer: &CombinedSerializer) -> PyResult<bool> {
@@ -93,12 +75,22 @@ fn exclude_default(value: &PyAny, extra: &Extra, serializer: &CombinedSerializer
     Ok(false)
 }
 
+#[derive(Debug, Clone)]
+pub(super) enum FieldsMode {
+    // typeddict with no extra items
+    SimpleDict,
+    // a model - get `__dict__` and `__pydantic_extra__` - `GeneralFieldsSerializer` will get a tuple
+    ModelExtra,
+    // typeddict with extra items - one dict with extra items
+    TypedDictAllow,
+}
+
 /// General purpose serializer for fields - used by dataclasses, models and typed_dicts
 #[derive(Debug, Clone)]
 pub struct GeneralFieldsSerializer {
     fields: AHashMap<String, SerField>,
     computed_fields: Option<ComputedFields>,
-    include_extra: bool,
+    mode: FieldsMode,
     // isize because we look up filter via `.hash()` which returns an isize
     filter: SchemaFilter<isize>,
     required_fields: usize,
@@ -107,26 +99,39 @@ pub struct GeneralFieldsSerializer {
 impl GeneralFieldsSerializer {
     pub(super) fn new(
         fields: AHashMap<String, SerField>,
-        include_extra: bool,
+        mode: FieldsMode,
         computed_fields: Option<ComputedFields>,
     ) -> Self {
         let required_fields = fields.values().filter(|f| f.required).count();
         Self {
             fields,
-            include_extra,
+            mode,
             filter: SchemaFilter::default(),
             computed_fields,
             required_fields,
         }
     }
 
     fn extract_dicts<'a>(&self, value: &'a PyAny) -> Option<(&'a PyDict, Option<&'a PyDict>)> {
-        if let Ok(main_dict) = value.downcast::<PyDict>() {
-            Some((main_dict, None))
-        } else if let Ok((main_dict, extra_dict)) = value.extract::<(&PyDict, &PyDict)>() {
-            Some((main_dict, Some(extra_dict)))
-        } else {
-            None
+        match self.mode {
+            FieldsMode::ModelExtra => {
+                if let Ok((main_dict, extra_dict)) = value.extract::<(&PyDict, &PyAny)>() {
+                    if let Ok(extra_dict) = extra_dict.downcast::<PyDict>() {
+                        Some((main_dict, Some(extra_dict)))
+                    } else {
+                        Some((main_dict, None))
+                    }
+                } else {
+                    None
+                }
+            }
+            _ => {
+                if let Ok(main_dict) = value.downcast::<PyDict>() {
+                    Some((main_dict, None))
+                } else {
+                    None
+                }
+            }
         }
     }
 }
@@ -163,32 +168,39 @@ impl TypeSerializer for GeneralFieldsSerializer {
             return infer_to_python(value, include, exclude, &td_extra);
         };
 
-        // NOTE! we maintain the order of the input dict assuming that's right
         let output_dict = PyDict::new(py);
         let mut used_req_fields: usize = 0;
 
+        // NOTE! we maintain the order of the input dict assuming that's right
         for (key, value) in main_dict {
+            let key_str = key_str(key)?;
+            let op_field = self.fields.get(key_str);
             if extra.exclude_none && value.is_none() {
+                if let Some(field) = op_field {
+                    if field.required {
+                        used_req_fields += 1;
+                    }
+                }
                 continue;
             }
+            let extra = Extra {
+                field_name: Some(key_str),
+                ..td_extra
+            };
             if let Some((next_include, next_exclude)) = self.filter.key_filter(key, include, exclude)? {
-                let extra = Extra {
-                    field_name: Some(key.extract()?),
-                    ..td_extra
-                };
-                if let Ok(key_py_str) = key.downcast::<PyString>() {
-                    let key_str = key_py_str.to_str()?;
-                    if let Some(field) = self.fields.get(key_str) {
-                        field.to_python(output_dict, value, next_include, next_exclude, &extra)?;
-
-                        if field.required {
-                            used_req_fields += 1;
+                if let Some(field) = op_field {
+                    if let Some(ref serializer) = field.serializer {
+                        if !exclude_default(value, &extra, serializer)? {
+                            let value = serializer.to_python(value, next_include, next_exclude, &extra)?;
+                            let output_key = field.get_key_py(output_dict.py(), &extra);
+                            output_dict.set_item(output_key, value)?;
                         }
-                        continue;
                     }
-                }
-                if self.include_extra {
-                    // TODO test this
+
+                    if field.required {
+                        used_req_fields += 1;
+                    }
+                } else if matches!(self.mode, FieldsMode::TypedDictAllow) {
                     let value = infer_to_python(value, next_include, next_exclude, &extra)?;
                     output_dict.set_item(key, value)?;
                 } else if extra.check.enabled() {
@@ -242,40 +254,37 @@ impl TypeSerializer for GeneralFieldsSerializer {
             model: extra.model.map_or_else(|| Some(value), Some),
             ..*extra
         };
-        let expected_len = match self.include_extra {
-            true => main_dict.len() + option_length!(self.computed_fields),
-            false => self.fields.len() + option_length!(extra_dict) + option_length!(self.computed_fields),
+        let expected_len = match self.mode {
+            FieldsMode::TypedDictAllow => main_dict.len() + option_length!(self.computed_fields),
+            _ => self.fields.len() + option_length!(extra_dict) + option_length!(self.computed_fields),
         };
         // NOTE! As above, we maintain the order of the input dict assuming that's right
         // we don't both with `used_fields` here because on unions, `to_python(..., mode='json')` is used
         let mut map = serializer.serialize_map(Some(expected_len))?;
 
         for (key, value) in main_dict {
-            let extra = Extra {
-                field_name: Some(key.extract().map_err(py_err_se_err)?),
-                ..td_extra
-            };
             if extra.exclude_none && value.is_none() {
                 continue;
             }
+            let key_str = key_str(key).map_err(py_err_se_err)?;
+            let extra = Extra {
+                field_name: Some(key_str),
+                ..td_extra
+            };
+
             let filter = self.filter.key_filter(key, include, exclude).map_err(py_err_se_err)?;
             if let Some((next_include, next_exclude)) = filter {
-                if let Ok(key_py_str) = key.downcast::<PyString>() {
-                    let key_str = key_py_str.to_str().map_err(py_err_se_err)?;
-                    if let Some(field) = self.fields.get(key_str) {
-                        if let Some(ref serializer) = field.serializer {
-                            if !exclude_default(value, &extra, serializer).map_err(py_err_se_err)? {
-                                let s = PydanticSerializer::new(value, serializer, next_include, next_exclude, &extra);
-                                let output_key = field.get_key_json(key_str, &extra);
-                                map.serialize_entry(&output_key, &s)?;
-                            }
-                            continue;
+                if let Some(field) = self.fields.get(key_str) {
+                    if let Some(ref serializer) = field.serializer {
+                        if !exclude_default(value, &extra, serializer).map_err(py_err_se_err)? {
+                            let s = PydanticSerializer::new(value, serializer, next_include, next_exclude, &extra);
+                            let output_key = field.get_key_json(key_str, &extra);
+                            map.serialize_entry(&output_key, &s)?;
                         }
                     }
-                }
-                if self.include_extra {
-                    let s = SerializeInfer::new(value, include, exclude, &extra);
+                } else if matches!(self.mode, FieldsMode::TypedDictAllow) {
                     let output_key = infer_json_key(key, &extra).map_err(py_err_se_err)?;
+                    let s = SerializeInfer::new(value, next_include, next_exclude, &extra);
                     map.serialize_entry(&output_key, &s)?
                 }
             }
@@ -303,3 +312,7 @@ impl TypeSerializer for GeneralFieldsSerializer {
         "fields"
     }
 }
+
+fn key_str(key: &PyAny) -> PyResult<&str> {
+    key.downcast::<PyString>()?.to_str()
+}

src/serializers/type_serializers/dataclass.rs

@@ -4,11 +4,11 @@ use pyo3::types::{PyDict, PyList, PyString};
 
 use ahash::AHashMap;
 
-use crate::build_tools::{py_error_type, SchemaDict};
+use crate::build_tools::{py_error_type, ExtraBehavior, SchemaDict};
 use crate::definitions::DefinitionsBuilder;
 
 use super::model::ModelSerializer;
-use super::{BuildSerializer, CombinedSerializer, ComputedFields, GeneralFieldsSerializer, SerField};
+use super::{BuildSerializer, CombinedSerializer, ComputedFields, FieldsMode, GeneralFieldsSerializer, SerField};
 
 pub struct DataclassArgsBuilder;
 
@@ -25,6 +25,11 @@ impl BuildSerializer for DataclassArgsBuilder {
         let fields_list: &PyList = schema.get_as_req(intern!(py, "fields"))?;
         let mut fields: AHashMap<String, SerField> = AHashMap::with_capacity(fields_list.len());
 
+        let fields_mode = match ExtraBehavior::from_schema_or_config(py, schema, config, ExtraBehavior::Ignore)? {
+            ExtraBehavior::Allow => FieldsMode::TypedDictAllow,
+            _ => FieldsMode::SimpleDict,
+        };
+
         for (index, item) in fields_list.iter().enumerate() {
             let field_info: &PyDict = item.downcast()?;
             let name: String = field_info.get_as_req(intern!(py, "name"))?;
@@ -45,7 +50,7 @@ impl BuildSerializer for DataclassArgsBuilder {
 
         let computed_fields = ComputedFields::new(schema)?;
 
-        Ok(GeneralFieldsSerializer::new(fields, false, computed_fields).into())
+        Ok(GeneralFieldsSerializer::new(fields, fields_mode, computed_fields).into())
     }
 }
 

src/serializers/type_serializers/mod.rs

@@ -27,7 +27,7 @@ pub(self) use super::computed_fields::ComputedFields;
 pub(self) use super::config::utf8_py_error;
 pub(self) use super::errors::{py_err_se_err, PydanticSerializationError};
 pub(self) use super::extra::{Extra, ExtraOwned, SerCheck, SerMode};
-pub(self) use super::fields::{GeneralFieldsSerializer, SerField};
+pub(self) use super::fields::{FieldsMode, GeneralFieldsSerializer, SerField};
 pub(self) use super::filter::{AnyFilter, SchemaFilter};
 pub(self) use super::infer::{
     infer_json_key, infer_json_key_known, infer_serialize, infer_serialize_known, infer_to_python,

src/serializers/type_serializers/model.rs

@@ -11,8 +11,8 @@ use crate::definitions::DefinitionsBuilder;
 
 use super::{
     infer_json_key, infer_json_key_known, infer_serialize, infer_to_python, object_to_dict, py_err_se_err,
-    BuildSerializer, CombinedSerializer, ComputedFields, Extra, GeneralFieldsSerializer, ObType, SerCheck, SerField,
-    TypeSerializer,
+    BuildSerializer, CombinedSerializer, ComputedFields, Extra, FieldsMode, GeneralFieldsSerializer, ObType, SerCheck,
+    SerField, TypeSerializer,
 };
 
 pub struct ModelFieldsBuilder;
@@ -27,10 +27,10 @@ impl BuildSerializer for ModelFieldsBuilder {
     ) -> PyResult<CombinedSerializer> {
         let py = schema.py();
 
-        let include_extra = matches!(
-            ExtraBehavior::from_schema_or_config(py, schema, config, ExtraBehavior::Ignore)?,
-            ExtraBehavior::Allow
-        );
+        let fields_mode = match has_extra(schema, config)? {
+            true => FieldsMode::ModelExtra,
+            false => FieldsMode::SimpleDict,
+        };
 
         let fields_dict: &PyDict = schema.get_as_req(intern!(py, "fields"))?;
         let mut fields: AHashMap<String, SerField> = AHashMap::with_capacity(fields_dict.len());
@@ -57,14 +57,15 @@ impl BuildSerializer for ModelFieldsBuilder {
 
         let computed_fields = ComputedFields::new(schema)?;
 
-        Ok(GeneralFieldsSerializer::new(fields, include_extra, computed_fields).into())
+        Ok(GeneralFieldsSerializer::new(fields, fields_mode, computed_fields).into())
     }
 }
 
 #[derive(Debug, Clone)]
 pub struct ModelSerializer {
     class: Py<PyType>,
     serializer: Box<CombinedSerializer>,
+    has_extra: bool,
     name: String,
 }
 
@@ -84,12 +85,19 @@ impl BuildSerializer for ModelSerializer {
         Ok(Self {
             class: class.into(),
             serializer,
+            has_extra: has_extra(schema, config)?,
             name: class.getattr(intern!(py, "__name__"))?.extract()?,
         }
         .into())
     }
 }
 
+fn has_extra(schema: &PyDict, config: Option<&PyDict>) -> PyResult<bool> {
+    let py = schema.py();
+    let extra_behaviour = ExtraBehavior::from_schema_or_config(py, schema, config, ExtraBehavior::Ignore)?;
+    Ok(matches!(extra_behaviour, ExtraBehavior::Allow))
+}
+
 impl ModelSerializer {
     fn allow_value(&self, value: &PyAny, extra: &Extra) -> PyResult<bool> {
         match extra.check {
@@ -103,16 +111,12 @@ impl ModelSerializer {
         let py = value.py();
         let dict = object_to_dict(value, true, extra)?;
 
-        let model_extra = match value.getattr(intern!(py, "__pydantic_extra__")) {
-            Ok(model_extra) => model_extra,
-            Err(_) => return Ok(dict),
-        };
-
-        if model_extra.is_none() {
-            Ok(dict)
-        } else {
+        if self.has_extra {
+            let model_extra = value.getattr(intern!(py, "__pydantic_extra__"))?;
             let py_tuple = (dict, model_extra).to_object(py);
             Ok(py_tuple.into_ref(py))
+        } else {
+            Ok(dict)
         }
     }
 }