Fix --export-ref-info with type variable with value restrictions

mypy/refinfo.py

@@ -5,6 +5,7 @@
 from mypy.nodes import (
     LDEF,
     Expression,
+    FuncDef,
     MemberExpr,
     MypyFile,
     NameExpr,
@@ -39,6 +40,14 @@ def visit_member_expr(self, expr: MemberExpr) -> None:
         super().visit_member_expr(expr)
         self.record_ref_expr(expr)
 
+    def visit_func_def(self, func: FuncDef) -> None:
+        if func.expanded:
+            for item in func.expanded:
+                if isinstance(item, FuncDef):
+                    super().visit_func_def(item)
+        else:
+            super().visit_func_def(func)
+
     def record_ref_expr(self, expr: RefExpr) -> None:
         fullname = None
         if expr.kind != LDEF and "." in expr.fullname:

mypy/test/test_ref_info.py

@@ -0,0 +1,45 @@
+"""Test exporting line-level reference information (undocumented feature)"""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+
+from mypy import build
+from mypy.modulefinder import BuildSource
+from mypy.options import Options
+from mypy.test.config import test_temp_dir
+from mypy.test.data import DataDrivenTestCase, DataSuite
+from mypy.test.helpers import assert_string_arrays_equal
+
+
+class RefInfoSuite(DataSuite):
+    required_out_section = True
+    files = ["ref-info.test"]
+
+    def run_case(self, testcase: DataDrivenTestCase) -> None:
+        options = Options()
+        options.use_builtins_fixtures = True
+        options.show_traceback = True
+        options.export_ref_info = True  # This is the flag we are testing
+
+        src = "\n".join(testcase.input)
+        result = build.build(
+            sources=[BuildSource("main", None, src)], options=options, alt_lib_path=test_temp_dir
+        )
+        assert not result.errors
+
+        major, minor = sys.version_info[:2]
+        ref_path = os.path.join(options.cache_dir, f"{major}.{minor}", "__main__.refs.json")
+
+        with open(ref_path) as refs_file:
+            data = json.load(refs_file)
+
+        a = []
+        for item in data:
+            a.append(f"{item['line']}:{item['column']}:{item['target']}")
+
+        assert_string_arrays_equal(
+            testcase.output, a, f"Invalid output ({testcase.file}, line {testcase.line})"
+        )

test-data/unit/ref-info.test

@@ -0,0 +1,83 @@
+[case testCallGlobalFunction]
+def f() -> None:
+    g()
+
+def g() -> None:
+    pass
+[out]
+2:4:__main__.g
+
+[case testCallMethod]
+def f() -> None:
+    c = C()
+    if int():
+        c.method()
+
+class C:
+    def method(self) -> None: pass
+[out]
+2:8:__main__.C
+3:7:builtins.int
+4:8:__main__.C.method
+
+[case testCallStaticMethod]
+class C:
+    def f(self) -> None:
+        C.static()
+        self.static()
+
+    @classmethod
+    def cm(cls) -> None:
+        cls.static()
+
+    @staticmethod
+    def static() -> None: pass
+[builtins fixtures/classmethod.pyi]
+[out]
+3:8:__main__.C
+3:8:__main__.C.static
+4:8:__main__.C.static
+8:8:__main__.C.static
+
+[case testCallClassMethod]
+class C:
+    def f(self) -> None:
+        C.cm()
+        self.cm()
+
+    @classmethod
+    def cm(cls) -> None:
+        cls.cm()
+[builtins fixtures/classmethod.pyi]
+[out]
+3:8:__main__.C
+3:8:__main__.C.cm
+4:8:__main__.C.cm
+8:8:__main__.C.cm
+
+[case testTypeVarWithValueRestriction]
+from typing import TypeVar
+
+T = TypeVar("T", "C", "D")
+
+def f(o: T) -> None:
+    f(o)
+    o.m()
+    o.x
+
+class C:
+    x: int
+    def m(self) -> None: pass
+
+class D:
+    x: str
+    def m(self) -> None: pass
+[out]
+3:4:typing.TypeVar
+3:0:__main__.T
+6:4:__main__.f
+7:4:__main__.C.m
+8:4:__main__.C.x
+6:4:__main__.f
+7:4:__main__.D.m
+8:4:__main__.D.x

test-data/unit/typexport-basic.test

@@ -1055,6 +1055,21 @@ CallExpr(7) : builtins.str
 NameExpr(7) : def (x: builtins.str) -> builtins.str
 NameExpr(7) : S
 
+[case testTypeVariableWithValueRestrictionInFunction]
+## NameExpr
+from typing import TypeVar
+
+T = TypeVar("T", int, str)
+
+def f(x: T) -> T:
+    y = 1
+    return x
+[out]
+NameExpr(7) : builtins.int
+NameExpr(7) : builtins.int
+NameExpr(8) : builtins.int
+NameExpr(8) : builtins.str
+
 
 -- Binary operations
 -- -----------------