Support discriminated union

changes/619-PrettyWood.md

@@ -0,0 +1 @@
+Add a discriminated union. See [the doc](https://pydantic-docs.helpmanual.io/usage/types/#discriminated-unions) for more information.

docs/examples/schema_ad_hoc.py

@@ -0,0 +1,20 @@
+from typing import Literal, Union
+
+from typing_extensions import Annotated
+
+from pydantic import BaseModel, Field, schema_json
+
+
+class Cat(BaseModel):
+    pet_type: Literal['cat']
+    cat_name: str
+
+
+class Dog(BaseModel):
+    pet_type: Literal['dog']
+    dog_name: str
+
+
+Pet = Annotated[Union[Cat, Dog], Field(discriminator='pet_type')]
+
+print(schema_json(Pet, title='The Pet Schema', indent=2))

docs/examples/types_union_discriminated.py

@@ -0,0 +1,30 @@
+from typing import Literal, Union
+
+from pydantic import BaseModel, Field, ValidationError
+
+
+class Cat(BaseModel):
+    pet_type: Literal['cat']
+    name: str
+
+
+class Dog(BaseModel):
+    pet_type: Literal['dog']
+    name: str
+
+
+class Lizard(BaseModel):
+    pet_type: Literal['reptile', 'lizard']
+    name: str
+
+
+class Model(BaseModel):
+    pet: Union[Cat, Dog, Lizard] = Field(..., discriminator='pet_type')
+    n: int
+
+
+print(Model(pet={'pet_type': 'dog', 'name': 'woof'}, n='1'))
+try:
+    Model(pet={'pet_type': 'dog'}, n='1')
+except ValidationError as e:
+    print(e)

docs/examples/types_union_discriminated_nested.py

@@ -0,0 +1,50 @@
+from typing import Literal, Union
+
+from typing_extensions import Annotated
+
+from pydantic import BaseModel, Field, ValidationError
+
+
+class BlackCat(BaseModel):
+    pet_type: Literal['cat']
+    color: Literal['black']
+    black_name: str
+
+
+class WhiteCat(BaseModel):
+    pet_type: Literal['cat']
+    color: Literal['white']
+    white_name: str
+
+
+# Can also be written with a custom root type
+#
+# class Cat(BaseModel):
+#   __root__: Annotated[Union[BlackCat, WhiteCat], Field(discriminator='color')]
+
+Cat = Annotated[Union[BlackCat, WhiteCat], Field(discriminator='color')]
+
+
+class Dog(BaseModel):
+    pet_type: Literal['dog']
+    name: str
+
+
+Pet = Annotated[Union[Cat, Dog], Field(discriminator='pet_type')]
+
+
+class Model(BaseModel):
+    pet: Pet
+    n: int
+
+
+m = Model(pet={'pet_type': 'cat', 'color': 'black', 'black_name': 'felix'}, n=1)
+print(m)
+try:
+    Model(pet={'pet_type': 'cat', 'color': 'red'}, n='1')
+except ValidationError as e:
+    print(e)
+try:
+    Model(pet={'pet_type': 'cat', 'color': 'black'}, n='1')
+except ValidationError as e:
+    print(e)

docs/usage/schema.md

@@ -37,6 +37,19 @@ The format of `$ref`s (`"#/definitions/FooBar"` above) can be altered by calling
 with the `ref_template` keyword argument, e.g. `ApplePie.schema(ref_template='/schemas/{model}.json#/')`, here `{model}`
 will be replaced with the model naming using `str.format()`.
 
+## Getting schema of a specified type
+
+_Pydantic_ includes two standalone utility functions `schema` and `schema_json` that can be used to
+apply the schema generation logic used for _pydantic_ models in a more ad-hoc way.
+These functions behave similarly to `BaseModel.schema` and `BaseModel.schema_json`,
+but work with arbitrary pydantic-compatible types.
+
+```py
+{!.tmp_examples/schema_ad_hoc.py!}
+```
+_(This script is complete, it should run "as is")_
+
+
 ## Field customisation
 
 Optionally, the `Field` function can be used to provide extra information about the field and validations.

docs/usage/types.md

@@ -262,6 +262,39 @@ _(This script is complete, it should run "as is")_
 
     See more details in [Required Fields](models.md#required-fields).
 
+#### Discriminated Unions (a.k.a. Tagged Unions)
+
+When `Union` is used with multiple submodels, you sometimes know exactly which submodel needs to
+be checked and validated and want to enforce this.
+To do that you can set the same field - let's call it `my_discriminator` - in each of the submodels
+with a discriminated value, which is one (or many) `Literal` value(s).
+For your `Union`, you can set the discriminator in its value: `Field(discriminator='my_discriminator')`.
+
+Setting a discriminated union has many benefits:
+
+- validation is faster since it is only attempted against one model
+- only one explicit error is raised in case of failure
+- the generated JSON schema implements the [associated OpenAPI specification](https://github.com/OAI/OpenAPI-Specification/blob/master/versions/3.0.2.md#discriminatorObject)
+
+```py
+{!.tmp_examples/types_union_discriminated.py!}
+```
+_(This script is complete, it should run "as is")_
+
+!!! note
+    Using the [Annotated Fields syntax](../schema/#typingannotated-fields) can be handy to regroup
+    the `Union` and `discriminator` information. See below for an example!
+
+#### Nested Discriminated Unions
+
+Only one discriminator can be set for a field but sometimes you want to combine multiple discriminators.
+In this case you can always create "intermediate" models with `__root__` and add your discriminator.
+
+```py
+{!.tmp_examples/types_union_discriminated_nested.py!}
+```
+_(This script is complete, it should run "as is")_
+
 ### Enums and Choices
 
 *pydantic* uses python's standard `enum` classes to define choices.

pydantic/__init__.py

@@ -62,6 +62,8 @@
     'parse_file_as',
     'parse_obj_as',
     'parse_raw_as',
+    'schema',
+    'schema_json',
     # types
     'NoneStr',
     'NoneBytes',

pydantic/errors.py

@@ -1,6 +1,6 @@
 from decimal import Decimal
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Set, Tuple, Type, Union
+from typing import TYPE_CHECKING, Any, Callable, Sequence, Set, Tuple, Type, Union
 
 from .typing import display_as_type
 
@@ -95,6 +95,8 @@
     'InvalidLengthForBrand',
     'InvalidByteSize',
     'InvalidByteSizeUnit',
+    'MissingDiscriminator',
+    'InvalidDiscriminator',
 )
 
 
@@ -581,3 +583,23 @@ class InvalidByteSize(PydanticValueError):
 
 class InvalidByteSizeUnit(PydanticValueError):
     msg_template = 'could not interpret byte unit: {unit}'
+
+
+class MissingDiscriminator(PydanticValueError):
+    code = 'discriminated_union.missing_discriminator'
+    msg_template = 'Discriminator {discriminator_key!r} is missing in value'
+
+
+class InvalidDiscriminator(PydanticValueError):
+    code = 'discriminated_union.invalid_discriminator'
+    msg_template = (
+        'No match for discriminator {discriminator_key!r} and value {discriminator_value!r} '
+        '(allowed values: {allowed_values})'
+    )
+
+    def __init__(self, *, discriminator_key: str, discriminator_value: Any, allowed_values: Sequence[Any]) -> None:
+        super().__init__(
+            discriminator_key=discriminator_key,
+            discriminator_value=discriminator_value,
+            allowed_values=', '.join(map(repr, allowed_values)),
+        )