[undocumented] Optionally export line-level information about references (#14805)

JukkaL · web-flow · commit de2613460a10 · 2023-02-28T18:34:44.000Z
When run with `--export-ref-info`, store line-level reference
information in the cache, in a JSON file with a `.refs.json` extension.
It includes the line numbers and targets of each RefExpr in the program.
This only works properly if incremental mode is disabled.

The target can either be a fullname, or `*.name` for an `name` attribute
reference where the type of the object is unknown.

This is an undocumented, experimental feature that may be useful for
certain tools, but it shouldn't be used in production use cases.
diff --git a/mypy/build.py b/mypy/build.py
@@ -2410,6 +2410,10 @@ def finish_passes(self) -> None:
             manager.report_file(self.tree, self.type_map(), self.options)
 
             self.update_fine_grained_deps(self.manager.fg_deps)
+
+            if manager.options.export_ref_info:
+                write_undocumented_ref_info(self, manager.metastore, manager.options)
+
             self.free_state()
             if not manager.options.fine_grained_incremental and not manager.options.preserve_asts:
                 free_tree(self.tree)
@@ -2941,6 +2945,7 @@ def dispatch(sources: list[BuildSource], manager: BuildManager, stdout: TextIO)
         dump_all_dependencies(
             manager.modules, manager.all_types, manager.options.python_version, manager.options
         )
+
     return graph
 
 
@@ -3616,3 +3621,22 @@ def is_silent_import_module(manager: BuildManager, path: str) -> bool:
         is_sub_path(path, dir)
         for dir in manager.search_paths.package_path + manager.search_paths.typeshed_path
     )
+
+
+def write_undocumented_ref_info(state: State, metastore: MetadataStore, options: Options) -> None:
+    # This exports some dependency information in a rather ad-hoc fashion, which
+    # can be helpful for some tools. This is all highly experimental and could be
+    # removed at any time.
+
+    from mypy.refinfo import get_undocumented_ref_info_json
+
+    if not state.tree:
+        # We need a full AST for this.
+        return
+
+    _, data_file, _ = get_cache_names(state.id, state.xpath, options)
+    ref_info_file = ".".join(data_file.split(".")[:-2]) + ".refs.json"
+    assert not ref_info_file.startswith(".")
+
+    deps_json = get_undocumented_ref_info_json(state.tree)
+    metastore.write(ref_info_file, json.dumps(deps_json))
diff --git a/mypy/main.py b/mypy/main.py
@@ -1017,6 +1017,8 @@ def add_invertible_flag(
     add_invertible_flag(
         "--allow-empty-bodies", default=False, help=argparse.SUPPRESS, group=internals_group
     )
+    # This undocumented feature exports limited line-level dependency information.
+    internals_group.add_argument("--export-ref-info", action="store_true", help=argparse.SUPPRESS)
 
     report_group = parser.add_argument_group(
         title="Report generation", description="Generate a report in the specified format."
diff --git a/mypy/options.py b/mypy/options.py
@@ -339,6 +339,9 @@ def __init__(self) -> None:
         self.disable_recursive_aliases = False
         # Deprecated reverse version of the above, do not use.
         self.enable_recursive_aliases = False
+        # Export line-level, limited, fine-grained dependency information in cache data
+        # (undocumented feature).
+        self.export_ref_info = False
 
         self.disable_bytearray_promotion = False
         self.disable_memoryview_promotion = False
diff --git a/mypy/refinfo.py b/mypy/refinfo.py
@@ -0,0 +1,35 @@
+"""Find line-level reference information from a mypy AST (undocumented feature)"""
+
+from __future__ import annotations
+
+from mypy.nodes import LDEF, MemberExpr, MypyFile, NameExpr, RefExpr
+from mypy.traverser import TraverserVisitor
+
+
+class RefInfoVisitor(TraverserVisitor):
+    def __init__(self) -> None:
+        super().__init__()
+        self.data: list[dict[str, object]] = []
+
+    def visit_name_expr(self, expr: NameExpr) -> None:
+        super().visit_name_expr(expr)
+        self.record_ref_expr(expr)
+
+    def visit_member_expr(self, expr: MemberExpr) -> None:
+        super().visit_member_expr(expr)
+        self.record_ref_expr(expr)
+
+    def record_ref_expr(self, expr: RefExpr) -> None:
+        fullname = None
+        if expr.kind != LDEF and "." in expr.fullname:
+            fullname = expr.fullname
+        elif isinstance(expr, MemberExpr) and not expr.fullname:
+            fullname = f"*.{expr.name}"
+        if fullname is not None:
+            self.data.append({"line": expr.line, "column": expr.column, "target": fullname})
+
+
+def get_undocumented_ref_info_json(tree: MypyFile) -> list[dict[str, object]]:
+    visitor = RefInfoVisitor()
+    tree.accept(visitor)
+    return visitor.data
