[Feature Request] Provide a discriminated union type (OpenAPI 3)

Feature Request

Pydantic currently has a decent support for union types through the typing.Union type from PEP484, but it does not currently cover all the cases covered by the JSONSchema and OpenAPI specifications, most likely because the two specifications diverge on those points.

OpenAPI supports something similar to tagged unions where a certain field is designated to serve as a "discriminator", which is then matched against literal values to determine which of multiple schemas to use for payload validation. In order to allow Pydantic to support those, I suppose there would have to be a specific type similar to typing.Union in order to specify what discriminator field to use and how to match it. Such a type would then be rendered into a schema object (oneOf) with an OpenAPI discriminator object built into it, as well as correctly validate incoming JSON into the correct type based on the value or the discriminator field. This change would only impact OpenAPI, as JSON schema (draft 7 onwards) uses conditional types instead, which would probably need to be the topic of a different feature request, as both methods appear mutually incompatible.

Implementation ideas

I'd imagine the final result to be something like this.

MyUnion = Union[Foo, Bar]
MyTaggedUnion = TaggedUnion(Union[Foo, Bar], discriminator='type', mapping={'foo': Foo, 'bar': Bar}))

Python doesn't have a feature like TypeScript to let you statically ensure that discriminator exists as a field for all variants of that union, though that shouldn't be a problem since this is going to be raised during validation regardless.

discriminator and mapping could also simply be added to Schema, though I'm not sure about whether it's a good idea to add OpenAPI-specific extensions there.

PEP 593 would also have been a nice alternative, since it would hypothetically allow tagged unions to be implemented as a regular union with annotations specific to Pydantic for that purpose, however it is only still a draft and most likely won't make it until Python 3.9 (if at all).

The validation component of this can already by accomplished via the const kwarg to Schema will be permitted as annotations via Literal once #582 is released.

This allows you to add a field which is constrained to one or multiple values then use that field to discriminate between to models.

However at the moment this doesn't extend to a discriminator object in the schema.

Perhaps it would be possible to either:

automatically spot the discriminator when building schema, or
add a property to Config or a kwarg to Schema to tell pydantic about it

?

Guessing what the discriminator field may be based on const or literal values might lead to unexpected behavior, using Config properties or schema parameters would probably be preferable.

A union field could have a discriminator parameter on its schema object to indicate which field to match against, which would then have to be const/literal values on each of the variants. Each literal value could then be rendered as a separate key/value pair in the mapping dictionary, which would raise an error in case of collision.

This would leave the default case (first key in the discriminator mapping is considered as the default option is nothing matches, IIRC) undefined, though...

I'll see if I can try coming up with something that can be easily validated against while not being too cumbersome to work with.

I've taken another look at the OpenAPI spec on this and tried a number of potential syntaxes.

The Literal syntax for Python 3.8 combined with an inheritance model flows really well. Considering discriminator objects are only allowed as fields of non-inlined schema objects in OpenAPI, there should be no issue with allowing them as BaseModel properties only. I'm not entirely sure, however, of how well mypy will tolerate shadowing with narrower types like that.

from pydantic import BaseModel, Schema
from typing import Literal

class Pet(BaseModel):
	petType: str = Schema(...)

	# One or the other
	__discriminator__ = 'petType'

	class Config:
		discriminator = 'petType'

class Cat(Pet):
	petType: Literal['cat', 'robot_cat'] # Should render as a string enum
	name: str = Schema(...)

class Dog(Pet):
	petType: Literal['dog', 'robot_dog']
	bark: str = Schema(...)

There would however be the problem of dealing with versions of Python without support for Literal, which could be accomplished with const as you've suggested, but this would mean having no more than one matched value per type. "Mapping keys MUST be string values, but tooling MAY convert response values to strings for comparison." would indicate that the discriminator property has to be of type string and that the values for subclasses can be validated as enums (of strings) as well.

