|
1 | 1 | import re
|
| 2 | +import sys |
2 | 3 | from enum import Enum
|
3 |
| -from typing import Union |
| 4 | +from typing import Generic, TypeVar, Union |
4 | 5 |
|
5 | 6 | import pytest
|
6 | 7 | from typing_extensions import Annotated, Literal
|
7 | 8 |
|
8 | 9 | from pydantic import BaseModel, Field, ValidationError
|
9 | 10 | from pydantic.errors import ConfigError
|
| 11 | +from pydantic.generics import GenericModel |
10 | 12 |
|
11 | 13 |
|
12 | 14 | def test_discriminated_union_only_union():
|
@@ -361,3 +363,36 @@ class Model(BaseModel):
|
361 | 363 | n: int
|
362 | 364 |
|
363 | 365 | assert isinstance(Model(**{'pet': {'pet_type': 'dog', 'name': 'Milou'}, 'n': 5}).pet, Dog)
|
| 366 | + |
| 367 | + |
| 368 | +@pytest.mark.skipif(sys.version_info < (3, 7), reason='generics only supported for python 3.7 and above') |
| 369 | +def test_generic(): |
| 370 | + T = TypeVar('T') |
| 371 | + |
| 372 | + class Success(GenericModel, Generic[T]): |
| 373 | + type: Literal['Success'] = 'Success' |
| 374 | + data: T |
| 375 | + |
| 376 | + class Failure(BaseModel): |
| 377 | + type: Literal['Failure'] = 'Failure' |
| 378 | + error_message: str |
| 379 | + |
| 380 | + class Container(GenericModel, Generic[T]): |
| 381 | + result: Union[Success[T], Failure] = Field(discriminator='type') |
| 382 | + |
| 383 | + with pytest.raises(ValidationError, match="Discriminator 'type' is missing in value"): |
| 384 | + Container[str].parse_obj({'result': {}}) |
| 385 | + |
| 386 | + with pytest.raises( |
| 387 | + ValidationError, |
| 388 | + match=re.escape("No match for discriminator 'type' and value 'Other' (allowed values: 'Success', 'Failure')"), |
| 389 | + ): |
| 390 | + Container[str].parse_obj({'result': {'type': 'Other'}}) |
| 391 | + |
| 392 | + with pytest.raises( |
| 393 | + ValidationError, match=re.escape('Container[str]\nresult -> Success[str] -> data\n field required') |
| 394 | + ): |
| 395 | + Container[str].parse_obj({'result': {'type': 'Success'}}) |
| 396 | + |
| 397 | + # coercion is done properly |
| 398 | + assert Container[str].parse_obj({'result': {'type': 'Success', 'data': 1}}).result.data == '1' |
0 commit comments