93 lines
2.7 KiB
Python
93 lines
2.7 KiB
Python
|
"""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
|