"""
API for analyzing Python stubs using mypy.
"""
import os
from os.path import abspath
from typing import Iterable, Optional, Set, Tuple
from mypy.build import BuildResult, State, build
from mypy.main import process_options
from mypy.nodes import (
GDEF,
MDEF,
Decorator,
FuncDef,
MypyFile,
OverloadedFuncDef,
SymbolNode,
TypeAlias,
TypeInfo,
TypeVarExpr,
Var,
)
from .types import RelevantSymbolNode
IGNORED_MODULE_SYMBOLS = [
# because docstrings are not generated by stubgen
"__doc__",
# created by the import system, do not need to be handwritten:
"__name__",
"__package__",
"__file__",
"__path__",
]
"""
Module level definitions that will not be collected by collect_types.
"""
def _mypy_analyze(
mypy_conf_path: str, root_path: str, stubs_path: Optional[str] = None
) -> BuildResult:
"""
Parse and analyze the types of the code in root_path.
:param mypy_conf_path: path to a mypy.ini
:param root_path: path to the code directory where the type analysis is started
:param stubs_path: path to the directory of stubs for mypy to use
:returns: Mypy's analysis result
"""
# The call to `build.build` is inspired by `mypy/mypy/main.py::main`
# `build` is not a documented public API
args = [
"--config-file",
mypy_conf_path,
"--no-incremental",
"--cache-dir=" + ("nul" if os.name == "nt" else "/dev/null"),
root_path,
]
sources, options = process_options(args)
if stubs_path is not None:
options = options.apply_changes({"mypy_path": [stubs_path]})
return build(sources, options, None, None)
[docs]def is_stubbed_module(module: State) -> bool:
"""
Check if a module's types were loaded from a stub.
:param module: The module to check
"""
return module.path is not None and module.path.endswith(".pyi")
[docs]def collect_types(
symbol_node: SymbolNode, collected_types: Optional[Set[str]] = None
) -> Iterable[RelevantSymbolNode]:
"""
Collect all relevant type definitions of the symbols in the given node.
:param symbol_node: any symbol node, e.g. MypyFile (BuildResult.graph.tree)
:param collected_types: used to avoid collecting duplicates
"""
if not collected_types:
collected_types = set()
# ignore builtins because we don't provide stubs for them
if "builtins" in symbol_node.fullname:
return
# do not collect types twice
if symbol_node.fullname in collected_types:
return
collected_types.add(symbol_node.fullname)
if isinstance(symbol_node, MypyFile):
# the symbol node represents a Python module
for symbol in symbol_node.names.values():
# only global and class member definitions are interesting
if symbol.kind not in [GDEF, MDEF]:
continue
if symbol.node and symbol.node.name in IGNORED_MODULE_SYMBOLS:
continue
if symbol.node and symbol.module_public:
yield from collect_types(symbol.node, collected_types)
elif isinstance(symbol_node, TypeInfo):
# the symbol represents a class definition
yield symbol_node
for class_member in symbol_node.names.values():
if class_member.node:
yield from collect_types(class_member.node, collected_types)
elif isinstance(
symbol_node,
(Decorator, FuncDef, OverloadedFuncDef, Var, TypeAlias, TypeVarExpr),
):
# the symbol represents a function definition,
# variable, type alias or generic TypeVar
yield symbol_node
else:
assert False, f"Unexpected symbol type {type(symbol_node)}"
[docs]def get_stub_types(
stubs_path: str, mypy_conf_path: str, root_path: Optional[str] = None
) -> Iterable[Tuple[RelevantSymbolNode, str]]:
"""
Analyze the stub files in stubs_path and return module
and class definitions of stubs as symbol nodes.
Only relevant symbol nodes (e.g. for variables, functions, classes, methods)
are returned.
They contain the type annotation information.
:param stubs_path: where all the stub files are located
:param mypy_conf_path: path to mypy.ini
:param root_path: path to the code directory where the type analysis is started
"""
stubs_path = abspath(stubs_path)
if root_path:
build_result = _mypy_analyze(mypy_conf_path, root_path, stubs_path)
else:
build_result = _mypy_analyze(mypy_conf_path, stubs_path)
stubbed_modules = {
module
for module in build_result.graph.values()
if module.path
and is_stubbed_module(module)
and module.path.startswith(stubs_path)
}
for module in stubbed_modules:
if module.tree:
assert module.path
yield from (
(stub_type, module.path) for stub_type in collect_types(module.tree)
)