Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1from __future__ import annotations as _annotations
3from typing import TYPE_CHECKING, Any, cast
5from pydantic_core import CoreSchema, core_schema
7from ..errors import PydanticUserError
8from . import _core_utils
9from ._core_utils import (
10 CoreSchemaField,
11)
13if TYPE_CHECKING:
14 from ..types import Discriminator
15 from ._core_metadata import CoreMetadata
18class MissingDefinitionForUnionRef(Exception):
19 """Raised when applying a discriminated union discriminator to a schema
20 requires a definition that is not yet defined
21 """
23 def __init__(self, ref: str) -> None:
24 self.ref = ref
25 super().__init__(f'Missing definition for ref {self.ref!r}')
28def set_discriminator_in_metadata(schema: CoreSchema, discriminator: Any) -> None:
29 metadata = cast('CoreMetadata', schema.setdefault('metadata', {}))
30 metadata['pydantic_internal_union_discriminator'] = discriminator
33def apply_discriminator(
34 schema: core_schema.CoreSchema,
35 discriminator: str | Discriminator,
36 definitions: dict[str, core_schema.CoreSchema] | None = None,
37) -> core_schema.CoreSchema:
38 """Applies the discriminator and returns a new core schema.
40 Args:
41 schema: The input schema.
42 discriminator: The name of the field which will serve as the discriminator.
43 definitions: A mapping of schema ref to schema.
45 Returns:
46 The new core schema.
48 Raises:
49 TypeError:
50 - If `discriminator` is used with invalid union variant.
51 - If `discriminator` is used with union type with one variant.
52 - If `discriminator` value mapped to multiple choices.
53 MissingDefinitionForUnionRef:
54 If the definition for ref is missing.
55 PydanticUserError:
56 - If a model in union doesn't have a discriminator field.
57 - If discriminator field has a non-string alias.
58 - If discriminator fields have different aliases.
59 - If discriminator field not of type `Literal`.
60 """
61 from ..types import Discriminator
63 if isinstance(discriminator, Discriminator):
64 if isinstance(discriminator.discriminator, str):
65 discriminator = discriminator.discriminator
66 else:
67 return discriminator._convert_schema(schema, definitions=definitions)
69 return _ApplyInferredDiscriminator(discriminator, definitions or {}).apply(schema)
72class _ApplyInferredDiscriminator:
73 """This class is used to convert an input schema containing a union schema into one where that union is
74 replaced with a tagged-union, with all the associated debugging and performance benefits.
76 This is done by:
77 * Validating that the input schema is compatible with the provided discriminator
78 * Introspecting the schema to determine which discriminator values should map to which union choices
79 * Handling various edge cases such as 'definitions', 'default', 'nullable' schemas, and more
81 I have chosen to implement the conversion algorithm in this class, rather than a function,
82 to make it easier to maintain state while recursively walking the provided CoreSchema.
83 """
85 def __init__(self, discriminator: str, definitions: dict[str, core_schema.CoreSchema]):
86 # `discriminator` should be the name of the field which will serve as the discriminator.
87 # It must be the python name of the field, and *not* the field's alias. Note that as of now,
88 # all members of a discriminated union _must_ use a field with the same name as the discriminator.
89 # This may change if/when we expose a way to manually specify the TaggedUnionSchema's choices.
90 self.discriminator = discriminator
92 # `definitions` should contain a mapping of schema ref to schema for all schemas which might
93 # be referenced by some choice
94 self.definitions = definitions
96 # `_discriminator_alias` will hold the value, if present, of the alias for the discriminator
97 #
98 # Note: following the v1 implementation, we currently disallow the use of different aliases
99 # for different choices. This is not a limitation of pydantic_core, but if we try to handle
100 # this, the inference logic gets complicated very quickly, and could result in confusing
101 # debugging challenges for users making subtle mistakes.
102 #
103 # Rather than trying to do the most powerful inference possible, I think we should eventually
104 # expose a way to more-manually control the way the TaggedUnionSchema is constructed through
105 # the use of a new type which would be placed as an Annotation on the Union type. This would
106 # provide the full flexibility/power of pydantic_core's TaggedUnionSchema where necessary for
107 # more complex cases, without over-complicating the inference logic for the common cases.
108 self._discriminator_alias: str | None = None
110 # `_should_be_nullable` indicates whether the converted union has `None` as an allowed value.
111 # If `None` is an acceptable value of the (possibly-wrapped) union, we ignore it while
112 # constructing the TaggedUnionSchema, but set the `_should_be_nullable` attribute to True.
113 # Once we have constructed the TaggedUnionSchema, if `_should_be_nullable` is True, we ensure
114 # that the final schema gets wrapped as a NullableSchema. This has the same semantics on the
115 # python side, but resolves the issue that `None` cannot correspond to any discriminator values.
116 self._should_be_nullable = False
118 # `_is_nullable` is used to track if the final produced schema will definitely be nullable;
119 # we set it to True if the input schema is wrapped in a nullable schema that we know will be preserved
120 # as an indication that, even if None is discovered as one of the union choices, we will not need to wrap
121 # the final value in another nullable schema.
122 #
123 # This is more complicated than just checking for the final outermost schema having type 'nullable' thanks
124 # to the possible presence of other wrapper schemas such as DefinitionsSchema, WithDefaultSchema, etc.
125 self._is_nullable = False
127 # `_should_allow_missing` and `_is_allowing_missing` play the same role as `_should_be_nullable`
128 # and `_is_nullable`, for the `MISSING` sentinel (represented by the 'missing-sentinel' wrapper schema).
129 self._should_allow_missing = False
130 self._is_allowing_missing = False
132 # `_choices_to_handle` serves as a stack of choices to add to the tagged union. Initially, choices
133 # from the union in the wrapped schema will be appended to this list, and the recursive choice-handling
134 # algorithm may add more choices to this stack as (nested) unions are encountered.
135 self._choices_to_handle: list[core_schema.CoreSchema] = []
137 # `_tagged_union_choices` is built during the call to `apply`, and will hold the choices to be included
138 # in the output TaggedUnionSchema that will replace the union from the input schema
139 self._tagged_union_choices: dict[Any, core_schema.CoreSchema] = {}
141 # `_used` is changed to True after applying the discriminator to prevent accidental reuse
142 self._used = False
144 def apply(self, schema: core_schema.CoreSchema) -> core_schema.CoreSchema:
145 """Return a new CoreSchema based on `schema` that uses a tagged-union with the discriminator provided
146 to this class.
148 Args:
149 schema: The input schema.
151 Returns:
152 The new core schema.
154 Raises:
155 TypeError:
156 - If `discriminator` is used with invalid union variant.
157 - If `discriminator` is used with union type with one variant.
158 - If `discriminator` value mapped to multiple choices.
159 ValueError:
160 If the definition for ref is missing.
161 PydanticUserError:
162 - If a model in union doesn't have a discriminator field.
163 - If discriminator field has a non-string alias.
164 - If discriminator fields have different aliases.
165 - If discriminator field not of type `Literal`.
166 """
167 assert not self._used
168 schema = self._apply_to_root(schema)
169 if self._should_be_nullable and not self._is_nullable:
170 schema = core_schema.nullable_schema(schema)
171 if self._should_allow_missing and not self._is_allowing_missing:
172 schema = core_schema.missing_sentinel_schema(schema)
173 self._used = True
174 return schema
176 def _apply_to_root(self, schema: core_schema.CoreSchema) -> core_schema.CoreSchema:
177 """This method handles the outer-most stage of recursion over the input schema:
178 unwrapping nullable or definitions schemas, and calling the `_handle_choice`
179 method iteratively on the choices extracted (recursively) from the possibly-wrapped union.
180 """
181 if schema['type'] == 'nullable':
182 self._is_nullable = True
183 wrapped = self._apply_to_root(schema['schema'])
184 nullable_wrapper = schema.copy()
185 nullable_wrapper['schema'] = wrapped
186 return nullable_wrapper
188 if schema['type'] == 'missing-sentinel' and 'schema' in schema:
189 self._is_allowing_missing = True
190 wrapped = self._apply_to_root(schema['schema'])
191 missing_sentinel_wrapper = schema.copy()
192 missing_sentinel_wrapper['schema'] = wrapped
193 return missing_sentinel_wrapper
195 if schema['type'] == 'definitions':
196 wrapped = self._apply_to_root(schema['schema'])
197 definitions_wrapper = schema.copy()
198 definitions_wrapper['schema'] = wrapped
199 return definitions_wrapper
201 if schema['type'] == 'definition-ref':
202 schema_ref = schema['schema_ref']
203 if schema_ref not in self.definitions: # pragma: no cover
204 raise MissingDefinitionForUnionRef(schema_ref)
206 def_schema = self.definitions[schema_ref]
207 # If using a referenceable union as discriminated (e.g. `type Pet = Cat | Dog; field: Pet = Field(discriminator=...)`):
208 if def_schema['type'] == 'union':
209 schema = def_schema.copy()
210 schema.pop('ref')
212 if schema['type'] != 'union':
213 # If the schema is not a union, it probably means it just had a single member and
214 # was flattened by pydantic_core.
215 # However, it still may make sense to apply the discriminator to this schema,
216 # as a way to get discriminated-union-style error messages, so we allow this here.
217 schema = core_schema.union_schema([schema])
219 # Reverse the choices list before extending the stack so that they get handled in the order they occur
220 choices_schemas = [v[0] if isinstance(v, tuple) else v for v in schema['choices'][::-1]]
221 self._choices_to_handle.extend(choices_schemas)
222 while self._choices_to_handle:
223 choice = self._choices_to_handle.pop()
224 self._handle_choice(choice)
226 if self._discriminator_alias is not None and self._discriminator_alias != self.discriminator:
227 # * We need to annotate `discriminator` as a union here to handle both branches of this conditional
228 # * We need to annotate `discriminator` as list[list[str | int]] and not list[list[str]] due to the
229 # invariance of list, and because list[list[str | int]] is the type of the discriminator argument
230 # to tagged_union_schema below
231 # * See the docstring of pydantic_core.core_schema.tagged_union_schema for more details about how to
232 # interpret the value of the discriminator argument to tagged_union_schema. (The list[list[str]] here
233 # is the appropriate way to provide a list of fallback attributes to check for a discriminator value.)
234 discriminator: str | list[list[str | int]] = [[self.discriminator], [self._discriminator_alias]]
235 else:
236 discriminator = self.discriminator
237 return core_schema.tagged_union_schema(
238 choices=self._tagged_union_choices,
239 discriminator=discriminator,
240 custom_error_type=schema.get('custom_error_type'),
241 custom_error_message=schema.get('custom_error_message'),
242 custom_error_context=schema.get('custom_error_context'),
243 strict=False,
244 from_attributes=True,
245 ref=schema.get('ref'),
246 metadata=schema.get('metadata'),
247 serialization=schema.get('serialization'),
248 )
250 def _handle_choice(self, choice: core_schema.CoreSchema) -> None:
251 """This method handles the "middle" stage of recursion over the input schema.
252 Specifically, it is responsible for handling each choice of the outermost union
253 (and any "coalesced" choices obtained from inner unions).
255 Here, "handling" entails:
256 * Coalescing nested unions and compatible tagged-unions
257 * Tracking the presence of 'none' and 'nullable' schemas occurring as choices
258 * Validating that each allowed discriminator value maps to a unique choice
259 * Updating the _tagged_union_choices mapping that will ultimately be used to build the TaggedUnionSchema.
260 """
261 if choice['type'] == 'definition-ref':
262 if choice['schema_ref'] not in self.definitions:
263 raise MissingDefinitionForUnionRef(choice['schema_ref'])
265 if choice['type'] == 'none':
266 self._should_be_nullable = True
267 elif choice['type'] == 'definitions':
268 self._handle_choice(choice['schema'])
269 elif choice['type'] == 'nullable':
270 self._should_be_nullable = True
271 self._handle_choice(choice['schema']) # unwrap the nullable schema
272 elif choice['type'] == 'missing-sentinel':
273 self._should_allow_missing = True
274 if (inner_schema := choice.get('schema')) is not None:
275 self._handle_choice(inner_schema) # unwrap the missing-sentinel schema
276 elif choice['type'] == 'union':
277 # Reverse the choices list before extending the stack so that they get handled in the order they occur
278 choices_schemas = [v[0] if isinstance(v, tuple) else v for v in choice['choices'][::-1]]
279 self._choices_to_handle.extend(choices_schemas)
280 elif choice['type'] not in {
281 'model',
282 'typed-dict',
283 'tagged-union',
284 'lax-or-strict',
285 'dataclass',
286 'dataclass-args',
287 'definition-ref',
288 } and not _core_utils.is_function_with_inner_schema(choice):
289 # We should eventually handle 'definition-ref' as well
290 err_str = f'The core schema type {choice["type"]!r} is not a valid discriminated union variant.'
291 if choice['type'] == 'list':
292 err_str += (
293 ' If you are making use of a list of union types, make sure the discriminator is applied to the '
294 'union type and not the list (e.g. `list[Annotated[<T> | <U>, Field(discriminator=...)]]`).'
295 )
296 raise TypeError(err_str)
297 else:
298 if choice['type'] == 'tagged-union' and self._is_discriminator_shared(choice):
299 # In this case, this inner tagged-union is compatible with the outer tagged-union,
300 # and its choices can be coalesced into the outer TaggedUnionSchema.
301 subchoices = [x for x in choice['choices'].values() if not isinstance(x, (str, int))]
302 # Reverse the choices list before extending the stack so that they get handled in the order they occur
303 self._choices_to_handle.extend(subchoices[::-1])
304 return
306 inferred_discriminator_values = self._infer_discriminator_values_for_choice(choice, source_name=None)
307 self._set_unique_choice_for_values(choice, inferred_discriminator_values)
309 def _is_discriminator_shared(self, choice: core_schema.TaggedUnionSchema) -> bool:
310 """This method returns a boolean indicating whether the discriminator for the `choice`
311 is the same as that being used for the outermost tagged union. This is used to
312 determine whether this TaggedUnionSchema choice should be "coalesced" into the top level,
313 or whether it should be treated as a separate (nested) choice.
314 """
315 inner_discriminator = choice['discriminator']
316 return inner_discriminator == self.discriminator or (
317 isinstance(inner_discriminator, list)
318 and (self.discriminator in inner_discriminator or [self.discriminator] in inner_discriminator)
319 )
321 def _infer_discriminator_values_for_choice( # noqa C901
322 self, choice: core_schema.CoreSchema, source_name: str | None
323 ) -> list[Any]:
324 """This function recurses over `choice`, extracting all discriminator values that should map to this choice.
326 `model_name` is accepted for the purpose of producing useful error messages.
327 """
328 if choice['type'] == 'definitions':
329 return self._infer_discriminator_values_for_choice(choice['schema'], source_name=source_name)
331 elif _core_utils.is_function_with_inner_schema(choice):
332 return self._infer_discriminator_values_for_choice(choice['schema'], source_name=source_name)
334 elif choice['type'] == 'lax-or-strict':
335 return sorted(
336 set(
337 self._infer_discriminator_values_for_choice(choice['lax_schema'], source_name=None)
338 + self._infer_discriminator_values_for_choice(choice['strict_schema'], source_name=None)
339 )
340 )
342 elif choice['type'] == 'tagged-union':
343 values: list[Any] = []
344 # Ignore str/int "choices" since these are just references to other choices
345 subchoices = [x for x in choice['choices'].values() if not isinstance(x, (str, int))]
346 for subchoice in subchoices:
347 subchoice_values = self._infer_discriminator_values_for_choice(subchoice, source_name=None)
348 values.extend(subchoice_values)
349 return values
351 elif choice['type'] == 'union':
352 values = []
353 for subchoice in choice['choices']:
354 subchoice_schema = subchoice[0] if isinstance(subchoice, tuple) else subchoice
355 subchoice_values = self._infer_discriminator_values_for_choice(subchoice_schema, source_name=None)
356 values.extend(subchoice_values)
357 return values
359 elif choice['type'] == 'nullable':
360 self._should_be_nullable = True
361 return self._infer_discriminator_values_for_choice(choice['schema'], source_name=None)
363 elif choice['type'] == 'missing-sentinel' and (inner_schema := choice.get('schema')) is not None:
364 self._should_allow_missing = True
365 return self._infer_discriminator_values_for_choice(inner_schema, source_name=None)
367 elif choice['type'] == 'model':
368 return self._infer_discriminator_values_for_choice(choice['schema'], source_name=choice['cls'].__name__)
370 elif choice['type'] == 'dataclass':
371 return self._infer_discriminator_values_for_choice(choice['schema'], source_name=choice['cls'].__name__)
373 elif choice['type'] == 'model-fields':
374 return self._infer_discriminator_values_for_model_choice(choice, source_name=source_name)
376 elif choice['type'] == 'dataclass-args':
377 return self._infer_discriminator_values_for_dataclass_choice(choice, source_name=source_name)
379 elif choice['type'] == 'typed-dict':
380 return self._infer_discriminator_values_for_typed_dict_choice(choice, source_name=source_name)
382 elif choice['type'] == 'definition-ref':
383 schema_ref = choice['schema_ref']
384 if schema_ref not in self.definitions:
385 raise MissingDefinitionForUnionRef(schema_ref)
386 return self._infer_discriminator_values_for_choice(self.definitions[schema_ref], source_name=source_name)
387 else:
388 err_str = f'The core schema type {choice["type"]!r} is not a valid discriminated union variant.'
389 if choice['type'] == 'list':
390 err_str += (
391 ' If you are making use of a list of union types, make sure the discriminator is applied to the '
392 'union type and not the list (e.g. `list[Annotated[<T> | <U>, Field(discriminator=...)]]`).'
393 )
394 raise TypeError(err_str)
396 def _infer_discriminator_values_for_typed_dict_choice(
397 self, choice: core_schema.TypedDictSchema, source_name: str | None = None
398 ) -> list[Any]:
399 """This method just extracts the _infer_discriminator_values_for_choice logic specific to TypedDictSchema
400 for the sake of readability.
401 """
402 source = 'TypedDict' if source_name is None else f'TypedDict {source_name!r}'
403 field = choice['fields'].get(self.discriminator)
404 if field is None:
405 raise PydanticUserError(
406 f'{source} needs a discriminator field for key {self.discriminator!r}', code='discriminator-no-field'
407 )
408 return self._infer_discriminator_values_for_field(field, source)
410 def _infer_discriminator_values_for_model_choice(
411 self, choice: core_schema.ModelFieldsSchema, source_name: str | None = None
412 ) -> list[Any]:
413 source = 'ModelFields' if source_name is None else f'Model {source_name!r}'
414 field = choice['fields'].get(self.discriminator)
415 if field is None:
416 raise PydanticUserError(
417 f'{source} needs a discriminator field for key {self.discriminator!r}', code='discriminator-no-field'
418 )
419 return self._infer_discriminator_values_for_field(field, source)
421 def _infer_discriminator_values_for_dataclass_choice(
422 self, choice: core_schema.DataclassArgsSchema, source_name: str | None = None
423 ) -> list[Any]:
424 source = 'DataclassArgs' if source_name is None else f'Dataclass {source_name!r}'
425 for field in choice['fields']:
426 if field['name'] == self.discriminator:
427 break
428 else:
429 raise PydanticUserError(
430 f'{source} needs a discriminator field for key {self.discriminator!r}', code='discriminator-no-field'
431 )
432 return self._infer_discriminator_values_for_field(field, source)
434 def _infer_discriminator_values_for_field(self, field: CoreSchemaField, source: str) -> list[Any]:
435 if field['type'] == 'computed-field':
436 # This should never occur as a discriminator, as it is only relevant to serialization
437 return []
438 alias = field.get('validation_alias', self.discriminator)
439 if not isinstance(alias, str):
440 raise PydanticUserError(
441 f'Alias {alias!r} is not supported in a discriminated union', code='discriminator-alias-type'
442 )
443 if self._discriminator_alias is None:
444 self._discriminator_alias = alias
445 elif self._discriminator_alias != alias:
446 raise PydanticUserError(
447 f'Aliases for discriminator {self.discriminator!r} must be the same '
448 f'(got {alias}, {self._discriminator_alias})',
449 code='discriminator-alias',
450 )
451 return self._infer_discriminator_values_for_inner_schema(field['schema'], source)
453 def _infer_discriminator_values_for_inner_schema(self, schema: core_schema.CoreSchema, source: str) -> list[Any]:
454 """When inferring discriminator values for a field, we typically extract the expected values from a literal
455 schema. This function does that, but also handles nested unions and defaults.
456 """
457 match schema:
458 case {'type': 'literal', 'expected': expected}:
459 return expected
460 case {'type': 'none'}:
461 return [None]
462 case {'type': 'union', 'choices': choices}:
463 # Generally when multiple values are allowed they should be placed in a single `Literal`, but
464 # we add this case to handle the situation where a field is annotated as a `Union` of `Literal`s.
465 # For example, this lets us handle `Union[Literal['key'], Union[Literal['Key'], Literal['KEY']]]`
466 values: list[Any] = []
467 for choice in choices:
468 choice_schema = choice[0] if isinstance(choice, tuple) else choice
469 choice_values = self._infer_discriminator_values_for_inner_schema(choice_schema, source)
470 values.extend(choice_values)
471 return values
472 case {'type': 'default', 'schema': inner_schema}:
473 # This will happen if the field has a default value; we ignore it while extracting the discriminator values
474 return self._infer_discriminator_values_for_inner_schema(inner_schema, source)
475 case {'type': 'function-after', 'schema': inner_schema}:
476 # After validators don't affect the discriminator values
477 return self._infer_discriminator_values_for_inner_schema(inner_schema, source)
478 case {'type': 'model', 'root_model': True, 'schema': inner_schema}:
479 # Support RootModel[Literal[...]] as discriminator field type
480 return self._infer_discriminator_values_for_inner_schema(inner_schema, source)
481 case {'type': 'function-before' | 'function-wrap' | 'function-plain' as schema_type}:
482 validator_type = repr(schema_type.split('-')[1])
483 raise PydanticUserError(
484 f'Cannot use a mode={validator_type} validator in the'
485 f' discriminator field {self.discriminator!r} of {source}',
486 code='discriminator-validator',
487 )
488 case _:
489 raise PydanticUserError(
490 f'{source} needs field {self.discriminator!r} to be of type `Literal`',
491 code='discriminator-needs-literal',
492 )
494 def _set_unique_choice_for_values(self, choice: core_schema.CoreSchema, values: list[Any]) -> None:
495 """This method updates `self.tagged_union_choices` so that all provided (discriminator) `values` map to the
496 provided `choice`, validating that none of these values already map to another (different) choice.
497 """
498 for discriminator_value in values:
499 if discriminator_value in self._tagged_union_choices:
500 # It is okay if `value` is already in tagged_union_choices as long as it maps to the same value.
501 # Because tagged_union_choices may map values to other values, we need to walk the choices dict
502 # until we get to a "real" choice, and confirm that is equal to the one assigned.
503 existing_choice = self._tagged_union_choices[discriminator_value]
504 if existing_choice != choice:
505 raise TypeError(
506 f'Value {discriminator_value!r} for discriminator '
507 f'{self.discriminator!r} mapped to multiple choices'
508 )
509 else:
510 self._tagged_union_choices[discriminator_value] = choice