✨ Implement MyPyVisitor

bump_pydantic/codemods/add_default_none.py

@@ -5,9 +5,7 @@
 from libcst.codemod import CodemodContext, VisitorBasedCodemodCommand
 from libcst.metadata import FullyQualifiedNameProvider, QualifiedName
 
-from bump_pydantic.codemods.class_def_visitor import ClassDefVisitor
-from bump_pydantic.markers.find_base_model import CONTEXT_KEY as BASE_MODEL_CONTEXT_KEY
-from bump_pydantic.markers.find_base_model import find_base_model
+from bump_pydantic.codemods.mypy_visitor import CONTEXT_KEY
 
 
 class AddDefaultNoneCommand(VisitorBasedCodemodCommand):
@@ -51,7 +49,7 @@ def visit_ClassDef(self, node: cst.ClassDef) -> None:
             return None
 
         fqn: QualifiedName = next(iter(fqn_set))  # type: ignore
-        if fqn.name in self.context.scratch[BASE_MODEL_CONTEXT_KEY]:
+        if self.context.scratch[CONTEXT_KEY].get(fqn.name, False):
             self.inside_base_model = True
 
     def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef) -> cst.ClassDef:
@@ -94,7 +92,6 @@ def leave_AnnAssign(self, original_node: cst.AnnAssign, updated_node: cst.AnnAss
     from tempfile import TemporaryDirectory
 
     from libcst.metadata import FullRepoManager
-    from rich.pretty import pprint
 
     with TemporaryDirectory(dir=os.getcwd()) as tmpdir:
         package_dir = f"{tmpdir}/package"
@@ -123,11 +120,8 @@ class Bar(Foo):
         wrapper = mrg.get_metadata_wrapper_for_path(module)
         context = CodemodContext(wrapper=wrapper)
 
-        command = ClassDefVisitor(context=context)
-        mod = wrapper.visit(command)
-
-        find_base_model(scratch=context.scratch)
-        pprint(context.scratch)
+        # classes = run_mypy_visitor(context=context)
+        # mod = wrapper.visit(command)
 
         command = AddDefaultNoneCommand(context=context)  # type: ignore[assignment]
         mod = wrapper.visit(command)

bump_pydantic/codemods/mypy_visitor.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+import sys
+from argparse import ArgumentParser
+
+from mypy.build import build
+from mypy.main import process_options
+from mypy.nodes import ClassDef
+from mypy.traverser import TraverserVisitor
+
+CONTEXT_KEY = "mypy_visitor"
+
+
+class MyPyVisitor(TraverserVisitor):
+    def __init__(self) -> None:
+        super().__init__()
+        self.classes: dict[str, bool] = {}
+
+    def visit_class_def(self, o: ClassDef) -> None:
+        super().visit_class_def(o)
+        self.classes[o.fullname] = o.info.has_base("pydantic.main.BaseModel")
+
+
+def run_mypy_visitor(arg_files: list[str]) -> dict[str, bool]:
+    files, opt = process_options(arg_files, stdout=sys.stdout, stderr=sys.stderr)
+
+    opt.export_types = True
+    opt.incremental = True
+    opt.fine_grained_incremental = True
+    opt.cache_fine_grained = True
+    opt.allow_redefinition = True
+    opt.local_partial_types = True
+
+    result = build(files, opt, stdout=sys.stdout, stderr=sys.stderr)
+
+    visitor = MyPyVisitor()
+    classes: dict[str, bool] = {}
+
+    for file in files:
+        tree = result.graph[file.module].tree
+        if tree:
+            tree.accept(visitor=visitor)
+            classes.update(visitor.classes)
+    return classes
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument("files", nargs="+")
+    args = parser.parse_args()
+
+    run_mypy_visitor(args.files)