Class AliasPrinter

mypy/stubgen.py:281–428

Visitor used to collect type aliases _and_ type variable definitions. Visit r.h.s of the definition to get the string representation of type alias.

class AliasPrinter(NodeVisitor[str]):
    """Visitor used to collect type aliases _and_ type variable definitions.

    Visit r.h.s of the definition to get the string representation of type alias.
    """

    def __init__(self, stubgen: ASTStubGenerator) -> None:
        self.stubgen = stubgen
        super().__init__()

    def visit_call_expr(self, node: CallExpr) -> str:
        # Call expressions are not usually types, but we also treat `X = TypeVar(...)` as a
        # type alias that has to be preserved (even if TypeVar is not the same as an alias)
        callee = node.callee.accept(self)
        args = []
        for name, arg, kind in zip(node.arg_names, node.args, node.arg_kinds):
            if kind == ARG_POS:
                args.append(arg.accept(self))
            elif kind == ARG_STAR:
                args.append("*" + arg.accept(self))
            elif kind == ARG_STAR2:
                args.append("**" + arg.accept(self))
            elif kind == ARG_NAMED:
                args.append(f"{name}={arg.accept(self)}")
            else:
                raise ValueError(f"Unknown argument kind {kind} in call")
        return f"{callee}({', '.join(args)})"

    def _visit_ref_expr(self, node: NameExpr | MemberExpr) -> str:
        fullname = self.stubgen.get_fullname(node)
        if fullname in TYPING_BUILTIN_REPLACEMENTS:
            return self.stubgen.add_name(TYPING_BUILTIN_REPLACEMENTS[fullname], require=False)
        qualname = get_qualified_name(node)
        self.stubgen.import_tracker.require_name(qualname)
        return qualname

    def visit_name_expr(self, node: NameExpr) -> str:
        return self._visit_ref_expr(node)

    def visit_member_expr(self, o: MemberExpr) -> str:
        return self._visit_ref_expr(o)

    def _visit_literal_node(
        self, node: StrExpr | BytesExpr | IntExpr | FloatExpr | ComplexExpr
    ) -> str:
        return repr(node.value)

    def visit_str_expr(self, node: StrExpr) -> str:
        return self._visit_literal_node(node)

    def visit_bytes_expr(self, node: BytesExpr) -> str:
        return f"b{self._visit_literal_node(node)}"

    def visit_int_expr(self, node: IntExpr) -> str:
        return self._visit_literal_node(node)

    def visit_float_expr(self, node: FloatExpr) -> str:
        return self._visit_literal_node(node)