-
Notifications
You must be signed in to change notification settings - Fork 15
/
test_pydantic.py
107 lines (82 loc) · 2.72 KB
/
test_pydantic.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
"""
This file tests pydantic support.
See https://github.com/snok/flake8-type-checking/issues/52
for discussion on the implementation.
"""
import textwrap
import pytest
from flake8_type_checking.codes import TC002
from tests import _get_error
@pytest.mark.parametrize(
'enabled, expected',
(
[True, {'2:0 ' + TC002.format(module='decimal.Decimal')}],
[False, {'2:0 ' + TC002.format(module='decimal.Decimal')}],
),
)
def test_non_pydantic_model(enabled, expected):
"""
A class cannot be a pydantic model if it doesn't have a base class,
so we should raise the same error here in both cases.
"""
example = textwrap.dedent(
'''
from decimal import Decimal
class X:
x: Decimal
'''
)
assert _get_error(example, error_code_filter='TC001,TC002', type_checking_pydantic_enabled=enabled) == expected
def test_class_with_base_class():
"""
Whenever a class inherits from anything, we need
to assume it might be a pydantic model, for which
we need to register annotations as uses.
"""
example = textwrap.dedent(
'''
from decimal import Decimal
class X(Y):
x: Decimal
'''
)
assert _get_error(example, error_code_filter='TC001,TC002', type_checking_pydantic_enabled=True) == set()
def test_complex_pydantic_model():
"""
Test actual Pydantic models, with different annotation types.
"""
example = textwrap.dedent(
'''
from __future__ import annotations
from datetime import datetime
from decimal import Decimal
from typing import TYPE_CHECKING
from pydantic import BaseModel, condecimal, validator
if TYPE_CHECKING:
from datetime import date
from typing import Union
def format_datetime(value: Union[str, datetime]) -> datetime:
if isinstance(value, str):
value = datetime.strptime(value, '%Y-%m-%dT%H:%M:%S.%f%z')
assert isinstance(value, datetime)
return value
class ModelBase(BaseModel):
id: int
created_at: datetime
updated_at: datetime
_format_datetime = validator('created_at', 'updated_at', pre=True, allow_reuse=True)(format_datetime)
class NestedModel(ModelBase):
z: Decimal
x: int
y: str
class FinalModel(ModelBase):
a: str
b: int
c: float
d: bool
e: date
f: NestedModel
g: condecimal(ge=Decimal(0)) = Decimal(0)
'''
)
assert _get_error(example, error_code_filter='TC001,TC002', type_checking_pydantic_enabled=True) == set()