From 5770c89c49859c6218b64e15f076ba3d6e289ab1 Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Mon, 19 Jun 2023 15:45:37 +0200 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20Support=20GenericModel=20to=20BaseM?= =?UTF-8?q?odel=20replacement=20(#12)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .github/workflows/main.yml | 2 - README.md | 29 +++++ bump_pydantic/codemods/__init__.py | 6 +- .../codemods/replace_generic_model.py | 63 +++++++++++ bump_pydantic/codemods/replace_imports.py | 1 - project/replace_generic.py | 9 ++ pyproject.toml | 13 +-- tests/test_generic_model.py | 102 ++++++++++++++++++ 8 files changed, 215 insertions(+), 10 deletions(-) create mode 100644 bump_pydantic/codemods/replace_generic_model.py create mode 100644 project/replace_generic.py create mode 100644 tests/test_generic_model.py diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 695367f..e7168fd 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -12,7 +12,6 @@ jobs: steps: - uses: actions/checkout@v2 - - name: set up python uses: actions/setup-python@v4 with: @@ -38,7 +37,6 @@ jobs: steps: - uses: actions/checkout@v2 - - name: set up python uses: actions/setup-python@v4 with: diff --git a/README.md b/README.md index 48de098..8f4b2fc 100644 --- a/README.md +++ b/README.md @@ -56,3 +56,32 @@ class User(BaseModel): ``` #### BP004: Replace `BaseModel` methods + + +#### BP005: Replace `GenericModel` by `BaseModel` + +- ✅ Replace `GenericModel` by `BaseModel`. + +The following code will be transformed: + +```py +from typing import Generic, TypeVar +from pydantic.generics import GenericModel + +T = TypeVar('T') + +class User(GenericModel, Generic[T]): + name: str +``` + +Into: + +```py +from typing import Generic, TypeVar +from pydantic.generics import GenericModel + +T = TypeVar('T') + +class User(BaseModel, Generic[T]): + name: str +``` diff --git a/bump_pydantic/codemods/__init__.py b/bump_pydantic/codemods/__init__.py index 0f2059b..2e2f7af 100644 --- a/bump_pydantic/codemods/__init__.py +++ b/bump_pydantic/codemods/__init__.py @@ -1,10 +1,11 @@ from typing import List, Type from libcst.codemod import ContextAwareTransformer -from libcst.codemod.visitors import AddImportsVisitor +from libcst.codemod.visitors import AddImportsVisitor, RemoveImportsVisitor from bump_pydantic.codemods.add_default_none import AddDefaultNoneCommand from bump_pydantic.codemods.replace_config import ReplaceConfigCodemod +from bump_pydantic.codemods.replace_generic_model import ReplaceGenericModelCommand from bump_pydantic.codemods.replace_imports import ReplaceImportsCodemod @@ -13,6 +14,9 @@ def gather_codemods() -> List[Type[ContextAwareTransformer]]: AddDefaultNoneCommand, ReplaceConfigCodemod, ReplaceImportsCodemod, + ReplaceGenericModelCommand, + # RemoveImportsVisitor needs to be the second to last. + RemoveImportsVisitor, # AddImportsVisitor needs to be the last. AddImportsVisitor, ] diff --git a/bump_pydantic/codemods/replace_generic_model.py b/bump_pydantic/codemods/replace_generic_model.py new file mode 100644 index 0000000..9607bca --- /dev/null +++ b/bump_pydantic/codemods/replace_generic_model.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +import libcst as cst +import libcst.matchers as m +from libcst.codemod import VisitorBasedCodemodCommand +from libcst.codemod.visitors import AddImportsVisitor, RemoveImportsVisitor + +GENERIC_MODEL_ARG = m.Arg(value=m.Name("GenericModel")) | m.Arg( + value=m.Attribute(value=m.Name("generics"), attr=m.Name("GenericModel")) +) + + +class ReplaceGenericModelCommand(VisitorBasedCodemodCommand): + @m.leave(m.ClassDef(bases=[m.ZeroOrMore(), GENERIC_MODEL_ARG, m.ZeroOrMore()])) + def leave_generic_model(self, original_node: cst.ClassDef, updated_node: cst.ClassDef) -> cst.ClassDef: + RemoveImportsVisitor.remove_unused_import(context=self.context, module="pydantic.generics", obj="GenericModel") + RemoveImportsVisitor.remove_unused_import(context=self.context, module="pydantic", obj="generics") + AddImportsVisitor.add_needed_import(context=self.context, module="pydantic", obj="BaseModel") + return updated_node.with_changes( + bases=[ + base if not m.matches(base, GENERIC_MODEL_ARG) else cst.Arg(value=cst.Name("BaseModel")) + for base in updated_node.bases + ] + ) + + +if __name__ == "__main__": + import textwrap + + from rich.console import Console + + console = Console() + + source = textwrap.dedent( + """ + from typing import Generic, TypeVar + + from pydantic.generics import GenericModel + + T = TypeVar("T") + + class Potato(GenericModel, Generic[T]): + ... + """ + ) + console.print(source) + # console.print("=" * 80) + + # mod = cst.parse_module(source) + # context = CodemodContext(filename="main.py") + + # wrapper = cst.MetadataWrapper(mod) + # command = ReplaceGenericModelCommand(context=context) + # mod = wrapper.visit(command) + + # wrapper = cst.MetadataWrapper(mod) + # command = RemoveImportsVisitor(context=context) # type: ignore[assignment] + # mod = wrapper.visit(command) + + # wrapper = cst.MetadataWrapper(mod) + # command = AddImportsVisitor(context=context) # type: ignore[assignment] + # mod = wrapper.visit(command) + # console.print(mod.code) diff --git a/bump_pydantic/codemods/replace_imports.py b/bump_pydantic/codemods/replace_imports.py index 7baf634..43e6507 100644 --- a/bump_pydantic/codemods/replace_imports.py +++ b/bump_pydantic/codemods/replace_imports.py @@ -141,7 +141,6 @@ class Potato(BaseSettings): context = CodemodContext(filename="main.py") wrapper = cst.MetadataWrapper(mod) command = ReplaceImportsCodemod(context=context) - console.print(mod) mod = wrapper.visit(command) wrapper = cst.MetadataWrapper(mod) diff --git a/project/replace_generic.py b/project/replace_generic.py new file mode 100644 index 0000000..5625f21 --- /dev/null +++ b/project/replace_generic.py @@ -0,0 +1,9 @@ +from typing import Generic, TypeVar + +from pydantic.generics import GenericModel + +T = TypeVar("T") + + +class User(GenericModel, Generic[T]): + name: str diff --git a/pyproject.toml b/pyproject.toml index 07bea8c..a1b849f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -76,14 +76,15 @@ isort = { known-first-party = ['bump_pydantic', 'tests'] } target-version = 'py38' [tool.coverage.run] -source_pkgs = ["bump_pydantic", "tests"] +source_pkgs = ["bump_pydantic"] branch = true [tool.coverage.report] show_missing = true skip_covered = true -exclude_lines = ["no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:"] - -[tool.coverage.paths] -source = ["bump_pydantic/"] -detached = true +exclude_lines = [ + "pragma: nocover", + "pragma: no cover", + "if __name__ == .__main__.:", + "if TYPE_CHECKING:", +] diff --git a/tests/test_generic_model.py b/tests/test_generic_model.py new file mode 100644 index 0000000..18c1a8d --- /dev/null +++ b/tests/test_generic_model.py @@ -0,0 +1,102 @@ +from libcst.codemod import CodemodTest + +from bump_pydantic.codemods.replace_generic_model import ReplaceGenericModelCommand + + +class TestReplaceGenericModelCommand(CodemodTest): + TRANSFORM = ReplaceGenericModelCommand + + def test_noop(self) -> None: + code = """ + from typing import Generic, TypeVar + + T = TypeVar("T") + + class Potato(Generic[T]): + ... + """ + self.assertCodemod(code, code) + + def test_generic_model(self) -> None: + before = """ + from typing import TypeVar + from pydantic.generics import GenericModel + + T = TypeVar("T") + + class Potato(GenericModel, Generic[T]): + ... + """ + after = """ + from typing import TypeVar + from pydantic import BaseModel + + T = TypeVar("T") + + class Potato(BaseModel, Generic[T]): + ... + """ + self.assertCodemod(before, after) + + def test_generic_model_multiple_bases(self) -> None: + before = """ + from typing import TypeVar + from pydantic.generics import GenericModel + + T = TypeVar("T") + + class Potato(GenericModel, Generic[T], object): + ... + """ + after = """ + from typing import TypeVar + from pydantic import BaseModel + + T = TypeVar("T") + + class Potato(BaseModel, Generic[T], object): + ... + """ + self.assertCodemod(before, after) + + def test_generic_model_second_base(self) -> None: + before = """ + from typing import TypeVar + from pydantic.generics import GenericModel + + T = TypeVar("T") + + class Potato(object, GenericModel, Generic[T]): + ... + """ + after = """ + from typing import TypeVar + from pydantic import BaseModel + + T = TypeVar("T") + + class Potato(object, BaseModel, Generic[T]): + ... + """ + self.assertCodemod(before, after) + + def test_generic_model_from_pydantic_import_generics(self) -> None: + before = """ + from typing import TypeVar + from pydantic import generics + + T = TypeVar("T") + + class Potato(generics.GenericModel, Generic[T]): + ... + """ + after = """ + from typing import TypeVar + from pydantic import BaseModel + + T = TypeVar("T") + + class Potato(BaseModel, Generic[T]): + ... + """ + self.assertCodemod(before, after)