✨ Support GenericModel to BaseModel replacement (#12)

.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:

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
+```

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,
     ]

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)

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)

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

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:",
+]

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)