The OpenAPI spec also mentions that the schema names can act as implicit discriminators, whether or not an explicit mapping is present (so 'Dog' and 'Cat' could technically be valid if there is no other constraint), however I don't really believe supporting that use case is of much concern.

Literal is supported in python 3.7 (and 3.6 if I recall correctly) by importing from the package typing_extensions.

Ah, well, that's sure to make things a lot simpler, then.

That leaves the question of whether to use model config or __dunder__ attributes to define the discriminator property. Pydantic has some examples for both styles and I can't find any info on when one should be prefered over the other in style guidelines.

EDIT: All __dunder__ attributes and methods are considered reserved by the interpreter and can break without warning, so adding more such attributes may not be such a good idea after all.

I might try filing a PR for this in the coming days if I have time.

@dmontagu

I will propose a PR to include additional schema (JSON Schema) data for models.

This will allow you to create the validation required using, e.g. Literal or the Generics functionality @dmontagu added. Then you can do the validation in your model using standard Pydantic validators.

And then you can describe it in the extra schema data (as JSON Schema/OpenAPI schema) as an OpenAPI discriminator, etc. purely for documentation.

I think adding discriminator support directly to Pydantic wouldn't be convenient, as the discriminator ideas are very specific to OpenAPI, are a bit constrained (not that generalizable), and have conflicts with some ideas in JSON Schema (if, else, etc).

