move check from Extra to SerializationState (#1862)

Author: davidhewitt

src/serializers/extra.rs

@@ -30,6 +30,8 @@ pub(crate) struct SerializationState<'py> {
     pub model: Option<Bound<'py, PyAny>>,
     /// The name of the field currently being serialized, if any
     pub field_name: Option<FieldName<'py>>,
+    /// Inside unions, checks are applied to attempt to select a preferred branch
+    pub check: SerCheck,
     pub include_exclude: (Option<Bound<'py, PyAny>>, Option<Bound<'py, PyAny>>),
 }
 
@@ -69,6 +71,7 @@ impl<'py> SerializationState<'py> {
             config,
             model: None,
             field_name: None,
+            check: SerCheck::None,
             include_exclude: (include, exclude),
         })
     }
@@ -86,19 +89,18 @@ impl SerializationState<'_> {
         })
     }
 
-    pub fn warn_fallback_py(&mut self, field_type: &str, value: &Bound<'_, PyAny>, extra: &Extra) -> PyResult<()> {
+    pub fn warn_fallback_py(&mut self, field_type: &str, value: &Bound<'_, PyAny>) -> PyResult<()> {
         self.warnings
-            .on_fallback_py(field_type, value, self.field_name.as_ref(), extra)
+            .on_fallback_py(field_type, value, self.field_name.as_ref(), self.check)
     }
 
-    pub fn warn_fallback_ser<'py, S: serde::ser::Serializer>(
+    pub fn warn_fallback_ser<S: serde::ser::Serializer>(
         &mut self,
         field_type: &str,
-        value: &Bound<'py, PyAny>,
-        extra: &Extra<'_, 'py>,
+        value: &Bound<'_, PyAny>,
     ) -> Result<(), S::Error> {
         self.warnings
-            .on_fallback_ser::<S>(field_type, value, self.field_name.as_ref(), extra)
+            .on_fallback_ser::<S>(field_type, value, self.field_name.as_ref(), self.check)
     }
 
     pub fn final_check(&self, py: Python) -> PyResult<()> {
@@ -160,8 +162,6 @@ pub(crate) struct Extra<'a, 'py> {
     pub exclude_none: bool,
     pub exclude_computed_fields: bool,
     pub round_trip: bool,
-    // the next two are used for union logic
-    pub check: SerCheck,
     pub serialize_unknown: bool,
     pub fallback: Option<&'a Bound<'py, PyAny>>,
     pub serialize_as_any: bool,
@@ -193,7 +193,6 @@ impl<'a, 'py> Extra<'a, 'py> {
             exclude_none,
             exclude_computed_fields,
             round_trip,
-            check: SerCheck::None,
             serialize_unknown,
             fallback,
             serialize_as_any,
@@ -283,7 +282,7 @@ impl ExtraOwned {
             round_trip: extra.round_trip,
             config: state.config,
             rec_guard: state.rec_guard.clone(),
-            check: extra.check,
+            check: state.check,
             model: state.model.as_ref().map(|model| model.clone().into()),
             field_name: state.field_name.as_ref().map(|name| match name {
                 FieldName::Root => FieldNameOwned::Root,
@@ -308,7 +307,6 @@ impl ExtraOwned {
             exclude_none: self.exclude_none,
             exclude_computed_fields: self.exclude_computed_fields,
             round_trip: self.round_trip,
-            check: self.check,
             serialize_unknown: self.serialize_unknown,
             fallback: self.fallback.as_ref().map(|m| m.bind(py)),
             serialize_as_any: self.serialize_as_any,
@@ -321,16 +319,17 @@ impl ExtraOwned {
             warnings: self.warnings.clone(),
             rec_guard: self.rec_guard.clone(),
             config: self.config,
+            model: self.model.as_ref().map(|m| m.bind(py).clone()),
             field_name: match &self.field_name {
                 Some(FieldNameOwned::Root) => Some(FieldName::Root),
                 Some(FieldNameOwned::Regular(b)) => Some(FieldName::Regular(b.bind(py).clone())),
                 None => None,
             },
+            check: self.check,
             include_exclude: (
                 self.include.as_ref().map(|m| m.bind(py).clone()),
                 self.exclude.as_ref().map(|m| m.bind(py).clone()),
             ),
-            model: self.model.as_ref().map(|m| m.bind(py).clone()),
         }
     }
 }
@@ -443,17 +442,17 @@ impl CollectWarnings {
         }
     }
 
-    pub fn on_fallback_py<'py>(
+    fn on_fallback_py(
         &mut self,
         field_type: &str,
-        value: &Bound<'py, PyAny>,
+        value: &Bound<'_, PyAny>,
         field_name: Option<&FieldName<'_>>,
-        extra: &Extra<'_, 'py>,
+        check: SerCheck,
     ) -> PyResult<()> {
         // special case for None as it's very common e.g. as a default value
         if value.is_none() {
             Ok(())
-        } else if extra.check.enabled() {
+        } else if check.enabled() {
             Err(PydanticSerializationUnexpectedValue::new_from_parts(
                 field_name.map(ToString::to_string),
                 Some(field_type.to_string()),
@@ -466,17 +465,17 @@ impl CollectWarnings {
         }
     }
 
-    pub fn on_fallback_ser<'py, S: serde::ser::Serializer>(
+    pub fn on_fallback_ser<S: serde::ser::Serializer>(
         &mut self,
         field_type: &str,
-        value: &Bound<'py, PyAny>,
+        value: &Bound<'_, PyAny>,
         field_name: Option<&FieldName<'_>>,
-        extra: &Extra<'_, 'py>,
+        check: SerCheck,
     ) -> Result<(), S::Error> {
         // special case for None as it's very common e.g. as a default value
         if value.is_none() {
             Ok(())
-        } else if extra.check.enabled() {
+        } else if check.enabled() {
             // note: I think this should never actually happen since we use `to_python(..., mode='json')` during
             // JSON serialization to "try" union branches, but it's here for completeness/correctness
             // in particular, in future we could allow errors instead of warnings on fallback

src/serializers/fields.rs

@@ -215,7 +215,7 @@ impl GeneralFieldsSerializer {
                         .filter(|_| !extra.serialize_as_any)
                         .unwrap_or_else(|| AnySerializer::get());
                     (&key, serializer)
-                } else if extra.check == SerCheck::Strict {
+                } else if state.check == SerCheck::Strict {
                     return Err(PydanticSerializationUnexpectedValue::new(
                         Some(format!("Unexpected field `{key}`")),
                         Some(key_str.to_string()),
@@ -233,7 +233,7 @@ impl GeneralFieldsSerializer {
             }
         }
 
-        if extra.check.enabled()
+        if state.check.enabled()
             // If any of these are true we can't count fields
             && !(extra.exclude_defaults || extra.exclude_unset || extra.exclude_none || extra.exclude_computed_fields || state.exclude().is_some())
             // Check for missing fields, we can't have extra fields here
@@ -399,7 +399,7 @@ impl TypeSerializer for GeneralFieldsSerializer {
         let model = state.model.clone().unwrap_or_else(|| value.clone());
 
         let Some((main_dict, extra_dict)) = self.extract_dicts(value) else {
-            state.warn_fallback_py(self.get_name(), value, extra)?;
+            state.warn_fallback_py(self.get_name(), value)?;
             return infer_to_python(value, state, extra);
         };
         let output_dict = self.main_to_python(
@@ -450,7 +450,7 @@ impl TypeSerializer for GeneralFieldsSerializer {
         extra: &Extra<'_, 'py>,
     ) -> Result<S::Ok, S::Error> {
         let Some((main_dict, extra_dict)) = self.extract_dicts(value) else {
-            state.warn_fallback_ser::<S>(self.get_name(), value, extra)?;
+            state.warn_fallback_ser::<S>(self.get_name(), value)?;
             return infer_serialize(value, serializer, state, extra);
         };
         let missing_sentinel = get_missing_sentinel_object(value.py());

src/serializers/infer.rs

@@ -648,6 +648,7 @@ pub(crate) fn call_pydantic_serializer<'py, T, E: From<PyErr>>(
         model: state.model.clone(),
         field_name: state.field_name.clone(),
         include_exclude: state.include_exclude.clone(),
+        check: state.check,
     };
 
     // Avoid falling immediately back into inference because we need to use the serializer

src/serializers/shared.rs

@@ -390,7 +390,7 @@ pub(crate) trait TypeSerializer: Send + Sync + Debug {
         match extra.ob_type_lookup.is_type(key, ObType::None) {
             IsType::Exact | IsType::Subclass => py_err!(PyTypeError; "`{}` not valid as object key", expected_type),
             IsType::False => {
-                state.warn_fallback_py(self.get_name(), key, extra)?;
+                state.warn_fallback_py(self.get_name(), key)?;
                 infer_json_key(key, state, extra)
             }
         }
@@ -680,7 +680,7 @@ impl<'py> DoSerialize<'py, Py<PyAny>, PyErr> for SerializeToPython {
         state: &mut SerializationState<'py>,
         extra: &Extra<'_, 'py>,
     ) -> PyResult<Py<PyAny>> {
-        state.warn_fallback_py(name, value, extra)?;
+        state.warn_fallback_py(name, value)?;
         infer_to_python(value, state, extra)
     }
 }
@@ -709,9 +709,7 @@ impl<'py, S: Serializer> DoSerialize<'py, S::Ok, WrappedSerError<S::Error>> for
         state: &mut SerializationState<'py>,
         extra: &Extra<'_, 'py>,
     ) -> Result<S::Ok, WrappedSerError<S::Error>> {
-        state
-            .warn_fallback_ser::<S>(name, value, extra)
-            .map_err(WrappedSerError)?;
+        state.warn_fallback_ser::<S>(name, value).map_err(WrappedSerError)?;
         infer_serialize(value, self.serializer, state, extra).map_err(WrappedSerError)
     }
 }

src/serializers/type_serializers/bytes.rs

@@ -82,7 +82,7 @@ impl TypeSerializer for BytesSerializer {
                 _ => Ok(value.clone().unbind()),
             },
             Err(_) => {
-                state.warn_fallback_py(self.get_name(), value, extra)?;
+                state.warn_fallback_py(self.get_name(), value)?;
                 infer_to_python(value, state, extra)
             }
         }
@@ -97,7 +97,7 @@ impl TypeSerializer for BytesSerializer {
         match key.downcast::<PyBytes>() {
             Ok(py_bytes) => self.bytes_mode.bytes_to_string(key.py(), py_bytes.as_bytes()),
             Err(_) => {
-                state.warn_fallback_py(self.get_name(), key, extra)?;
+                state.warn_fallback_py(self.get_name(), key)?;
                 infer_json_key(key, state, extra)
             }
         }
@@ -113,7 +113,7 @@ impl TypeSerializer for BytesSerializer {
         match value.downcast::<PyBytes>() {
             Ok(py_bytes) => self.bytes_mode.serialize_bytes(py_bytes.as_bytes(), serializer),
             Err(_) => {
-                state.warn_fallback_ser::<S>(self.get_name(), value, extra)?;
+                state.warn_fallback_ser::<S>(self.get_name(), value)?;
                 infer_serialize(value, serializer, state, extra)
             }
         }

src/serializers/type_serializers/complex.rs

@@ -42,7 +42,7 @@ impl TypeSerializer for ComplexSerializer {
                 _ => Ok(value.clone().unbind()),
             },
             Err(_) => {
-                state.warn_fallback_py(self.get_name(), value, extra)?;
+                state.warn_fallback_py(self.get_name(), value)?;
                 infer_to_python(value, state, extra)
             }
         }
@@ -70,7 +70,7 @@ impl TypeSerializer for ComplexSerializer {
                 Ok(serializer.collect_str::<String>(&s)?)
             }
             Err(_) => {
-                state.warn_fallback_ser::<S>(self.get_name(), value, extra)?;
+                state.warn_fallback_ser::<S>(self.get_name(), value)?;
                 infer_serialize(value, serializer, state, extra)
             }
         }

src/serializers/type_serializers/dataclass.rs

@@ -122,8 +122,8 @@ impl BuildSerializer for DataclassSerializer {
 }
 
 impl DataclassSerializer {
-    fn allow_value(&self, value: &Bound<'_, PyAny>, extra: &Extra) -> PyResult<bool> {
-        match extra.check {
+    fn allow_value(&self, value: &Bound<'_, PyAny>, state: &SerializationState<'_>) -> PyResult<bool> {
+        match state.check {
             SerCheck::Strict => Ok(value.get_type().is(self.class.bind(value.py()))),
             SerCheck::Lax => value.is_instance(self.class.bind(value.py())),
             SerCheck::None => value.hasattr(intern!(value.py(), "__dataclass_fields__")),
@@ -152,7 +152,7 @@ impl TypeSerializer for DataclassSerializer {
         extra: &Extra<'_, 'py>,
     ) -> PyResult<Py<PyAny>> {
         let state = &mut state.scoped_set(|s| &mut s.model, Some(value.clone()));
-        if self.allow_value(value, extra)? {
+        if self.allow_value(value, state)? {
             let py = value.py();
             if let CombinedSerializer::Fields(ref fields_serializer) = *self.serializer {
                 let output_dict: Bound<PyDict> =
@@ -166,7 +166,7 @@ impl TypeSerializer for DataclassSerializer {
             }
         } else {
             // FIXME: probably don't want to have state.model set here, should move the scoped_set above?
-            state.warn_fallback_py(self.get_name(), value, extra)?;
+            state.warn_fallback_py(self.get_name(), value)?;
             infer_to_python(value, state, extra)
         }
     }
@@ -177,10 +177,10 @@ impl TypeSerializer for DataclassSerializer {
         state: &mut SerializationState<'py>,
         extra: &Extra<'_, 'py>,
     ) -> PyResult<Cow<'a, str>> {
-        if self.allow_value(key, extra)? {
+        if self.allow_value(key, state)? {
             infer_json_key_known(ObType::Dataclass, key, state, extra)
         } else {
-            state.warn_fallback_py(&self.name, key, extra)?;
+            state.warn_fallback_py(&self.name, key)?;
             infer_json_key(key, state, extra)
         }
     }
@@ -193,7 +193,7 @@ impl TypeSerializer for DataclassSerializer {
         extra: &Extra<'_, 'py>,
     ) -> Result<S::Ok, S::Error> {
         let state = &mut state.scoped_set(|s| &mut s.model, Some(value.clone()));
-        if self.allow_value(value, extra).map_err(py_err_se_err)? {
+        if self.allow_value(value, state).map_err(py_err_se_err)? {
             if let CombinedSerializer::Fields(ref fields_serializer) = *self.serializer {
                 let expected_len = self.fields.len() + fields_serializer.computed_field_count();
                 let mut map = fields_serializer.main_serde_serialize(
@@ -211,7 +211,7 @@ impl TypeSerializer for DataclassSerializer {
             }
         } else {
             // FIXME: probably don't want to have state.model set here, should move the scoped_set above?
-            state.warn_fallback_ser::<S>(self.get_name(), value, extra)?;
+            state.warn_fallback_ser::<S>(self.get_name(), value)?;
             infer_serialize(value, serializer, state, extra)
         }
     }

src/serializers/type_serializers/datetime_etc.rs

@@ -125,7 +125,7 @@ macro_rules! build_temporal_serializer {
                         _ => Ok(value.clone().unbind()),
                     },
                     _ => {
-                        state.warn_fallback_py(self.get_name(), value, extra)?;
+                        state.warn_fallback_py(self.get_name(), value)?;
                         infer_to_python(value, state, extra)
                     }
                 }
@@ -140,7 +140,7 @@ macro_rules! build_temporal_serializer {
                 match $downcast(key) {
                     Ok(py_value) => Ok(self.temporal_mode.$json_key_fn(py_value)?),
                     Err(_) => {
-                        state.warn_fallback_py(self.get_name(), key, extra)?;
+                        state.warn_fallback_py(self.get_name(), key)?;
                         infer_json_key(key, state, extra)
                     }
                 }
@@ -156,7 +156,7 @@ macro_rules! build_temporal_serializer {
                 match $downcast(value) {
                     Ok(py_value) => self.temporal_mode.$serialize_fn(py_value, serializer),
                     Err(_) => {
-                        state.warn_fallback_ser::<S>(self.get_name(), value, extra)?;
+                        state.warn_fallback_ser::<S>(self.get_name(), value)?;
                         infer_serialize(value, serializer, state, extra)
                     }
                 }

src/serializers/type_serializers/decimal.rs

@@ -44,7 +44,7 @@ impl TypeSerializer for DecimalSerializer {
         match extra.ob_type_lookup.is_type(value, ObType::Decimal) {
             IsType::Exact | IsType::Subclass => infer_to_python_known(ObType::Decimal, value, state, extra),
             IsType::False => {
-                state.warn_fallback_py(self.get_name(), value, extra)?;
+                state.warn_fallback_py(self.get_name(), value)?;
                 infer_to_python(value, state, extra)
             }
         }
@@ -59,7 +59,7 @@ impl TypeSerializer for DecimalSerializer {
         match extra.ob_type_lookup.is_type(key, ObType::Decimal) {
             IsType::Exact | IsType::Subclass => infer_json_key_known(ObType::Decimal, key, state, extra),
             IsType::False => {
-                state.warn_fallback_py(self.get_name(), key, extra)?;
+                state.warn_fallback_py(self.get_name(), key)?;
                 infer_json_key(key, state, extra)
             }
         }
@@ -75,7 +75,7 @@ impl TypeSerializer for DecimalSerializer {
         match extra.ob_type_lookup.is_type(value, ObType::Decimal) {
             IsType::Exact | IsType::Subclass => infer_serialize_known(ObType::Decimal, value, serializer, state, extra),
             IsType::False => {
-                state.warn_fallback_ser::<S>(self.get_name(), value, extra)?;
+                state.warn_fallback_ser::<S>(self.get_name(), value)?;
                 infer_serialize(value, serializer, state, extra)
             }
         }

src/serializers/type_serializers/dict.rs

@@ -105,7 +105,7 @@ impl TypeSerializer for DictSerializer {
                 Ok(new_dict.into())
             }
             Err(_) => {
-                state.warn_fallback_py(self.get_name(), value, extra)?;
+                state.warn_fallback_py(self.get_name(), value)?;
                 infer_to_python(value, state, extra)
             }
         }
@@ -145,7 +145,7 @@ impl TypeSerializer for DictSerializer {
                 map.end()
             }
             Err(_) => {
-                state.warn_fallback_ser::<S>(self.get_name(), value, extra)?;
+                state.warn_fallback_ser::<S>(self.get_name(), value)?;
                 infer_serialize(value, serializer, state, extra)
             }
         }