Flatten CoreSchema types to get a single discriminant key

Author: adriangb

generate_self_schema.py

@@ -197,26 +197,9 @@ def main() -> None:
         assert m, f'Unknown schema type: {type_}'
         key = m.group(1)
         value = get_schema(s)
-        if key == 'function':
-            mode = value['fields']['mode']['schema']['expected']
-            if mode == ['plain']:
-                key = 'function-plain'
-            elif mode == ['wrap']:
-                key = 'function-wrap'
-        elif key == 'tuple':
-            if value['fields']['mode']['schema']['expected'] == ['positional']:
-                key = 'tuple-positional'
-            else:
-                key = 'tuple-variable'
-
         choices[key] = value
 
-    schema = {
-        'type': 'tagged-union',
-        'ref': 'root-schema',
-        'discriminator': 'self-schema-discriminator',
-        'choices': choices,
-    }
+    schema = {'type': 'tagged-union', 'ref': 'root-schema', 'discriminator': 'type', 'choices': choices}
     python_code = (
         f'# this file is auto-generated by generate_self_schema.py, DO NOT edit manually\nself_schema = {schema}\n'
     )

pydantic_core/core_schema.py

@@ -1094,8 +1094,7 @@ def list_schema(
 
 
 class TuplePositionalSchema(TypedDict, total=False):
-    type: Required[Literal['tuple']]
-    mode: Required[Literal['positional']]
+    type: Required[Literal['tuple-positional']]
     items_schema: Required[List[CoreSchema]]
     extra_schema: CoreSchema
     strict: bool
@@ -1134,8 +1133,7 @@ def tuple_positional_schema(
         serialization: Custom serialization schema
     """
     return dict_not_none(
-        type='tuple',
-        mode='positional',
+        type='tuple-positional',
         items_schema=list(items_schema),
         extra_schema=extra_schema,
         strict=strict,
@@ -1146,8 +1144,7 @@ def tuple_positional_schema(
 
 
 class TupleVariableSchema(TypedDict, total=False):
-    type: Required[Literal['tuple']]
-    mode: Literal['variable']
+    type: Required[Literal['tuple-variable']]
     items_schema: CoreSchema
     min_length: int
     max_length: int
@@ -1187,8 +1184,7 @@ def tuple_variable_schema(
         serialization: Custom serialization schema
     """
     return dict_not_none(
-        type='tuple',
-        mode='variable',
+        type='tuple-variable',
         items_schema=items_schema,
         min_length=min_length,
         max_length=max_length,
@@ -1449,24 +1445,30 @@ def __call__(self, __input_value: Any, __info: ValidationInfo) -> Any:  # pragma
         ...
 
 
-class FunctionSchema(TypedDict, total=False):
-    type: Required[Literal['function']]
-    mode: Required[Literal['before', 'after']]
+class _FunctionSchema(TypedDict, total=False):
     function: Required[ValidatorFunction]
     schema: Required[CoreSchema]
     ref: str
     metadata: Any
     serialization: SerSchema
 
 
+class FunctionBeforeSchema(_FunctionSchema, total=False):
+    type: Required[Literal['function-before']]
+
+
+class FunctionAfterSchema(_FunctionSchema, total=False):
+    type: Required[Literal['function-after']]
+
+
 def function_before_schema(
     function: ValidatorFunction,
     schema: CoreSchema,
     *,
     ref: str | None = None,
     metadata: Any = None,
     serialization: SerSchema | None = None,
-) -> FunctionSchema:
+) -> FunctionBeforeSchema:
     """
     Returns a schema that calls a validator function before validating the provided schema, e.g.:
 
@@ -1492,8 +1494,7 @@ def fn(v: Any, info: core_schema.ValidationInfo) -> str:
         serialization: Custom serialization schema
     """
     return dict_not_none(
-        type='function',
-        mode='before',
+        type='function-before',
         function=function,
         schema=schema,
         ref=ref,
@@ -1509,7 +1510,7 @@ def function_after_schema(
     ref: str | None = None,
     metadata: Any = None,
     serialization: SerSchema | None = None,
-) -> FunctionSchema:
+) -> FunctionAfterSchema:
     """
     Returns a schema that calls a validator function after validating the provided schema, e.g.:
 
@@ -1533,13 +1534,7 @@ def fn(v: str, info: core_schema.ValidationInfo) -> str:
         serialization: Custom serialization schema
     """
     return dict_not_none(
-        type='function',
-        mode='after',
-        function=function,
-        schema=schema,
-        ref=ref,
-        metadata=metadata,
-        serialization=serialization,
+        type='function-after', function=function, schema=schema, ref=ref, metadata=metadata, serialization=serialization
     )
 
 
@@ -1556,8 +1551,7 @@ def __call__(
 
 
 class FunctionWrapSchema(TypedDict, total=False):
-    type: Required[Literal['function']]
-    mode: Required[Literal['wrap']]
+    type: Required[Literal['function-wrap']]
     function: Required[WrapValidatorFunction]
     schema: Required[CoreSchema]
     ref: str
@@ -1597,19 +1591,12 @@ def fn(v: str, validator: core_schema.CallableValidator, info: core_schema.Valid
         serialization: Custom serialization schema
     """
     return dict_not_none(
-        type='function',
-        mode='wrap',
-        function=function,
-        schema=schema,
-        ref=ref,
-        metadata=metadata,
-        serialization=serialization,
+        type='function-wrap', function=function, schema=schema, ref=ref, metadata=metadata, serialization=serialization
     )
 
 
 class FunctionPlainSchema(TypedDict, total=False):
-    type: Required[Literal['function']]
-    mode: Required[Literal['plain']]
+    type: Required[Literal['function-plain']]
     function: Required[ValidatorFunction]
     ref: str
     metadata: Any
@@ -1641,7 +1628,7 @@ def fn(v: str, info: core_schema.ValidationInfo) -> str:
         serialization: Custom serialization schema
     """
     return dict_not_none(
-        type='function', mode='plain', function=function, ref=ref, metadata=metadata, serialization=serialization
+        type='function-plain', function=function, ref=ref, metadata=metadata, serialization=serialization
     )
 
 
@@ -2795,7 +2782,8 @@ def definition_reference_schema(
     FrozenSetSchema,
     GeneratorSchema,
     DictSchema,
-    FunctionSchema,
+    FunctionBeforeSchema,
+    FunctionAfterSchema,
     FunctionWrapSchema,
     FunctionPlainSchema,
     WithDefaultSchema,
@@ -2836,12 +2824,16 @@ def definition_reference_schema(
     'is-subclass',
     'callable',
     'list',
-    'tuple',
+    'tuple-positional',
+    'tuple-variable',
     'set',
     'frozenset',
     'generator',
     'dict',
-    'function',
+    'function-before',
+    'function-after',
+    'function-wrap',
+    'function-plain',
     'default',
     'nullable',
     'union',

src/serializers/shared.rs

@@ -78,18 +78,12 @@ combined_serializer! {
         // hence they're here.
         Function: super::type_serializers::function::FunctionPlainSerializer;
         FunctionWrap: super::type_serializers::function::FunctionWrapSerializer;
-        // `TuplePositionalSerializer` & `TupleVariableSerializer` are created by
-        // `TupleBuilder` based on the `mode` parameter.
-        TuplePositional: super::type_serializers::tuple::TuplePositionalSerializer;
-        TupleVariable: super::type_serializers::tuple::TupleVariableSerializer;
     }
     // `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
     find_only: {
-        super::type_serializers::tuple::TupleBuilder;
         super::type_serializers::union::TaggedUnionBuilder;
         super::type_serializers::other::ChainBuilder;
-        super::type_serializers::other::FunctionBuilder;
         super::type_serializers::other::CustomErrorBuilder;
         super::type_serializers::other::CallBuilder;
         super::type_serializers::other::LaxOrStrictBuilder;
@@ -100,6 +94,10 @@ combined_serializer! {
         super::type_serializers::definitions::DefinitionsBuilder;
         super::type_serializers::dataclass::DataclassArgsBuilder;
         super::type_serializers::dataclass::DataclassBuilder;
+        super::type_serializers::function::FunctionBeforeSerializerBuilder;
+        super::type_serializers::function::FunctionAfterSerializerBuilder;
+        super::type_serializers::function::FunctionPlainSerializerBuilder;
+        super::type_serializers::function::FunctionWrapSerializerBuilder;
     }
     // `both` means the struct is added to both the `CombinedSerializer` enum and the match statement in
     // `find_serializer` so they can be used via a `type` str.
@@ -132,6 +130,8 @@ combined_serializer! {
         Union: super::type_serializers::union::UnionSerializer;
         Literal: super::type_serializers::literal::LiteralSerializer;
         Recursive: super::type_serializers::definitions::DefinitionRefSerializer;
+        TuplePositional: super::type_serializers::tuple::TuplePositionalSerializer;
+        TupleVariable: super::type_serializers::tuple::TupleVariableSerializer;
     }
 }
 
@@ -150,11 +150,15 @@ impl CombinedSerializer {
                 Some("function-plain") => {
                     // `function` is a special case, not included in `find_serializer` since it means something
                     // different in `schema.type`
-                    return super::type_serializers::function::FunctionPlainSerializer::new_combined(ser_schema)
-                        .map_err(|err| py_error_type!("Error building `function-plain` serializer:\n  {}", err));
+                    return super::type_serializers::function::FunctionPlainSerializer::build(
+                        ser_schema,
+                        config,
+                        build_context,
+                    )
+                    .map_err(|err| py_error_type!("Error building `function-plain` serializer:\n  {}", err));
                 }
                 Some("function-wrap") => {
-                    return super::type_serializers::function::FunctionWrapSerializer::new_combined(
+                    return super::type_serializers::function::FunctionWrapSerializer::build(
                         ser_schema,
                         config,
                         build_context,

src/serializers/type_serializers/function.rs

@@ -21,6 +21,54 @@ use super::{
     PydanticSerializationError, TypeSerializer,
 };
 
+pub struct FunctionBeforeSerializerBuilder;
+
+impl BuildSerializer for FunctionBeforeSerializerBuilder {
+    const EXPECTED_TYPE: &'static str = "function-before";
+
+    fn build(
+        schema: &PyDict,
+        config: Option<&PyDict>,
+        build_context: &mut BuildContext<CombinedSerializer>,
+    ) -> PyResult<CombinedSerializer> {
+        let py = schema.py();
+        // `before` schemas will obviously have type from `schema` since the validator is called second
+        let schema = schema.get_as_req(intern!(py, "schema"))?;
+        CombinedSerializer::build(schema, config, build_context)
+    }
+}
+
+pub struct FunctionAfterSerializerBuilder;
+
+impl BuildSerializer for FunctionAfterSerializerBuilder {
+    const EXPECTED_TYPE: &'static str = "function-after";
+    fn build(
+        schema: &PyDict,
+        config: Option<&PyDict>,
+        build_context: &mut BuildContext<CombinedSerializer>,
+    ) -> PyResult<CombinedSerializer> {
+        let py = schema.py();
+        // while `before` schemas have an obvious type, for
+        // `after` schemas it's less, clear but the default will be the same type, and the user/lib can always
+        // override the serializer
+        let schema = schema.get_as_req(intern!(py, "schema"))?;
+        CombinedSerializer::build(schema, config, build_context)
+    }
+}
+
+pub struct FunctionPlainSerializerBuilder;
+
+impl BuildSerializer for FunctionPlainSerializerBuilder {
+    const EXPECTED_TYPE: &'static str = "function-plain";
+    fn build(
+        schema: &PyDict,
+        config: Option<&PyDict>,
+        build_context: &mut BuildContext<CombinedSerializer>,
+    ) -> PyResult<CombinedSerializer> {
+        super::any::AnySerializer::build(schema, config, build_context)
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct FunctionPlainSerializer {
     func: PyObject,
@@ -30,8 +78,13 @@ pub struct FunctionPlainSerializer {
     when_used: WhenUsed,
 }
 
-impl FunctionPlainSerializer {
-    pub fn new_combined(schema: &PyDict) -> PyResult<CombinedSerializer> {
+impl BuildSerializer for FunctionPlainSerializer {
+    const EXPECTED_TYPE: &'static str = "function-plain";
+    fn build(
+        schema: &PyDict,
+        _config: Option<&PyDict>,
+        _build_context: &mut BuildContext<CombinedSerializer>,
+    ) -> PyResult<CombinedSerializer> {
         let py = schema.py();
         let function = schema.get_as_req::<&PyAny>(intern!(py, "function"))?;
         let function_name = function_name(function)?;
@@ -46,7 +99,9 @@ impl FunctionPlainSerializer {
         }
         .into())
     }
+}
 
+impl FunctionPlainSerializer {
     fn call(
         &self,
         value: &PyAny,
@@ -177,6 +232,19 @@ macro_rules! function_type_serializer {
 
 function_type_serializer!(FunctionPlainSerializer);
 
+pub struct FunctionWrapSerializerBuilder;
+
+impl BuildSerializer for FunctionWrapSerializerBuilder {
+    const EXPECTED_TYPE: &'static str = "function-wrap";
+    fn build(
+        schema: &PyDict,
+        config: Option<&PyDict>,
+        build_context: &mut BuildContext<CombinedSerializer>,
+    ) -> PyResult<CombinedSerializer> {
+        super::any::AnySerializer::build(schema, config, build_context)
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct FunctionWrapSerializer {
     serializer: Box<CombinedSerializer>,
@@ -187,8 +255,9 @@ pub struct FunctionWrapSerializer {
     when_used: WhenUsed,
 }
 
-impl FunctionWrapSerializer {
-    pub fn new_combined(
+impl BuildSerializer for FunctionWrapSerializer {
+    const EXPECTED_TYPE: &'static str = "function-wrap";
+    fn build(
         schema: &PyDict,
         config: Option<&PyDict>,
         build_context: &mut BuildContext<CombinedSerializer>,
@@ -211,7 +280,9 @@ impl FunctionWrapSerializer {
         }
         .into())
     }
+}
 
+impl FunctionWrapSerializer {
     fn call(
         &self,
         value: &PyAny,