add polymorphic_serialization config for models and dataclasses

Author: davidhewitt

src/serializers/mod.rs

@@ -24,6 +24,7 @@ mod fields;
 mod filter;
 mod infer;
 mod ob_type;
+mod polymorphism_trampoline;
 mod prebuilt;
 pub mod ser;
 mod shared;

src/serializers/polymorphism_trampoline.rs

@@ -0,0 +1,94 @@
+use std::{borrow::Cow, sync::Arc};
+
+use pyo3::{prelude::*, types::PyType};
+
+use crate::serializers::{
+    errors::unwrap_ser_error,
+    infer::call_pydantic_serializer,
+    shared::{serialize_to_json, serialize_to_python, DoSerialize, TypeSerializer},
+    CombinedSerializer, SerializationState,
+};
+
+/// The polymorphism trampoline detects subclasses of its target type and dispatches to their
+/// `__pydantic_serializer__` serializer for serialization.
+///
+/// This exists as a separate structure to allow for cases such as model serializers where the
+/// inner serializer may just be a function serializer and so cannot handle polymorphism itself.
+pub struct PolymorphismTrampoline {
+    class: Py<PyType>,
+    /// Inner serializer used when the type is not a subclass (responsible for any fallback etc)
+    serializer: Arc<CombinedSerializer>,
+    /// Whether polymorphic serialization is enabled from config
+    enabled_from_config: bool,
+}
+
+impl_py_gc_traverse!(PolymorphismTrampoline { class, serializer });
+
+impl std::fmt::Debug for PolymorphismTrampoline {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        // to avoid breaking the repr from before we inserted the trampoline
+        self.serializer.fmt(f)
+    }
+}
+
+impl PolymorphismTrampoline {
+    pub fn new(class: Py<PyType>, serializer: Arc<CombinedSerializer>, enabled_from_config: bool) -> Self {
+        Self {
+            class,
+            serializer,
+            enabled_from_config,
+        }
+    }
+
+    fn is_subclass(&self, value: &Bound<'_, PyAny>) -> PyResult<bool> {
+        Ok(!value.get_type().is(&self.class) && value.is_instance(self.class.bind(value.py()))?)
+    }
+
+    fn serialize<'py, T, E: From<PyErr>>(
+        &self,
+        value: &Bound<'py, PyAny>,
+        state: &mut SerializationState<'_, 'py>,
+        do_serialize: impl DoSerialize<'py, T, E>,
+    ) -> Result<T, E> {
+        // FIXME: allow per-call override of polymorphic serialization?
+        if self.enabled_from_config && self.is_subclass(value)? {
+            call_pydantic_serializer(value, state, do_serialize)
+        } else {
+            do_serialize.serialize_no_infer(&self.serializer, value, state)
+        }
+    }
+}
+
+impl TypeSerializer for PolymorphismTrampoline {
+    fn to_python<'py>(
+        &self,
+        value: &Bound<'py, PyAny>,
+        state: &mut SerializationState<'_, 'py>,
+    ) -> PyResult<Py<PyAny>> {
+        self.serialize(value, state, serialize_to_python())
+    }
+
+    fn json_key<'a, 'py>(
+        &self,
+        key: &'a Bound<'py, PyAny>,
+        state: &mut SerializationState<'_, 'py>,
+    ) -> PyResult<Cow<'a, str>> {
+        // json key serialization for models and dataclasses was always polymorphic anyway
+        // FIXME: make this consistent with the other cases?
+        self.serializer.json_key(key, state)
+    }
+
+    fn serde_serialize<'py, S: serde::ser::Serializer>(
+        &self,
+        value: &Bound<'py, PyAny>,
+        serializer: S,
+        state: &mut SerializationState<'_, 'py>,
+    ) -> Result<S::Ok, S::Error> {
+        self.serialize(value, state, serialize_to_json(serializer))
+            .map_err(unwrap_ser_error)
+    }
+
+    fn get_name(&self) -> &str {
+        self.serializer.get_name()
+    }
+}

