Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: support generic models with discriminated union #3551

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
4 changes: 4 additions & 0 deletions pydantic/fields.py
Expand Up @@ -732,6 +732,10 @@ def prepare_discriminated_union_sub_fields(self) -> None:
Note that this process can be aborted if a `ForwardRef` is encountered
"""
assert self.discriminator_key is not None

if self.type_.__class__ is DeferredType:
return

assert self.sub_fields is not None
sub_fields_mapping: Dict[str, 'ModelField'] = {}
all_aliases: Set[str] = set()
Expand Down
37 changes: 36 additions & 1 deletion tests/test_discrimated_union.py
@@ -1,12 +1,14 @@
import re
import sys
from enum import Enum
from typing import Union
from typing import Generic, TypeVar, Union

import pytest
from typing_extensions import Annotated, Literal

from pydantic import BaseModel, Field, ValidationError
from pydantic.errors import ConfigError
from pydantic.generics import GenericModel


def test_discriminated_union_only_union():
Expand Down Expand Up @@ -361,3 +363,36 @@ class Model(BaseModel):
n: int

assert isinstance(Model(**{'pet': {'pet_type': 'dog', 'name': 'Milou'}, 'n': 5}).pet, Dog)


@pytest.mark.skipif(sys.version_info < (3, 7), reason='generics only supported for python 3.7 and above')
def test_generic():
T = TypeVar('T')

class Success(GenericModel, Generic[T]):
type: Literal['Success'] = 'Success'
data: T

class Failure(BaseModel):
type: Literal['Failure'] = 'Failure'
error_message: str

class Container(GenericModel, Generic[T]):
result: Union[Success[T], Failure] = Field(discriminator='type')

with pytest.raises(ValidationError, match="Discriminator 'type' is missing in value"):
Container[str].parse_obj({'result': {}})

with pytest.raises(
ValidationError,
match=re.escape("No match for discriminator 'type' and value 'Other' (allowed values: 'Success', 'Failure')"),
):
Container[str].parse_obj({'result': {'type': 'Other'}})

with pytest.raises(
ValidationError, match=re.escape('Container[str]\nresult -> Success[str] -> data\n field required')
):
Container[str].parse_obj({'result': {'type': 'Success'}})

# coercion is done properly
assert Container[str].parse_obj({'result': {'type': 'Success', 'data': 1}}).result.data == '1'