From de2613460a1031f276fcb142ca417ad8c7781e36 Mon Sep 17 00:00:00 2001
From: Jukka Lehtosalo <jukka.lehtosalo@iki.fi>
Date: Tue, 28 Feb 2023 18:34:44 +0000
Subject: [PATCH] [undocumented] Optionally export line-level information about
 references (#14805)

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.
---
 mypy/build.py   | 24 ++++++++++++++++++++++++
 mypy/main.py    |  2 ++
 mypy/options.py |  3 +++
 mypy/refinfo.py | 35 +++++++++++++++++++++++++++++++++++
 4 files changed, 64 insertions(+)
 create mode 100644 mypy/refinfo.py

diff --git a/mypy/build.py b/mypy/build.py
index a4817d1866c7..c99999d32a05 100644
--- 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
index 47dea2ae9797..3f5e02ec3f79 100644
--- 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
index 92c96a92c531..077e0d4ed90a 100644
--- 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
new file mode 100644
index 000000000000..4262824f8f97
--- /dev/null
+++ 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