src/serializers/shared.rs

@@ -19,6 +19,7 @@ use crate::build_tools::py_schema_error_type;
 use crate::definitions::DefinitionsBuilder;
 use crate::py_gc::PyGcTraverse;
 use crate::serializers::errors::WrappedSerError;
+use crate::serializers::polymorphism_trampoline::PolymorphismTrampoline;
 use crate::serializers::ser::PythonSerializer;
 use crate::serializers::type_serializers::any::AnySerializer;
 use crate::tools::{py_err, SchemaDict};
@@ -91,6 +92,9 @@ combined_serializer! {
         Fields: super::fields::GeneralFieldsSerializer;
         // prebuilt serializers are manually constructed, and thus manually added to the `CombinedSerializer` enum
         Prebuilt: super::prebuilt::PrebuiltSerializer;
+        // polymorphism trampoline is manually constructed to wrap models and dataclasses with
+        // polymorphic serialization
+        PolymorphismTrampoline: super::polymorphism_trampoline::PolymorphismTrampoline;
     }
     // `find_only` is for type_serializers which are built directly via the `type` key and `find_serializer`
     // but aren't actually used for serialization, e.g. their `build` method must return another serializer
@@ -162,7 +166,8 @@ impl CombinedSerializer {
         config: Option<&Bound<'_, PyDict>>,
         definitions: &mut DefinitionsBuilder<Arc<CombinedSerializer>>,
     ) -> PyResult<Arc<CombinedSerializer>> {
-        Self::_build(schema, config, definitions, false)
+        let serializer = Self::_build(schema, config, definitions, false)?;
+        Self::maybe_wrap_in_polymorphism_trampoline(serializer, schema)
     }
 
     fn _build(
@@ -230,6 +235,34 @@ impl CombinedSerializer {
         Self::find_serializer(type_, schema, config, definitions)
     }
 
+    fn maybe_wrap_in_polymorphism_trampoline(
+        serializer: Arc<CombinedSerializer>,
+        schema: &Bound<'_, PyDict>,
+    ) -> PyResult<Arc<CombinedSerializer>> {
+        let py = schema.py();
+        let type_: Bound<'_, PyString> = schema.get_as_req(intern!(py, "type"))?;
+        let type_ = type_.to_str()?;
+
+        if type_ == "model" || type_ == "dataclass" {
+            // Get polymorphic serialization config
+            let config = schema.get_as::<Bound<'_, PyDict>>(intern!(py, "config"))?;
+            let polymorphic_serialization: bool = config
+                .and_then(|cfg| cfg.get_as(intern!(py, "polymorphic_serialization")).transpose())
+                .unwrap_or(Ok(false))?;
+
+            Ok(Arc::new(
+                PolymorphismTrampoline::new(
+                    schema.get_as_req(intern!(py, "cls"))?,
+                    serializer,
+                    polymorphic_serialization,
+                )
+                .into(),
+            ))
+        } else {
+            Ok(serializer)
+        }
+    }
+
     /// Main recursive way to call serializers, supports possible recursive type inference by
     /// switching to type inference mode eagerly.
     pub fn to_python<'py>(
@@ -308,7 +341,8 @@ impl BuildSerializer for CombinedSerializer {
         config: Option<&Bound<'_, PyDict>>,
         definitions: &mut DefinitionsBuilder<Arc<CombinedSerializer>>,
     ) -> PyResult<Arc<CombinedSerializer>> {
-        Self::_build(schema, config, definitions, true)
+        let serializer = Self::_build(schema, config, definitions, true)?;
+        Self::maybe_wrap_in_polymorphism_trampoline(serializer, schema)
     }
 }
 
