|
45 | 45 | import unittest.mock |
46 | 46 | from contextlib import contextmanager |
47 | 47 | from typing import ( |
48 | | - TYPE_CHECKING, |
49 | 48 | Any, |
50 | 49 | Callable, |
51 | 50 | Dict, |
|
57 | 56 | ) |
58 | 57 |
|
59 | 58 | from parameterized import parameterized |
60 | | - |
61 | | -from synapse._pydantic_compat import HAS_PYDANTIC_V2 |
62 | | - |
63 | | -if TYPE_CHECKING or HAS_PYDANTIC_V2: |
64 | | - from pydantic.v1 import ( |
65 | | - BaseModel as PydanticBaseModel, |
66 | | - conbytes, |
67 | | - confloat, |
68 | | - conint, |
69 | | - constr, |
70 | | - ) |
71 | | - from pydantic.v1.typing import get_args |
72 | | -else: |
73 | | - from pydantic import ( |
74 | | - BaseModel as PydanticBaseModel, |
75 | | - conbytes, |
76 | | - confloat, |
77 | | - conint, |
78 | | - constr, |
79 | | - ) |
80 | | - from pydantic.typing import get_args |
81 | | - |
82 | 59 | from typing_extensions import ParamSpec |
83 | 60 |
|
| 61 | +from synapse._pydantic_compat import ( |
| 62 | + BaseModel as PydanticBaseModel, |
| 63 | + conbytes, |
| 64 | + confloat, |
| 65 | + conint, |
| 66 | + constr, |
| 67 | + get_args, |
| 68 | +) |
| 69 | + |
84 | 70 | logger = logging.getLogger(__name__) |
85 | 71 |
|
86 | 72 | CONSTRAINED_TYPE_FACTORIES_WITH_STRICT_FLAG: List[Callable] = [ |
@@ -183,22 +169,16 @@ def monkeypatch_pydantic() -> Generator[None, None, None]: |
183 | 169 | # Most Synapse code ought to import the patched objects directly from |
184 | 170 | # `pydantic`. But we also patch their containing modules `pydantic.main` and |
185 | 171 | # `pydantic.types` for completeness. |
186 | | - patch_basemodel1 = unittest.mock.patch( |
187 | | - "pydantic.BaseModel", new=PatchedBaseModel |
188 | | - ) |
189 | | - patch_basemodel2 = unittest.mock.patch( |
190 | | - "pydantic.main.BaseModel", new=PatchedBaseModel |
| 172 | + patch_basemodel = unittest.mock.patch( |
| 173 | + "synapse._pydantic_compat.BaseModel", new=PatchedBaseModel |
191 | 174 | ) |
192 | | - patches.enter_context(patch_basemodel1) |
193 | | - patches.enter_context(patch_basemodel2) |
| 175 | + patches.enter_context(patch_basemodel) |
194 | 176 | for factory in CONSTRAINED_TYPE_FACTORIES_WITH_STRICT_FLAG: |
195 | 177 | wrapper: Callable = make_wrapper(factory) |
196 | | - patch1 = unittest.mock.patch(f"pydantic.{factory.__name__}", new=wrapper) |
197 | | - patch2 = unittest.mock.patch( |
198 | | - f"pydantic.types.{factory.__name__}", new=wrapper |
| 178 | + patch = unittest.mock.patch( |
| 179 | + f"synapse._pydantic_compat.{factory.__name__}", new=wrapper |
199 | 180 | ) |
200 | | - patches.enter_context(patch1) |
201 | | - patches.enter_context(patch2) |
| 181 | + patches.enter_context(patch) |
202 | 182 | yield |
203 | 183 |
|
204 | 184 |
|
|
0 commit comments