mypy/refinfo.py

"""Find line-level reference information from a mypy AST (undocumented feature)"""

from __future__ import annotations

from mypy.nodes import (
    LDEF,
    Expression,
    FuncDef,
    MemberExpr,
    MypyFile,
    NameExpr,
    RefExpr,
    SymbolNode,
    TypeInfo,
)
from mypy.traverser import TraverserVisitor
from mypy.typeops import tuple_fallback
from mypy.types import (
    FunctionLike,
    Instance,
    TupleType,
    Type,
    TypeType,
    TypeVarLikeType,
    get_proper_type,
)


class RefInfoVisitor(TraverserVisitor):
    def __init__(self, type_map: dict[Expression, Type]) -> None:
        super().__init__()
        self.type_map = type_map
        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 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:
            fullname = expr.fullname
        elif isinstance(expr, MemberExpr):
            typ = self.type_map.get(expr.expr)
            sym = None
            if isinstance(expr.expr, RefExpr):
                sym = expr.expr.node
            if typ:
                tfn = type_fullname(typ, sym)
                if tfn:
                    fullname = f"{tfn}.{expr.name}"
            if not fullname:
                fullname = f"*.{expr.name}"
        if fullname is not None:
            self.data.append({"line": expr.line, "column": expr.column, "target": fullname})


def type_fullname(typ: Type, node: SymbolNode | None = None) -> str | None:
    typ = get_proper_type(typ)
    if isinstance(typ, Instance):
        return typ.type.fullname
    elif isinstance(typ, TypeType):
        return type_fullname(typ.item)
    elif isinstance(typ, FunctionLike) and typ.is_type_obj():
        if isinstance(node, TypeInfo):
            return node.fullname
        return type_fullname(typ.fallback)
    elif isinstance(typ, TupleType):
        return type_fullname(tuple_fallback(typ))
    elif isinstance(typ, TypeVarLikeType):
        return type_fullname(typ.upper_bound)
    return None


def get_undocumented_ref_info_json(
    tree: MypyFile, type_map: dict[Expression, Type]
) -> list[dict[str, object]]:
    visitor = RefInfoVisitor(type_map)
    tree.accept(visitor)
    return visitor.data