@@ -356,6 +390,7 @@ impl PyGcTraverse for CombinedSerializer {
             CombinedSerializer::Tuple(inner) => inner.py_gc_traverse(visit),
             CombinedSerializer::Uuid(inner) => inner.py_gc_traverse(visit),
             CombinedSerializer::Complex(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::PolymorphismTrampoline(inner) => inner.py_gc_traverse(visit),
         }
     }
 }

tests/serializers/test_dataclasses.py

@@ -286,3 +286,46 @@ class Foo:
     )
     s = SchemaSerializer(schema)
     assert s.to_python(Foo(my_field='hello'), by_alias=runtime) == expected
+
+
+@pytest.mark.parametrize('config', [True, False])
+@pytest.mark.parametrize('runtime', [True, False])
+def test_polymorphic_serialization(config: bool, runtime: bool) -> None:
+    @dataclasses.dataclass
+    class ClassA:
+        a: int
+
+    @dataclasses.dataclass
+    class ClassB(ClassA):
+        b: str
+
+    schema_a = core_schema.dataclass_schema(
+        ClassA,
+        core_schema.dataclass_args_schema(
+            'ClassA', [core_schema.dataclass_field(name='a', schema=core_schema.int_schema())]
+        ),
+        ['a'],
+        config=core_schema.CoreConfig(polymorphic_serialization=config),
+    )
+
+    schema_b = core_schema.dataclass_schema(
+        ClassB,
+        core_schema.dataclass_args_schema(
+            'ClassB',
+            [
+                core_schema.dataclass_field(name='a', schema=core_schema.int_schema()),
+                core_schema.dataclass_field(name='b', schema=core_schema.str_schema()),
+            ],
+        ),
+        ['a', 'b'],
+    )
+
+    ClassA.__pydantic_serializer__ = SchemaSerializer(schema_a)
+    ClassB.__pydantic_serializer__ = SchemaSerializer(schema_b)
+
+    assert ClassA.__pydantic_serializer__.to_python(ClassA(123)) == {'a': 123}
+
+    if config:
+        assert ClassA.__pydantic_serializer__.to_python(ClassB(123, 'test')) == {'a': 123, 'b': 'test'}
+    else:
+        assert ClassA.__pydantic_serializer__.to_python(ClassB(123, 'test')) == {'a': 123}

tests/serializers/test_model.py

@@ -1284,3 +1284,42 @@ def __init__(self, my_field: int) -> None:
     )
     s = SchemaSerializer(schema)
     assert s.to_python(Model(1), by_alias=runtime) == expected
+
+
+@pytest.mark.parametrize('config', [True, False])
+@pytest.mark.parametrize('runtime', [True, False])
+def test_polymorphic_serialization(config: bool, runtime: bool) -> None:
+    class ModelA:
+        def __init__(self, a: int) -> None:
+            self.a = a
+
+    class ModelB(ModelA):
+        def __init__(self, a: int, b: str) -> None:
+            super().__init__(a)
+            self.b = b
+
+    schema_a = core_schema.model_schema(
+        ModelA,
+        core_schema.model_fields_schema({'a': core_schema.model_field(core_schema.int_schema())}),
+        config=core_schema.CoreConfig(polymorphic_serialization=config),
+    )
+
+    schema_b = core_schema.model_schema(
+        ModelB,
+        core_schema.model_fields_schema(
+            {
+                'a': core_schema.model_field(core_schema.int_schema()),
+                'b': core_schema.model_field(core_schema.str_schema()),
+            }
+        ),
+    )
+
+    ModelA.__pydantic_serializer__ = SchemaSerializer(schema_a)
+    ModelB.__pydantic_serializer__ = SchemaSerializer(schema_b)
+
+    assert ModelA.__pydantic_serializer__.to_python(ModelA(123)) == {'a': 123}
+
+    if config:
+        assert ModelA.__pydantic_serializer__.to_python(ModelB(123, 'test')) == {'a': 123, 'b': 'test'}
+    else:
+        assert ModelA.__pydantic_serializer__.to_python(ModelB(123, 'test')) == {'a': 123}