✨ Implement MyPyVisitor (#51)

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)