1"""
2Test pydantic's compliance with mypy.
3
4Do a little skipping about with types to demonstrate its usage.
5"""
6import json
7import sys
8from datetime import date, datetime
9from pathlib import Path
10from typing import Any, Dict, Generic, List, Optional, TypeVar
11from uuid import UUID
12
13from pydantic import (
14    UUID1,
15    BaseModel,
16    DirectoryPath,
17    FilePath,
18    Json,
19    NegativeFloat,
20    NegativeInt,
21    NoneStr,
22    NonNegativeFloat,
23    NonNegativeInt,
24    NonPositiveFloat,
25    NonPositiveInt,
26    PositiveFloat,
27    PositiveInt,
28    PyObject,
29    StrictBool,
30    StrictBytes,
31    StrictFloat,
32    StrictInt,
33    StrictStr,
34    root_validator,
35    validate_arguments,
36    validator,
37)
38from pydantic.fields import Field, PrivateAttr
39from pydantic.generics import GenericModel
40from pydantic.typing import ForwardRef
41
42
43class Flags(BaseModel):
44    strict_bool: StrictBool = False
45
46    def __str__(self) -> str:
47        return f'flag={self.strict_bool}'
48
49
50class Model(BaseModel):
51    age: int
52    first_name = 'John'
53    last_name: NoneStr = None
54    signup_ts: Optional[datetime] = None
55    list_of_ints: List[int]
56
57    @validator('age')
58    def check_age(cls, value: int) -> int:
59        assert value < 100, 'too old'
60        return value
61
62    @root_validator
63    def root_check(cls, values: Dict[str, Any]) -> Dict[str, Any]:
64        return values
65
66    @root_validator(pre=True, allow_reuse=False, skip_on_failure=False)
67    def pre_root_check(cls, values: Dict[str, Any]) -> Dict[str, Any]:
68        return values
69
70
71def dog_years(age: int) -> int:
72    return age * 7
73
74
75def day_of_week(dt: datetime) -> int:
76    return dt.date().isoweekday()
77
78
79m = Model(age=21, list_of_ints=[1, '2', b'3'])
80
81assert m.age == 21, m.age
82m.age = 42
83assert m.age == 42, m.age
84assert m.first_name == 'John', m.first_name
85assert m.last_name is None, m.last_name
86assert m.list_of_ints == [1, 2, 3], m.list_of_ints
87
88dog_age = dog_years(m.age)
89assert dog_age == 294, dog_age
90
91
92m = Model(age=2, first_name=b'Woof', last_name=b'Woof', signup_ts='2017-06-07 00:00', list_of_ints=[1, '2', b'3'])
93
94assert m.first_name == 'Woof', m.first_name
95assert m.last_name == 'Woof', m.last_name
96assert m.signup_ts == datetime(2017, 6, 7), m.signup_ts
97assert day_of_week(m.signup_ts) == 3
98
99
100data = {'age': 10, 'first_name': 'Alena', 'last_name': 'Sousova', 'list_of_ints': [410]}
101m_from_obj = Model.parse_obj(data)
102
103assert isinstance(m_from_obj, Model)
104assert m_from_obj.age == 10
105assert m_from_obj.first_name == data['first_name']
106assert m_from_obj.last_name == data['last_name']
107assert m_from_obj.list_of_ints == data['list_of_ints']
108
109m_from_raw = Model.parse_raw(json.dumps(data))
110
111assert isinstance(m_from_raw, Model)
112assert m_from_raw.age == m_from_obj.age
113assert m_from_raw.first_name == m_from_obj.first_name
114assert m_from_raw.last_name == m_from_obj.last_name
115assert m_from_raw.list_of_ints == m_from_obj.list_of_ints
116
117m_copy = m_from_obj.copy()
118
119assert isinstance(m_from_raw, Model)
120assert m_copy.age == m_from_obj.age
121assert m_copy.first_name == m_from_obj.first_name
122assert m_copy.last_name == m_from_obj.last_name
123assert m_copy.list_of_ints == m_from_obj.list_of_ints
124
125
126if sys.version_info >= (3, 7):
127    T = TypeVar('T')
128
129    class WrapperModel(GenericModel, Generic[T]):
130        payload: T
131
132    int_instance = WrapperModel[int](payload=1)
133    int_instance.payload += 1
134    assert int_instance.payload == 2
135
136    str_instance = WrapperModel[str](payload='a')
137    str_instance.payload += 'a'
138    assert str_instance.payload == 'aa'
139
140    model_instance = WrapperModel[Model](payload=m)
141    model_instance.payload.list_of_ints.append(4)
142    assert model_instance.payload.list_of_ints == [1, 2, 3, 4]
143
144
145class WithField(BaseModel):
146    age: int
147    first_name: str = Field('John', const=True)
148
149
150# simple decorator
151@validate_arguments
152def foo(a: int, *, c: str = 'x') -> str:
153    return c * a
154
155
156foo(1, c='thing')
157foo(1)
158
159
160# nested decorator should not produce an error
161@validate_arguments(config={'arbitrary_types_allowed': True})
162def bar(a: int, *, c: str = 'x') -> str:
163    return c * a
164
165
166bar(1, c='thing')
167bar(1)
168
169
170class Foo(BaseModel):
171    a: int
172
173
174FooRef = ForwardRef('Foo')
175
176
177class MyConf(BaseModel):
178    str_pyobject: PyObject = Field('datetime.date')
179    callable_pyobject: PyObject = Field(date)
180
181
182conf = MyConf()
183var1: date = conf.str_pyobject(2020, 12, 20)
184var2: date = conf.callable_pyobject(2111, 1, 1)
185
186
187class MyPrivateAttr(BaseModel):
188    _private_field: str = PrivateAttr()
189
190
191class PydanticTypes(BaseModel):
192    # Boolean
193    my_strict_bool: StrictBool = True
194    # Integer
195    my_positive_int: PositiveInt = 1
196    my_negative_int: NegativeInt = -1
197    my_non_positive_int: NonPositiveInt = -1
198    my_non_negative_int: NonNegativeInt = 1
199    my_strict_int: StrictInt = 1
200    # Float
201    my_positive_float: PositiveFloat = 1.1
202    my_negative_float: NegativeFloat = -1.1
203    my_non_positive_float: NonPositiveFloat = -1.1
204    my_non_negative_float: NonNegativeFloat = 1.1
205    my_strict_float: StrictFloat = 1.1
206    # Bytes
207    my_strict_bytes: StrictBytes = b'pika'
208    # String
209    my_strict_str: StrictStr = 'pika'
210    # PyObject
211    my_pyobject_str: PyObject = 'datetime.date'  # type: ignore
212    my_pyobject_callable: PyObject = date
213    # UUID
214    my_uuid1: UUID1 = UUID('a8098c1a-f86e-11da-bd1a-00112444be1e')
215    my_uuid1_str: UUID1 = 'a8098c1a-f86e-11da-bd1a-00112444be1e'  # type: ignore
216    # Path
217    my_file_path: FilePath = Path(__file__)
218    my_file_path_str: FilePath = __file__  # type: ignore
219    my_dir_path: DirectoryPath = Path('.')
220    my_dir_path_str: DirectoryPath = '.'  # type: ignore
221    # Json
222    my_json: Json = '{"hello": "world"}'
223
224    class Config:
225        validate_all = True
226
227
228validated = PydanticTypes()
229validated.my_pyobject_str(2021, 1, 1)
230validated.my_pyobject_callable(2021, 1, 1)
231validated.my_uuid1.hex
232validated.my_uuid1_str.hex
233validated.my_file_path.absolute()
234validated.my_file_path_str.absolute()
235validated.my_dir_path.absolute()
236validated.my_dir_path_str.absolute()
237