Skip to content

Commit 507e8cc

Browse files
JensHeinrichKludex
JensHeinrich
and
Kludex
authored
Add TODO notes to __modify_schema__ and __get_validators__ (#116)
Co-authored-by: JensHeinrich <github.com/JensHeinrich> Co-authored-by: Marcelo Trylesinski <[email protected]>
1 parent e0a6ba9 commit 507e8cc

File tree

6 files changed

+256
-2
lines changed

6 files changed

+256
-2
lines changed

‎README.md

Lines changed: 41 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@ Bump Pydantic is a tool to help you migrate your code from Pydantic V1 to V2.
2727
- [BP006: Replace `__root__` by `RootModel`](#bp006-replace-__root__-by-rootmodel)
2828
- [BP007: Replace decorators](#bp007-replace-decorators)
2929
- [BP008: Replace `con*` functions by `Annotated` versions](#bp008-replace-con-functions-by-annotated-versions)
30+
- [BP009: Mark pydantic "protocol" functions in custom types with proper TODOs](bp009-mark-pydantic-protocol-functions-in-custom-types-with-proper-todos)
31+
3032
- [License](#license)
3133

3234
---
@@ -301,7 +303,44 @@ class User(BaseModel):
301303
name: Annotated[str, StringConstraints(min_length=1)]
302304
```
303305

304-
<!-- ### BP009: Replace `pydantic.parse_obj_as` by `pydantic.TypeAdapter`
306+
### BP009: Mark Pydantic "protocol" functions in custom types with proper TODOs
307+
308+
- ✅ Mark `__get_validators__` as to be replaced by `__get_pydantic_core_schema__`.
309+
- ✅ Mark `__modify_schema__` as to be replaced by `__get_pydantic_json_schema__`.
310+
311+
The following code will be transformed:
312+
313+
```py
314+
class SomeThing:
315+
@classmethod
316+
def __get_validators__(cls):
317+
yield from []
318+
319+
@classmethod
320+
def __modify_schema__(cls, field_schema, field):
321+
if field:
322+
field_schema['example'] = "Weird example"
323+
```
324+
325+
Into:
326+
327+
```py
328+
class SomeThing:
329+
@classmethod
330+
# TODO[pydantic]: We couldn't refactor `__get_validators__`, please create the `__get_pydantic_core_schema__` manually.
331+
# Check https://docs.pydantic.dev/latest/migration/#defining-custom-types for more information.
332+
def __get_validators__(cls):
333+
yield from []
334+
335+
@classmethod
336+
# TODO[pydantic]: We couldn't refactor `__modify_schema__`, please create the `__get_pydantic_json_schema__` manually.
337+
# Check https://docs.pydantic.dev/latest/migration/#defining-custom-types for more information.
338+
def __modify_schema__(cls, field_schema, field):
339+
if field:
340+
field_schema['example'] = "Weird example"
341+
```
342+
343+
<!-- ### BP010: Replace `pydantic.parse_obj_as` by `pydantic.TypeAdapter`
305344
306345
- ✅ Replace `pydantic.parse_obj_as(T, obj)` to `pydantic.TypeAdapter(T).validate_python(obj)`.
307346
@@ -344,7 +383,7 @@ class Users(BaseModel):
344383
users = TypeAdapter(Users).validate_python({'users': [{'name': 'John'}]})
345384
``` -->
346385

347-
<!-- ### BP010: Replace `PyObject` by `ImportString`
386+
<!-- ### BP011: Replace `PyObject` by `ImportString`
348387
349388
- ✅ Replace `PyObject` by `ImportString`.
350389

‎bump_pydantic/codemods/__init__.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66

77
from bump_pydantic.codemods.add_default_none import AddDefaultNoneCommand
88
from bump_pydantic.codemods.con_func import ConFuncCallCommand
9+
from bump_pydantic.codemods.custom_types import CustomTypeCodemod
910
from bump_pydantic.codemods.field import FieldCodemod
1011
from bump_pydantic.codemods.replace_config import ReplaceConfigCodemod
1112
from bump_pydantic.codemods.replace_generic_model import ReplaceGenericModelCommand
@@ -31,6 +32,8 @@ class Rule(str, Enum):
3132
"""Replace `@validator` with `@field_validator`."""
3233
BP008 = "BP008"
3334
"""Replace `con*` functions by `Annotated` versions."""
35+
BP009 = "BP009"
36+
"""Mark Pydantic "protocol" functions in custom types with proper TODOs."""
3437

3538

3639
def gather_codemods(disabled: List[Rule]) -> List[Type[ContextAwareTransformer]]:
@@ -61,6 +64,9 @@ def gather_codemods(disabled: List[Rule]) -> List[Type[ContextAwareTransformer]]
6164
if Rule.BP007 not in disabled:
6265
codemods.append(ValidatorCodemod)
6366

67+
if Rule.BP009 not in disabled:
68+
codemods.append(CustomTypeCodemod)
69+
6470
# Those codemods need to be the last ones.
6571
codemods.extend([RemoveImportsVisitor, AddImportsVisitor])
6672
return codemods

‎bump_pydantic/codemods/custom_types.py

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
import libcst as cst
2+
from libcst import matchers as m
3+
from libcst.codemod import CodemodContext, VisitorBasedCodemodCommand
4+
from libcst.codemod.visitors import AddImportsVisitor
5+
6+
PREFIX_COMMENT = "# TODO[pydantic]: "
7+
REFACTOR_COMMENT = f"{PREFIX_COMMENT}We couldn't refactor `{{old_name}}`, please create the `{{new_name}}` manually."
8+
GET_VALIDATORS_COMMENT = REFACTOR_COMMENT.format(old_name="__get_validators__", new_name="__get_pydantic_core_schema__")
9+
MODIFY_SCHEMA_COMMENT = REFACTOR_COMMENT.format(old_name="__modify_schema__", new_name="__get_pydantic_json_schema__")
10+
COMMENT_BY_FUNC_NAME = {"__get_validators__": GET_VALIDATORS_COMMENT, "__modify_schema__": MODIFY_SCHEMA_COMMENT}
11+
CHECK_LINK_COMMENT = "# Check https://docs.pydantic.dev/latest/migration/#defining-custom-types for more information."
12+
13+
GET_VALIDATORS_FUNCTION = m.FunctionDef(name=m.Name("__get_validators__"))
14+
MODIFY_SCHEMA_FUNCTION = m.FunctionDef(name=m.Name("__modify_schema__"))
15+
16+
17+
class CustomTypeCodemod(VisitorBasedCodemodCommand):
18+
@m.leave(MODIFY_SCHEMA_FUNCTION | GET_VALIDATORS_FUNCTION)
19+
def leave_modify_schema_func(
20+
self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef
21+
) -> cst.FunctionDef:
22+
for line in [*updated_node.leading_lines, *updated_node.lines_after_decorators]:
23+
if m.matches(line, m.EmptyLine(comment=m.Comment(value=CHECK_LINK_COMMENT))):
24+
return updated_node
25+
26+
comment = COMMENT_BY_FUNC_NAME[updated_node.name.value]
27+
return updated_node.with_changes(
28+
lines_after_decorators=[
29+
*updated_node.lines_after_decorators,
30+
cst.EmptyLine(comment=cst.Comment(value=(comment))),
31+
cst.EmptyLine(comment=cst.Comment(value=(CHECK_LINK_COMMENT))),
32+
]
33+
)
34+
35+
36+
if __name__ == "__main__":
37+
import textwrap
38+
39+
from rich.console import Console
40+
41+
console = Console()
42+
43+
source = textwrap.dedent(
44+
"""
45+
class SomeThing:
46+
@classmethod
47+
def __get_validators__(cls):
48+
yield from []
49+
return
50+
51+
@classmethod
52+
def __modify_schema__(
53+
cls, field_schema: Dict[str, Any], field: Optional[ModelField]
54+
):
55+
if field:
56+
field_schema['example'] = "Weird example"
57+
"""
58+
)
59+
console.print(source)
60+
console.print("=" * 80)
61+
62+
mod = cst.parse_module(source)
63+
context = CodemodContext(filename="main.py")
64+
wrapper = cst.MetadataWrapper(mod)
65+
command = CustomTypeCodemod(context=context)
66+
# console.print(mod)
67+
68+
mod = wrapper.visit(command)
69+
wrapper = cst.MetadataWrapper(mod)
70+
command = AddImportsVisitor(context=context) # type: ignore[assignment]
71+
mod = wrapper.visit(command)
72+
console.print(mod.code)

‎tests/integration/cases/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from .base_settings import cases as base_settings_cases
66
from .con_func import cases as con_func_cases
77
from .config_to_model import cases as config_to_model_cases
8+
from .custom_types import cases as custom_types_cases
89
from .field import cases as generic_model_cases
910
from .folder_inside_folder import cases as folder_inside_folder_cases
1011
from .is_base_model import cases as is_base_model_cases
@@ -28,6 +29,7 @@
2829
*folder_inside_folder_cases,
2930
*unicode_cases,
3031
*con_func_cases,
32+
*custom_types_cases,
3133
]
3234
before = Folder("project", *[case.source for case in cases])
3335
expected = Folder("project", *[case.expected for case in cases])

‎tests/integration/cases/custom_types.py

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
from ..case import Case
2+
from ..file import File
3+
4+
cases = [
5+
Case(
6+
name="Mark __get_validators__",
7+
source=File(
8+
"mark_get_validators.py",
9+
content=[
10+
"class SomeThing:",
11+
" @classmethod",
12+
" def __get_validators__(cls):",
13+
" yield from []",
14+
" return",
15+
],
16+
),
17+
expected=File(
18+
"mark_get_validators.py",
19+
content=[
20+
"class SomeThing:",
21+
" @classmethod",
22+
" # TODO[pydantic]: We couldn't refactor `__get_validators__`, please create the `__get_pydantic_core_schema__` manually.", # noqa: E501
23+
" # Check https://docs.pydantic.dev/latest/migration/#defining-custom-types for more information.",
24+
" def __get_validators__(cls):",
25+
" yield from []",
26+
" return",
27+
],
28+
),
29+
),
30+
Case(
31+
name="Mark __modify_schema__",
32+
source=File(
33+
"mark_modify_schema.py",
34+
content=[
35+
"class SomeThing:",
36+
" @classmethod",
37+
" def __modify_schema__(",
38+
" cls, field_schema: Dict[str, Any], field: Optional[ModelField]",
39+
" ):",
40+
" if field:",
41+
" field_schema['example'] = \"Weird example\"",
42+
],
43+
),
44+
expected=File(
45+
"mark_modify_schema.py",
46+
content=[
47+
"class SomeThing:",
48+
" @classmethod",
49+
" # TODO[pydantic]: We couldn't refactor `__modify_schema__`, please create the `__get_pydantic_json_schema__` manually.", # noqa: E501
50+
" # Check https://docs.pydantic.dev/latest/migration/#defining-custom-types for more information.",
51+
" def __modify_schema__(",
52+
" cls, field_schema: Dict[str, Any], field: Optional[ModelField]",
53+
" ):",
54+
" if field:",
55+
" field_schema['example'] = \"Weird example\"",
56+
],
57+
),
58+
),
59+
]

‎tests/unit/test_custom_types.py

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
from libcst.codemod import CodemodTest
2+
3+
from bump_pydantic.codemods.custom_types import CustomTypeCodemod
4+
5+
6+
class TestArbitraryClassCommand(CodemodTest):
7+
TRANSFORM = CustomTypeCodemod
8+
9+
maxDiff = None
10+
11+
def test_mark_get_validators(self) -> None:
12+
before = """
13+
class SomeThing:
14+
@classmethod
15+
def __get_validators__(cls):
16+
yield from []
17+
return
18+
"""
19+
after = """
20+
class SomeThing:
21+
@classmethod
22+
# TODO[pydantic]: We couldn't refactor `__get_validators__`, please create the `__get_pydantic_core_schema__` manually.
23+
# Check https://docs.pydantic.dev/latest/migration/#defining-custom-types for more information.
24+
def __get_validators__(cls):
25+
yield from []
26+
return
27+
""" # noqa: E501
28+
self.assertCodemod(before, after)
29+
30+
def test_mark_modify_schema(self) -> None:
31+
before = """
32+
class SomeThing:
33+
@classmethod
34+
def __modify_schema__(
35+
cls, field_schema: Dict[str, Any], field: Optional[ModelField]
36+
):
37+
if field:
38+
field_schema['example'] = "Weird example"
39+
"""
40+
after = """
41+
class SomeThing:
42+
@classmethod
43+
# TODO[pydantic]: We couldn't refactor `__modify_schema__`, please create the `__get_pydantic_json_schema__` manually.
44+
# Check https://docs.pydantic.dev/latest/migration/#defining-custom-types for more information.
45+
def __modify_schema__(
46+
cls, field_schema: Dict[str, Any], field: Optional[ModelField]
47+
):
48+
if field:
49+
field_schema['example'] = "Weird example"
50+
""" # noqa: E501
51+
self.assertCodemod(before, after)
52+
53+
def test_already_commented(self) -> None:
54+
before = """
55+
class SomeThing:
56+
@classmethod
57+
# TODO[pydantic]: We couldn't refactor `__modify_schema__`, please create the `__get_pydantic_json_schema__` manually.
58+
# Check https://docs.pydantic.dev/latest/migration/#defining-custom-types for more information.
59+
def __modify_schema__(
60+
cls, field_schema: Dict[str, Any], field: Optional[ModelField]
61+
):
62+
if field:
63+
field_schema['example'] = "Weird example"
64+
""" # noqa: E501
65+
after = """
66+
class SomeThing:
67+
@classmethod
68+
# TODO[pydantic]: We couldn't refactor `__modify_schema__`, please create the `__get_pydantic_json_schema__` manually.
69+
# Check https://docs.pydantic.dev/latest/migration/#defining-custom-types for more information.
70+
def __modify_schema__(
71+
cls, field_schema: Dict[str, Any], field: Optional[ModelField]
72+
):
73+
if field:
74+
field_schema['example'] = "Weird example"
75+
""" # noqa: E501
76+
self.assertCodemod(before, after)

0 commit comments

Comments
 (0)