Flatten CoreSchema types to get a single discriminant key (#450)

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

@@ -1139,8 +1139,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
@@ -1179,8 +1178,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,
@@ -1191,8 +1189,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
@@ -1232,8 +1229,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,
@@ -1509,24 +1505,26 @@ class GeneralValidatorFunctionSchema(TypedDict):
     function: GeneralValidatorFunction
 
 
-class FunctionSchema(TypedDict, total=False):
-    type: Required[Literal['function']]
+class _FunctionSchema(TypedDict, total=False):
     function: Required[Union[FieldValidatorFunctionSchema, GeneralValidatorFunctionSchema]]
-    mode: Required[Literal['before', 'after']]
     schema: Required[CoreSchema]
     ref: str
     metadata: Any
     serialization: SerSchema
 
 
+class FunctionBeforeSchema(_FunctionSchema, total=False):
+    type: Required[Literal['function-before']]
+
+
 def field_before_validation_function(
     function: FieldValidatorFunction,
     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 **model field** schema, e.g.:
@@ -1556,8 +1554,7 @@ def fn(v: bytes, info: core_schema.ModelFieldValidationInfo) -> str:
         serialization: Custom serialization schema
     """
     return dict_not_none(
-        type='function',
-        mode='before',
+        type='function-before',
         function={'type': 'field', 'function': function},
         schema=schema,
         ref=ref,
@@ -1573,7 +1570,7 @@ def general_before_validation_function(
     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.:
 
@@ -1599,8 +1596,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={'type': 'general', 'function': function},
         schema=schema,
         ref=ref,
@@ -1609,14 +1605,18 @@ def fn(v: Any, info: core_schema.ValidationInfo) -> str:
     )
 
 
+class FunctionAfterSchema(_FunctionSchema, total=False):
+    type: Required[Literal['function-after']]
+
+
 def field_after_validation_function(
     function: FieldValidatorFunction,
     schema: CoreSchema,
     *,
     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 **model field** schema, e.g.:
@@ -1646,8 +1646,7 @@ def fn(v: str, info: core_schema.ModelFieldValidationInfo) -> str:
         serialization: Custom serialization schema
     """
     return dict_not_none(
-        type='function',
-        mode='after',
+        type='function-after',
         function={'type': 'field', 'function': function},
         schema=schema,
         ref=ref,
@@ -1663,7 +1662,7 @@ def general_after_validation_function(
     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.:
 
@@ -1687,8 +1686,7 @@ def fn(v: str, info: core_schema.ValidationInfo) -> str:
         serialization: Custom serialization schema
     """
     return dict_not_none(
-        type='function',
-        mode='after',
+        type='function-after',
         function={'type': 'general', 'function': function},
         schema=schema,
         ref=ref,
@@ -1727,9 +1725,8 @@ class GeneralWrapValidatorFunctionSchema(TypedDict):
 
 
 class WrapFunctionSchema(TypedDict, total=False):
-    type: Required[Literal['function']]
+    type: Required[Literal['function-wrap']]
     function: Required[Union[GeneralWrapValidatorFunctionSchema, FieldWrapValidatorFunctionSchema]]
-    mode: Required[Literal['wrap']]
     schema: Required[CoreSchema]
     ref: str
     metadata: Any
@@ -1768,8 +1765,7 @@ def fn(v: str, validator: core_schema.CallableValidator, info: core_schema.Valid
         serialization: Custom serialization schema
     """
     return dict_not_none(
-        type='function',
-        mode='wrap',
+        type='function-wrap',
         function={'type': 'general', 'function': function},
         schema=schema,
         ref=ref,
@@ -1817,8 +1813,7 @@ def fn(v: bytes, validator: core_schema.CallableValidator, info: core_schema.Mod
         serialization: Custom serialization schema
     """
     return dict_not_none(
-        type='function',
-        mode='wrap',
+        type='function-wrap',
         function={'type': 'field', 'function': function},
         schema=schema,
         ref=ref,
@@ -1828,8 +1823,7 @@ def fn(v: bytes, validator: core_schema.CallableValidator, info: core_schema.Mod
 
 
 class PlainFunctionSchema(TypedDict, total=False):
-    type: Required[Literal['function']]
-    mode: Required[Literal['plain']]
+    type: Required[Literal['function-plain']]
     function: Required[Union[FieldValidatorFunctionSchema, GeneralValidatorFunctionSchema]]
     ref: str
     metadata: Any
@@ -1865,8 +1859,7 @@ def fn(v: str, info: core_schema.ValidationInfo) -> str:
         serialization: Custom serialization schema
     """
     return dict_not_none(
-        type='function',
-        mode='plain',
+        type='function-plain',
         function={'type': 'general', 'function': function},
         ref=ref,
         metadata=metadata,
@@ -1909,8 +1902,7 @@ def fn(v: Any, info: core_schema.ModelFieldValidationInfo) -> str:
         serialization: Custom serialization schema
     """
     return dict_not_none(
-        type='function',
-        mode='plain',
+        type='function-plain',
         function={'type': 'field', 'function': function},
         ref=ref,
         metadata=metadata,
@@ -3068,7 +3060,8 @@ def definition_reference_schema(
     FrozenSetSchema,
     GeneratorSchema,
     DictSchema,
-    FunctionSchema,
+    FunctionAfterSchema,
+    FunctionBeforeSchema,
     WrapFunctionSchema,
     PlainFunctionSchema,
     WithDefaultSchema,
@@ -3109,12 +3102,16 @@ def definition_reference_schema(
     'is-subclass',
     'callable',
     'list',
-    'tuple',
+    'tuple-positional',
+    'tuple-variable',
     'set',
     'frozenset',
     'generator',
     'dict',
-    'function',
+    'function-after',
+    'function-before',
+    '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,