Class AnnotationPrinter

mypy/stubutil.py:227–323

Visitor used to print existing annotations in a file. The main difference from TypeStrVisitor is a better treatment of unbound types. Notes: * This visitor doesn't add imports necessary for annotations, this is done separately by ImportTracker. * It can print all kinds of

class AnnotationPrinter(TypeStrVisitor):
    """Visitor used to print existing annotations in a file.

    The main difference from TypeStrVisitor is a better treatment of
    unbound types.

    Notes:
    * This visitor doesn't add imports necessary for annotations, this is done separately
      by ImportTracker.
    * It can print all kinds of types, but the generated strings may not be valid (notably
      callable types) since it prints the same string that reveal_type() does.
    * For Instance types it prints the fully qualified names.
    """

    # TODO: Generate valid string representation for callable types.
    # TODO: Use short names for Instances.
    def __init__(
        self,
        stubgen: BaseStubGenerator,
        known_modules: list[str] | None = None,
        local_modules: list[str] | None = None,
    ) -> None:
        super().__init__(options=mypy.options.Options())
        self.stubgen = stubgen
        self.known_modules = known_modules
        self.local_modules = local_modules or ["builtins"]

    def visit_any(self, t: AnyType) -> str:
        s = super().visit_any(t)
        self.stubgen.import_tracker.require_name(s)
        return s

    def visit_unbound_type(self, t: UnboundType) -> str:
        s = t.name
        fullname = self.stubgen.resolve_name(s)
        if fullname == "typing.Union":
            return " | ".join([item.accept(self) for item in t.args])
        if fullname == "typing.Optional":
            if len(t.args) == 1:
                return f"{t.args[0].accept(self)} | None"
            return self.stubgen.add_name("_typeshed.Incomplete")
        if fullname in TYPING_BUILTIN_REPLACEMENTS:
            s = self.stubgen.add_name(TYPING_BUILTIN_REPLACEMENTS[fullname], require=True)
        if self.known_modules is not None and "." in s:
            # see if this object is from any of the modules that we're currently processing.
            # reverse sort so that subpackages come before parents: e.g. "foo.bar" before "foo".
            for module_name in self.local_modules + sorted(self.known_modules, reverse=True):
                if s.startswith(module_name + "."):
                    if module_name in self.local_modules:
                        s = s[len(module_name) + 1 :]
                    arg_module = module_name
                    break
            else:
                arg_module = s[: s.rindex(".")]
            if arg_module not in self.local_modules:
                self.stubgen.import_tracker.add_import(arg_module, require=True)
        elif s == "NoneType":
            # when called without analysis all types are unbound, so this won't hit