But combining these things I described above (with the extra schema I'll PR) you should be able to achieve what you need, with the OpenAPI documentation you expect @sm-Fifteen .

@sm-Fifteen

@sm-Fifteen I think you can now perform the validation as you need in a validator and generate the schema the way you want it using schema_extra. Could you check if that solves your use case?

Thinking about this more (and running into a similar problem myself), I think some kind of discriminator field is required:

To speed up parsing so that validation only needs to be attempted against one model
To make the error message less verbose - currently if you try to validate against multiple models, all but one are likely to have multiple errors. All those errors get added to the error output making it extremely verbose.

Personally I think this should be done by adding a discriminator argument to Field / Schema (#577) rather than creating a custom Union type that will never play well with mypy or IDEs.

So Usage would be something like

class Foo(BaseModel):
    model_type: Literal['foo']

class Bar(BaseModel):
    model_type: Literal['bar']

class MyModel(BaseModel):
    foobar: Union[Foo, Bar] = Field(..., descriminator='model_type')

@samuelcolvin

@samuelcolvin @PrettyWood Any updates on this? I'm in the progress of switching away from GraphQL (for a lot of reasons, unrelated) but having proper support for tagged unions is kinda an issue. Also that tagged unions aren't properly exposed to OpenAPI.

Not sure if there are any bounty programs that pydantic supports, but I'd be happy to put up some $ for this

@ghostbody

@ghostbody Inspired our solution. We used type Literal from typing_extensions and an Enum class.

@unique
class SelectFieldTypes(str, Enum):
    date = "date"
    multiple_choice = "multiple-choice"
    multiple_select = "multiple-select"


class DateModel(BaseModel):
    id: str
    type: Literal[SelectFieldTypes.date]
    text: str
    description: str
    required: bool
    value: Optional[str]


class MultipleChoice(BaseModel):
    id: str
    type: Literal[SelectFieldTypes.multiple_choice]
    text: str
    description: str
    required: bool
    options: List[Option]
    value: Optional[str]


class MultipleSelect(BaseModel):
    id: str
    type: Literal[SelectFieldTypes.multiple_select]
    text: str
    description: str
    required: bool
    options: List[Option]
    value: Optional[str]

class SelectTemplate(BaseModel):
    id: str
    text: str
    fields_: List[
        Union[
            DateModel,
            MultipleChoice,
            MultipleSelect,
        ]
    ] = Field(..., alias="fields")

Hi everyone 😃
I just opened a small PR to see if that's the desired behaviour. If I understand correctly it's really meant to be used with Literal (like in TypeScript) to have

faster validation (I hence added directly in the metaclass the discriminator mapping)
smaller error
generated schema with the discriminator key

If that's all we want I guess my PR should be a good first draft 👍
Feeback more than welcome 😃
Do not hesitate to try it out

@PrettyWood

@PrettyWood That PR looks great. I'll give it a try later!

@PrettyWood

@PrettyWood: That looks like a pretty good way of handling things, having types with literal fields that could be validated just fine without the use of discriminators, and then specifying a discriminator field on some union of those types use on the parent union, in a way that builds on top of regular validation. This is actually similar to what's has been proposed as a potential replacement for discriminator on the OpenAPI side of things (see OAI/OpenAPI-Specification#2143 (comment)).

I'm not seeing any logic or tests for handling such discriminated unions as the root element, though. Is this supported by this PR?

@sm-Fifteen

@sm-Fifteen

    class Cat(BaseModel):
        pet_type: Literal['cat']
        name: str

    class Dog(BaseModel):
        pet_type: Literal['dog']
        name: str

    class Pet(BaseModel):
        __root__: Union[Cat, Dog] = Field(..., discriminator='pet_type')

    my_dog = Pet.parse_obj({'pet_type': 'dog', 'name': 'woof'}).__root__
    assert isinstance(my_dog, Dog)

works. Not sure if that's what you had in mind

@sm-Fifteen

Btw @sm-Fifteen in fact it should be even easier with #2147.
I added two new utils schema and schema_json in the same PR.
Now you can write directly this

Pet = Annotated[Union[Cat, Dog], Field(discriminator='pet_type')]

try:
    parse_obj_as(Pet, {'pet_type': 'dog', 'nam': 'woof'})
except ValidationError as e:
    print(e)
# 1 validation error for ParsingModel[Annotated[Union[__main__.Cat, __main__.Dog], FieldInfo(default=Ellipsis, extra={'discriminator': 'pet_type'})]]
# __root__ (Dog) -> name
#   field required (type=value_error.missing)

print(schema_json(Pet, title='Pet', indent=2))
# {
#   "title": "Pet",
#   "discriminator": {
#     "propertyName": "pet_type",
#     "mapping": {
#       "cat": "#/definitions/Cat",
#       "dog": "#/definitions/Dog"
#     }
#   },
#   "anyOf": [
#     {
#       "$ref": "#/definitions/Cat"
#     },
#     {
#       "$ref": "#/definitions/Dog"
#     }
#   ],
#   "definitions": {
#  ...

@sm-Fifteen

Btw @sm-Fifteen in fact it should be even easier with #2147.
I added two new utils schema and schema_json in the same PR.

I wasn't super convinced by the other root union example, although I figured it was at least serviceable, but using PEP 593 Annotated types now that they are available really does make for something fairly ergonomic. I like it.

Thank you so much for that PR. Testing it out, is seems like you can call Cat(name='Felix') directly without having to supply the tag if you add default values to the tags:

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')]

# parse_obj_as(Pet, dict(name='Felix'))  # fails as expected
parse_obj_as(Pet, dict(pet_type='cat', name='Felix'))
Cat(name='Felix')

Is this an expected use case or are there problems with this pattern?

@KevOrr

@KevOrr Yep the discriminator is used for the schema and to improve validation (faster and more explicit)
But IMO it's normal to be able to set the default value of your discriminator field to instantiate directly your submodels easily

I've been using this branch for weeks to add a discriminator to several types and it seems like it would be a handy compliment to Union. Is there more work to be done to solidify some part of it? I'm here to help if it is possible. The branch is falling further behind by the day and it seems like such a great feature.

@PrettyWood

#619 (comment)
This is great, thank you! @PrettyWood! Works for me like this even with a plain Union in pydantic 1.8.2.

However, with this code, it seems as if (1) (even non-pre) root validators on the non-matching classes are still executed, though, and (2) they don't get passed the discriminator field value:

    class Cat(BaseModel):
        pet_type: Literal['cat']
        name: str
        age: int
        
        @root_validator
        def height_constraints(cls, values: Dict[str, Any]) -> Dict[str, Any]:
            print(values)
            assert str(values["age"]) in values["name"]  # silly check
            return values

    class Dog(BaseModel):
        pet_type: Literal['dog']
        name: str

    class Pet(BaseModel):
        __root__: Union[Cat, Dog] = Field(..., discriminator='pet_type')

    my_dog = Pet.parse_obj({'pet_type': 'dog', 'name': 'woof'}).__root__
    assert isinstance(my_dog, Dog)

Running this code prints:

{'name': 'woof'}
Traceback (most recent call last):
  File "test_pyd.py", line 23, in <module>
    my_dog = Pet.parse_obj({'pet_type': 'dog', 'name': 'woof'}).__root__
  File "pydantic/main.py", line 578, in pydantic.main.BaseModel.parse_obj
  File "pydantic/main.py", line 404, in pydantic.main.BaseModel.__init__
  File "pydantic/main.py", line 1040, in pydantic.main.validate_model
  File "pydantic/fields.py", line 723, in pydantic.fields.ModelField.validate
  File "pydantic/fields.py", line 899, in pydantic.fields.ModelField._validate_singleton
  File "pydantic/fields.py", line 723, in pydantic.fields.ModelField.validate
  File "pydantic/fields.py", line 906, in pydantic.fields.ModelField._validate_singleton
  File "pydantic/fields.py", line 913, in pydantic.fields.ModelField._apply_validators
  File "pydantic/class_validators.py", line 310, in pydantic.class_validators._generic_validator_basic.lambda12
  File "pydantic/main.py", line 735, in pydantic.main.BaseModel.validate
  File "pydantic/main.py", line 404, in pydantic.main.BaseModel.__init__
  File "pydantic/main.py", line 1066, in pydantic.main.validate_model
  File "test_pyd.py", line 13, in height_constraints
    assert str(values["age"]) in values["name"]  # silly check
KeyError: 'age'

Does anyone know if that is being addressed in #2336?

@tgpfeiffer

Hi @tgpfeiffer
I can't reproduce with your code snippet and the branch f/discriminated-union.
The whole purpose of the discriminated union is to NOT validate extra classes so if this is actually happening I will definitely look into it

Navigation Menu

Repository navigation

[Feature Request] Provide a discriminated union type (OpenAPI 3) #619

Feature Request

Implementation ideas

37 remaining items

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

[Feature Request] Provide a discriminated union type (OpenAPI 3) #619

Description

Feature Request

Implementation ideas

Activity

samuelcolvin commented on Jun 25, 2019

sm-Fifteen commented on Jun 27, 2019

sm-Fifteen commented on Jun 30, 2019

dmontagu commented on Jun 30, 2019

sm-Fifteen commented on Jul 2, 2019

tiangolo commented on Jul 13, 2019

tiangolo commented on Aug 7, 2019

samuelcolvin commented on Aug 11, 2019

samuelcolvin commented on Aug 11, 2019

37 remaining items

levrik commented on Feb 9, 2021

GavanWilhite commented on Feb 9, 2021

wgriffin13 commented on Feb 9, 2021

PrettyWood commented on Feb 9, 2021

levrik commented on Feb 10, 2021

sm-Fifteen commented on Feb 10, 2021

PrettyWood commented on Feb 10, 2021

PrettyWood commented on Feb 23, 2021

sm-Fifteen commented on Feb 24, 2021

KevOrr commented on Mar 11, 2021

PrettyWood commented on Mar 14, 2021

LyleScott commented on Jul 5, 2021

tgpfeiffer commented on Aug 26, 2021

PrettyWood commented on Aug 26, 2021

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions