Use return_schema for computed fields (#595)

Author: dmontagu

pydantic_core/core_schema.py

@@ -415,23 +415,25 @@ def model_ser_schema(cls: Type[Any], schema: CoreSchema) -> ModelSerSchema:
 class ComputedField(TypedDict, total=False):
     type: Required[Literal['computed-field']]
     property_name: Required[str]
-    json_return_type: JsonReturnTypes
+    return_schema: Required[CoreSchema]
     alias: str
+    metadata: Any
 
 
 def computed_field(
-    property_name: str, *, json_return_type: JsonReturnTypes | None = None, alias: str | None = None
+    property_name: str, return_schema: CoreSchema, *, alias: str | None = None, metadata: Any = None
 ) -> ComputedField:
     """
     ComputedFields are properties of a model or dataclass that are included in serialization.
 
     Args:
         property_name: The name of the property on the model or dataclass
-        json_return_type: The type that the property returns if `mode='json'`
+        return_schema: The schema used for the type returned by the computed field
         alias: The name to use in the serialized output
+        metadata: Any other information you want to include with the schema, not used by pydantic-core
     """
     return dict_not_none(
-        type='computed-field', property_name=property_name, json_return_type=json_return_type, alias=alias
+        type='computed-field', property_name=property_name, return_schema=return_schema, alias=alias, metadata=metadata
     )
 
 
@@ -3677,6 +3679,8 @@ def definition_reference_schema(
     'definition-ref',
 ]
 
+CoreSchemaFieldType = Literal['model-field', 'dataclass-field', 'typed-dict-field', 'computed-field']
+
 
 # used in _pydantic_core.pyi::PydanticKnownError
 # to update this, call `pytest -k test_all_errors` and copy the output

src/serializers/computed_fields.rs

@@ -1,30 +1,31 @@
 use pyo3::intern;
 use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyList, PyString};
-
 use serde::ser::SerializeMap;
 use serde::Serialize;
 
-use crate::build_tools::SchemaDict;
+use crate::build_tools::{py_error_type, SchemaDict};
+use crate::definitions::DefinitionsBuilder;
 use crate::serializers::filter::SchemaFilter;
+use crate::serializers::shared::{BuildSerializer, CombinedSerializer, PydanticSerializer, TypeSerializer};
 
 use super::errors::py_err_se_err;
-use super::infer::{infer_serialize, infer_serialize_known, infer_to_python, infer_to_python_known};
-use super::ob_type::ObType;
 use super::Extra;
 
-use super::type_serializers::function::get_json_return_type;
-
 #[derive(Debug, Clone)]
 pub(super) struct ComputedFields(Vec<ComputedField>);
 
 impl ComputedFields {
-    pub fn new(schema: &PyDict) -> PyResult<Option<Self>> {
+    pub fn new(
+        schema: &PyDict,
+        config: Option<&PyDict>,
+        definitions: &mut DefinitionsBuilder<CombinedSerializer>,
+    ) -> PyResult<Option<Self>> {
         let py = schema.py();
         if let Some(computed_fields) = schema.get_as::<&PyList>(intern!(py, "computed_fields"))? {
             let computed_fields = computed_fields
                 .iter()
-                .map(ComputedField::new)
+                .map(|field| ComputedField::new(field, config, definitions))
                 .collect::<PyResult<Vec<_>>>()?;
             Ok(Some(Self(computed_fields)))
         } else {
@@ -88,22 +89,28 @@ impl ComputedFields {
 struct ComputedField {
     property_name: String,
     property_name_py: Py<PyString>,
-    return_ob_type: Option<ObType>,
+    serializer: CombinedSerializer,
     alias: String,
     alias_py: Py<PyString>,
 }
 
 impl ComputedField {
-    pub fn new(schema: &PyAny) -> PyResult<Self> {
+    pub fn new(
+        schema: &PyAny,
+        config: Option<&PyDict>,
+        definitions: &mut DefinitionsBuilder<CombinedSerializer>,
+    ) -> PyResult<Self> {
         let py = schema.py();
         let schema: &PyDict = schema.downcast()?;
         let property_name: &PyString = schema.get_as_req(intern!(py, "property_name"))?;
-        let return_ob_type = get_json_return_type(schema)?;
+        let return_schema = schema.get_as_req(intern!(py, "return_schema"))?;
+        let serializer = CombinedSerializer::build(return_schema, config, definitions)
+            .map_err(|e| py_error_type!("Computed field `{}`:\n  {}", property_name, e))?;
         let alias_py: &PyString = schema.get_as(intern!(py, "alias"))?.unwrap_or(property_name);
         Ok(Self {
             property_name: property_name.extract()?,
             property_name_py: property_name.into_py(py),
-            return_ob_type,
+            serializer,
             alias: alias_py.extract()?,
             alias_py: alias_py.into_py(py),
         })
@@ -124,11 +131,9 @@ impl ComputedField {
         if let Some((next_include, next_exclude)) = filter.key_filter(property_name_py, include, exclude)? {
             let next_value = model.getattr(property_name_py)?;
 
-            // TODO fix include & exclude
-            let value = match self.return_ob_type {
-                Some(ref ob_type) => infer_to_python_known(ob_type, next_value, next_include, next_exclude, extra),
-                None => infer_to_python(next_value, next_include, next_exclude, extra),
-            }?;
+            let value = self
+                .serializer
+                .to_python(next_value, next_include, next_exclude, extra)?;
             let key = match extra.by_alias {
                 true => self.alias_py.as_ref(py),
                 false => property_name_py,
@@ -152,12 +157,13 @@ impl<'py> Serialize for ComputedFieldSerializer<'py> {
         let py = self.model.py();
         let property_name_py = self.computed_field.property_name_py.as_ref(py);
         let next_value = self.model.getattr(property_name_py).map_err(py_err_se_err)?;
-
-        match self.computed_field.return_ob_type {
-            Some(ref ob_type) => {
-                infer_serialize_known(ob_type, next_value, serializer, self.include, self.exclude, self.extra)
-            }
-            None => infer_serialize(next_value, serializer, self.include, self.exclude, self.extra),
-        }
+        let s = PydanticSerializer::new(
+            next_value,
+            &self.computed_field.serializer,
+            self.include,
+            self.exclude,
+            self.extra,
+        );
+        s.serialize(serializer)
     }
 }

src/serializers/type_serializers/dataclass.rs

@@ -48,7 +48,7 @@ impl BuildSerializer for DataclassArgsBuilder {
             }
         }
 
-        let computed_fields = ComputedFields::new(schema)?;
+        let computed_fields = ComputedFields::new(schema, config, definitions)?;
 
         Ok(GeneralFieldsSerializer::new(fields, fields_mode, computed_fields).into())
     }

src/serializers/type_serializers/model.rs

@@ -55,7 +55,7 @@ impl BuildSerializer for ModelFieldsBuilder {
             }
         }
 
-        let computed_fields = ComputedFields::new(schema)?;
+        let computed_fields = ComputedFields::new(schema, config, definitions)?;
 
         Ok(GeneralFieldsSerializer::new(fields, fields_mode, computed_fields).into())
     }

src/serializers/type_serializers/typed_dict.rs

@@ -53,7 +53,7 @@ impl BuildSerializer for TypedDictBuilder {
             }
         }
 
-        let computed_fields = ComputedFields::new(schema)?;
+        let computed_fields = ComputedFields::new(schema, config, definitions)?;
 
         Ok(GeneralFieldsSerializer::new(fields, fields_mode, computed_fields).into())
     }

tests/serializers/test_dataclasses.py

@@ -105,7 +105,7 @@ def c(self) -> str:
                 core_schema.dataclass_field(name='a', schema=core_schema.str_schema()),
                 core_schema.dataclass_field(name='b', schema=core_schema.bytes_schema()),
             ],
-            computed_fields=[core_schema.computed_field('c')],
+            computed_fields=[core_schema.computed_field('c', core_schema.str_schema())],
         ),
     )
     s = SchemaSerializer(schema)

tests/serializers/test_model.py

@@ -546,7 +546,7 @@ def area(self) -> bytes:
                     'width': core_schema.model_field(core_schema.int_schema()),
                     'height': core_schema.model_field(core_schema.int_schema()),
                 },
-                computed_fields=[core_schema.computed_field('area', json_return_type='bytes')],
+                computed_fields=[core_schema.computed_field('area', core_schema.bytes_schema())],
             ),
         )
     )
@@ -578,8 +578,8 @@ def volume(self) -> int:
                     'height': core_schema.model_field(core_schema.int_schema()),
                 },
                 computed_fields=[
-                    core_schema.computed_field('area', alias='Area'),
-                    core_schema.computed_field('volume'),
+                    core_schema.computed_field('area', core_schema.int_schema(), alias='Area'),
+                    core_schema.computed_field('volume', core_schema.int_schema()),
                 ],
             ),
         )
@@ -608,7 +608,7 @@ def area(self) -> int:
                     'width': core_schema.model_field(core_schema.int_schema()),
                     'height': core_schema.model_field(core_schema.int_schema()),
                 },
-                computed_fields=[core_schema.computed_field('area')],
+                computed_fields=[core_schema.computed_field('area', core_schema.int_schema())],
             ),
         )
     )
@@ -627,7 +627,7 @@ class Model:
             Model,
             core_schema.model_fields_schema(
                 {'width': core_schema.model_field(core_schema.int_schema())},
-                computed_fields=[core_schema.computed_field('area', json_return_type='bytes')],
+                computed_fields=[core_schema.computed_field('area', core_schema.bytes_schema())],
             ),
         )
     )
@@ -655,7 +655,7 @@ def area(self) -> int:
             Model,
             core_schema.model_fields_schema(
                 {'width': core_schema.model_field(core_schema.int_schema())},
-                computed_fields=[core_schema.computed_field('area', json_return_type='bytes')],
+                computed_fields=[core_schema.computed_field('area', core_schema.bytes_schema())],
             ),
         )
     )
@@ -684,7 +684,7 @@ def b(self):
             Model,
             core_schema.model_fields_schema(
                 {'a': core_schema.model_field(core_schema.int_schema())},
-                computed_fields=[core_schema.computed_field('b')],
+                computed_fields=[core_schema.computed_field('b', core_schema.list_schema())],
             ),
         )
     )
@@ -734,8 +734,8 @@ def random_n(self) -> int:
             core_schema.model_fields_schema(
                 {'side': core_schema.model_field(core_schema.float_schema())},
                 computed_fields=[
-                    core_schema.computed_field('area', json_return_type='float'),
-                    core_schema.computed_field('random_n', alias='The random number', json_return_type='int'),
+                    core_schema.computed_field('area', core_schema.float_schema()),
+                    core_schema.computed_field('random_n', core_schema.int_schema(), alias='The random number'),
                 ],
             ),
         )

tests/test.rs

@@ -42,7 +42,9 @@ schema = {
     "cls": A,
     "config": {},
     "schema": {
-        "computed_fields": [{"property_name": "b", "type": "computed-field"}],
+        "computed_fields": [
+            {"property_name": "b", "return_schema": {"type": "any"}, "type": "computed-field"}
+        ],
         "fields": {},
         "type": "model-fields",
     },