| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994 |
- #!/usr/bin/env python3
- """Generator of dynamically typed draft stubs for arbitrary modules.
- The logic of this script can be split in three steps:
- * parsing options and finding sources:
- - use runtime imports be default (to find also C modules)
- - or use mypy's mechanisms, if importing is prohibited
- * (optionally) semantically analysing the sources using mypy (as a single set)
- * emitting the stubs text:
- - for Python modules: from ASTs using StubGenerator
- - for C modules using runtime introspection and (optionally) Sphinx docs
- During first and third steps some problematic files can be skipped, but any
- blocking error during second step will cause the whole program to stop.
- Basic usage:
- $ stubgen foo.py bar.py some_directory
- => Generate out/foo.pyi, out/bar.pyi, and stubs for some_directory (recursively).
- $ stubgen -m urllib.parse
- => Generate out/urllib/parse.pyi.
- $ stubgen -p urllib
- => Generate stubs for whole urlib package (recursively).
- For C modules, you can get more precise function signatures by parsing .rst (Sphinx)
- documentation for extra information. For this, use the --doc-dir option:
- $ stubgen --doc-dir <DIR>/Python-3.4.2/Doc/library -m curses
- Note: The generated stubs should be verified manually.
- TODO:
- - maybe use .rst docs also for Python modules
- - maybe export more imported names if there is no __all__ (this affects ssl.SSLError, for example)
- - a quick and dirty heuristic would be to turn this on if a module has something like
- 'from x import y as _y'
- - we don't seem to always detect properties ('closed' in 'io', for example)
- """
- from __future__ import annotations
- import argparse
- import glob
- import keyword
- import os
- import os.path
- import sys
- import traceback
- from collections import defaultdict
- from typing import Final, Iterable, Mapping
- import mypy.build
- import mypy.mixedtraverser
- import mypy.parse
- import mypy.traverser
- import mypy.util
- from mypy.build import build
- from mypy.errors import CompileError, Errors
- from mypy.find_sources import InvalidSourceList, create_source_list
- from mypy.modulefinder import (
- BuildSource,
- FindModuleCache,
- ModuleNotFoundReason,
- SearchPaths,
- default_lib_path,
- )
- from mypy.moduleinspect import ModuleInspect
- from mypy.nodes import (
- ARG_NAMED,
- ARG_POS,
- ARG_STAR,
- ARG_STAR2,
- IS_ABSTRACT,
- NOT_ABSTRACT,
- AssignmentStmt,
- Block,
- BytesExpr,
- CallExpr,
- ClassDef,
- ComparisonExpr,
- ComplexExpr,
- Decorator,
- DictExpr,
- EllipsisExpr,
- Expression,
- FloatExpr,
- FuncBase,
- FuncDef,
- IfStmt,
- Import,
- ImportAll,
- ImportFrom,
- IndexExpr,
- IntExpr,
- ListExpr,
- MemberExpr,
- MypyFile,
- NameExpr,
- OpExpr,
- OverloadedFuncDef,
- Statement,
- StrExpr,
- TupleExpr,
- TypeInfo,
- UnaryExpr,
- Var,
- )
- from mypy.options import Options as MypyOptions
- from mypy.stubdoc import Sig, find_unique_signatures, parse_all_signatures
- from mypy.stubgenc import (
- DocstringSignatureGenerator,
- ExternalSignatureGenerator,
- FallbackSignatureGenerator,
- SignatureGenerator,
- generate_stub_for_c_module,
- )
- from mypy.stubutil import (
- CantImport,
- common_dir_prefix,
- fail_missing,
- find_module_path_and_all_py3,
- generate_guarded,
- remove_misplaced_type_comments,
- report_missing,
- walk_packages,
- )
- from mypy.traverser import (
- all_yield_expressions,
- has_return_statement,
- has_yield_expression,
- has_yield_from_expression,
- )
- from mypy.types import (
- OVERLOAD_NAMES,
- TPDICT_NAMES,
- TYPED_NAMEDTUPLE_NAMES,
- AnyType,
- CallableType,
- Instance,
- NoneType,
- TupleType,
- Type,
- TypeList,
- TypeStrVisitor,
- UnboundType,
- UnionType,
- get_proper_type,
- )
- from mypy.visitor import NodeVisitor
- TYPING_MODULE_NAMES: Final = ("typing", "typing_extensions")
- # Common ways of naming package containing vendored modules.
- VENDOR_PACKAGES: Final = ["packages", "vendor", "vendored", "_vendor", "_vendored_packages"]
- # Avoid some file names that are unnecessary or likely to cause trouble (\n for end of path).
- BLACKLIST: Final = [
- "/six.py\n", # Likely vendored six; too dynamic for us to handle
- "/vendored/", # Vendored packages
- "/vendor/", # Vendored packages
- "/_vendor/",
- "/_vendored_packages/",
- ]
- # Special-cased names that are implicitly exported from the stub (from m import y as y).
- EXTRA_EXPORTED: Final = {
- "pyasn1_modules.rfc2437.univ",
- "pyasn1_modules.rfc2459.char",
- "pyasn1_modules.rfc2459.univ",
- }
- # These names should be omitted from generated stubs.
- IGNORED_DUNDERS: Final = {
- "__all__",
- "__author__",
- "__version__",
- "__about__",
- "__copyright__",
- "__email__",
- "__license__",
- "__summary__",
- "__title__",
- "__uri__",
- "__str__",
- "__repr__",
- "__getstate__",
- "__setstate__",
- "__slots__",
- }
- # These methods are expected to always return a non-trivial value.
- METHODS_WITH_RETURN_VALUE: Final = {
- "__ne__",
- "__eq__",
- "__lt__",
- "__le__",
- "__gt__",
- "__ge__",
- "__hash__",
- "__iter__",
- }
- # These magic methods always return the same type.
- KNOWN_MAGIC_METHODS_RETURN_TYPES: Final = {
- "__len__": "int",
- "__length_hint__": "int",
- "__init__": "None",
- "__del__": "None",
- "__bool__": "bool",
- "__bytes__": "bytes",
- "__format__": "str",
- "__contains__": "bool",
- "__complex__": "complex",
- "__int__": "int",
- "__float__": "float",
- "__index__": "int",
- }
- class Options:
- """Represents stubgen options.
- This class is mutable to simplify testing.
- """
- def __init__(
- self,
- pyversion: tuple[int, int],
- no_import: bool,
- doc_dir: str,
- search_path: list[str],
- interpreter: str,
- parse_only: bool,
- ignore_errors: bool,
- include_private: bool,
- output_dir: str,
- modules: list[str],
- packages: list[str],
- files: list[str],
- verbose: bool,
- quiet: bool,
- export_less: bool,
- ) -> None:
- # See parse_options for descriptions of the flags.
- self.pyversion = pyversion
- self.no_import = no_import
- self.doc_dir = doc_dir
- self.search_path = search_path
- self.interpreter = interpreter
- self.decointerpreter = interpreter
- self.parse_only = parse_only
- self.ignore_errors = ignore_errors
- self.include_private = include_private
- self.output_dir = output_dir
- self.modules = modules
- self.packages = packages
- self.files = files
- self.verbose = verbose
- self.quiet = quiet
- self.export_less = export_less
- class StubSource:
- """A single source for stub: can be a Python or C module.
- A simple extension of BuildSource that also carries the AST and
- the value of __all__ detected at runtime.
- """
- def __init__(
- self, module: str, path: str | None = None, runtime_all: list[str] | None = None
- ) -> None:
- self.source = BuildSource(path, module, None)
- self.runtime_all = runtime_all
- self.ast: MypyFile | None = None
- @property
- def module(self) -> str:
- return self.source.module
- @property
- def path(self) -> str | None:
- return self.source.path
- # What was generated previously in the stub file. We keep track of these to generate
- # nicely formatted output (add empty line between non-empty classes, for example).
- EMPTY: Final = "EMPTY"
- FUNC: Final = "FUNC"
- CLASS: Final = "CLASS"
- EMPTY_CLASS: Final = "EMPTY_CLASS"
- VAR: Final = "VAR"
- NOT_IN_ALL: Final = "NOT_IN_ALL"
- # Indicates that we failed to generate a reasonable output
- # for a given node. These should be manually replaced by a user.
- ERROR_MARKER: Final = "<ERROR>"
- 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: StubGenerator) -> None:
- super().__init__(options=mypy.options.Options())
- self.stubgen = stubgen
- 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
- self.stubgen.import_tracker.require_name(s)
- if t.args:
- s += f"[{self.args_str(t.args)}]"
- return s
- def visit_none_type(self, t: NoneType) -> str:
- return "None"
- def visit_type_list(self, t: TypeList) -> str:
- return f"[{self.list_str(t.items)}]"
- def visit_union_type(self, t: UnionType) -> str:
- return " | ".join([item.accept(self) for item in t.items])
- def args_str(self, args: Iterable[Type]) -> str:
- """Convert an array of arguments to strings and join the results with commas.
- The main difference from list_str is the preservation of quotes for string
- arguments
- """
- types = ["builtins.bytes", "builtins.str"]
- res = []
- for arg in args:
- arg_str = arg.accept(self)
- if isinstance(arg, UnboundType) and arg.original_str_fallback in types:
- res.append(f"'{arg_str}'")
- else:
- res.append(arg_str)
- return ", ".join(res)
- 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: StubGenerator) -> 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_name_expr(self, node: NameExpr) -> str:
- self.stubgen.import_tracker.require_name(node.name)
- return node.name
- def visit_member_expr(self, o: MemberExpr) -> str:
- node: Expression = o
- trailer = ""
- while isinstance(node, MemberExpr):
- trailer = "." + node.name + trailer
- node = node.expr
- if not isinstance(node, NameExpr):
- return ERROR_MARKER
- self.stubgen.import_tracker.require_name(node.name)
- return node.name + trailer
- def visit_str_expr(self, node: StrExpr) -> str:
- return repr(node.value)
- def visit_index_expr(self, node: IndexExpr) -> str:
- base = node.base.accept(self)
- index = node.index.accept(self)
- if len(index) > 2 and index.startswith("(") and index.endswith(")"):
- index = index[1:-1]
- return f"{base}[{index}]"
- def visit_tuple_expr(self, node: TupleExpr) -> str:
- return f"({', '.join(n.accept(self) for n in node.items)})"
- def visit_list_expr(self, node: ListExpr) -> str:
- return f"[{', '.join(n.accept(self) for n in node.items)}]"
- def visit_dict_expr(self, o: DictExpr) -> str:
- dict_items = []
- for key, value in o.items:
- # This is currently only used for TypedDict where all keys are strings.
- assert isinstance(key, StrExpr)
- dict_items.append(f"{key.accept(self)}: {value.accept(self)}")
- return f"{{{', '.join(dict_items)}}}"
- def visit_ellipsis(self, node: EllipsisExpr) -> str:
- return "..."
- def visit_op_expr(self, o: OpExpr) -> str:
- return f"{o.left.accept(self)} {o.op} {o.right.accept(self)}"
- class ImportTracker:
- """Record necessary imports during stub generation."""
- def __init__(self) -> None:
- # module_for['foo'] has the module name where 'foo' was imported from, or None if
- # 'foo' is a module imported directly; examples
- # 'from pkg.m import f as foo' ==> module_for['foo'] == 'pkg.m'
- # 'from m import f' ==> module_for['f'] == 'm'
- # 'import m' ==> module_for['m'] == None
- # 'import pkg.m' ==> module_for['pkg.m'] == None
- # ==> module_for['pkg'] == None
- self.module_for: dict[str, str | None] = {}
- # direct_imports['foo'] is the module path used when the name 'foo' was added to the
- # namespace.
- # import foo.bar.baz ==> direct_imports['foo'] == 'foo.bar.baz'
- # ==> direct_imports['foo.bar'] == 'foo.bar.baz'
- # ==> direct_imports['foo.bar.baz'] == 'foo.bar.baz'
- self.direct_imports: dict[str, str] = {}
- # reverse_alias['foo'] is the name that 'foo' had originally when imported with an
- # alias; examples
- # 'import numpy as np' ==> reverse_alias['np'] == 'numpy'
- # 'import foo.bar as bar' ==> reverse_alias['bar'] == 'foo.bar'
- # 'from decimal import Decimal as D' ==> reverse_alias['D'] == 'Decimal'
- self.reverse_alias: dict[str, str] = {}
- # required_names is the set of names that are actually used in a type annotation
- self.required_names: set[str] = set()
- # Names that should be reexported if they come from another module
- self.reexports: set[str] = set()
- def add_import_from(self, module: str, names: list[tuple[str, str | None]]) -> None:
- for name, alias in names:
- if alias:
- # 'from {module} import {name} as {alias}'
- self.module_for[alias] = module
- self.reverse_alias[alias] = name
- else:
- # 'from {module} import {name}'
- self.module_for[name] = module
- self.reverse_alias.pop(name, None)
- self.direct_imports.pop(alias or name, None)
- def add_import(self, module: str, alias: str | None = None) -> None:
- if alias:
- # 'import {module} as {alias}'
- self.module_for[alias] = None
- self.reverse_alias[alias] = module
- else:
- # 'import {module}'
- name = module
- # add module and its parent packages
- while name:
- self.module_for[name] = None
- self.direct_imports[name] = module
- self.reverse_alias.pop(name, None)
- name = name.rpartition(".")[0]
- def require_name(self, name: str) -> None:
- self.required_names.add(name.split(".")[0])
- def reexport(self, name: str) -> None:
- """Mark a given non qualified name as needed in __all__.
- This means that in case it comes from a module, it should be
- imported with an alias even is the alias is the same as the name.
- """
- self.require_name(name)
- self.reexports.add(name)
- def import_lines(self) -> list[str]:
- """The list of required import lines (as strings with python code)."""
- result = []
- # To summarize multiple names imported from a same module, we collect those
- # in the `module_map` dictionary, mapping a module path to the list of names that should
- # be imported from it. the names can also be alias in the form 'original as alias'
- module_map: Mapping[str, list[str]] = defaultdict(list)
- for name in sorted(self.required_names):
- # If we haven't seen this name in an import statement, ignore it
- if name not in self.module_for:
- continue
- m = self.module_for[name]
- if m is not None:
- # This name was found in a from ... import ...
- # Collect the name in the module_map
- if name in self.reverse_alias:
- name = f"{self.reverse_alias[name]} as {name}"
- elif name in self.reexports:
- name = f"{name} as {name}"
- module_map[m].append(name)
- else:
- # This name was found in an import ...
- # We can already generate the import line
- if name in self.reverse_alias:
- source = self.reverse_alias[name]
- result.append(f"import {source} as {name}\n")
- elif name in self.reexports:
- assert "." not in name # Because reexports only has nonqualified names
- result.append(f"import {name} as {name}\n")
- else:
- result.append(f"import {self.direct_imports[name]}\n")
- # Now generate all the from ... import ... lines collected in module_map
- for module, names in sorted(module_map.items()):
- result.append(f"from {module} import {', '.join(sorted(names))}\n")
- return result
- def find_defined_names(file: MypyFile) -> set[str]:
- finder = DefinitionFinder()
- file.accept(finder)
- return finder.names
- class DefinitionFinder(mypy.traverser.TraverserVisitor):
- """Find names of things defined at the top level of a module."""
- # TODO: Assignment statements etc.
- def __init__(self) -> None:
- # Short names of things defined at the top level.
- self.names: set[str] = set()
- def visit_class_def(self, o: ClassDef) -> None:
- # Don't recurse into classes, as we only keep track of top-level definitions.
- self.names.add(o.name)
- def visit_func_def(self, o: FuncDef) -> None:
- # Don't recurse, as we only keep track of top-level definitions.
- self.names.add(o.name)
- def find_referenced_names(file: MypyFile) -> set[str]:
- finder = ReferenceFinder()
- file.accept(finder)
- return finder.refs
- class ReferenceFinder(mypy.mixedtraverser.MixedTraverserVisitor):
- """Find all name references (both local and global)."""
- # TODO: Filter out local variable and class attribute references
- def __init__(self) -> None:
- # Short names of things defined at the top level.
- self.refs: set[str] = set()
- def visit_block(self, block: Block) -> None:
- if not block.is_unreachable:
- super().visit_block(block)
- def visit_name_expr(self, e: NameExpr) -> None:
- self.refs.add(e.name)
- def visit_instance(self, t: Instance) -> None:
- self.add_ref(t.type.fullname)
- super().visit_instance(t)
- def visit_unbound_type(self, t: UnboundType) -> None:
- if t.name:
- self.add_ref(t.name)
- def visit_tuple_type(self, t: TupleType) -> None:
- # Ignore fallback
- for item in t.items:
- item.accept(self)
- def visit_callable_type(self, t: CallableType) -> None:
- # Ignore fallback
- for arg in t.arg_types:
- arg.accept(self)
- t.ret_type.accept(self)
- def add_ref(self, fullname: str) -> None:
- self.refs.add(fullname.split(".")[-1])
- class StubGenerator(mypy.traverser.TraverserVisitor):
- """Generate stub text from a mypy AST."""
- def __init__(
- self,
- _all_: list[str] | None,
- include_private: bool = False,
- analyzed: bool = False,
- export_less: bool = False,
- ) -> None:
- # Best known value of __all__.
- self._all_ = _all_
- self._output: list[str] = []
- self._decorators: list[str] = []
- self._import_lines: list[str] = []
- # Current indent level (indent is hardcoded to 4 spaces).
- self._indent = ""
- # Stack of defined variables (per scope).
- self._vars: list[list[str]] = [[]]
- # What was generated previously in the stub file.
- self._state = EMPTY
- self._toplevel_names: list[str] = []
- self._include_private = include_private
- self.import_tracker = ImportTracker()
- # Was the tree semantically analysed before?
- self.analyzed = analyzed
- # Disable implicit exports of package-internal imports?
- self.export_less = export_less
- # Add imports that could be implicitly generated
- self.import_tracker.add_import_from("typing", [("NamedTuple", None)])
- # Names in __all__ are required
- for name in _all_ or ():
- if name not in IGNORED_DUNDERS:
- self.import_tracker.reexport(name)
- self.defined_names: set[str] = set()
- # Short names of methods defined in the body of the current class
- self.method_names: set[str] = set()
- def visit_mypy_file(self, o: MypyFile) -> None:
- self.module = o.fullname # Current module being processed
- self.path = o.path
- self.defined_names = find_defined_names(o)
- self.referenced_names = find_referenced_names(o)
- known_imports = {
- "_typeshed": ["Incomplete"],
- "typing": ["Any", "TypeVar", "NamedTuple"],
- "collections.abc": ["Generator"],
- "typing_extensions": ["TypedDict", "ParamSpec", "TypeVarTuple"],
- }
- for pkg, imports in known_imports.items():
- for t in imports:
- if t not in self.defined_names:
- alias = None
- else:
- alias = "_" + t
- self.import_tracker.add_import_from(pkg, [(t, alias)])
- super().visit_mypy_file(o)
- undefined_names = [name for name in self._all_ or [] if name not in self._toplevel_names]
- if undefined_names:
- if self._state != EMPTY:
- self.add("\n")
- self.add("# Names in __all__ with no definition:\n")
- for name in sorted(undefined_names):
- self.add(f"# {name}\n")
- def visit_overloaded_func_def(self, o: OverloadedFuncDef) -> None:
- """@property with setters and getters, @overload chain and some others."""
- overload_chain = False
- for item in o.items:
- if not isinstance(item, Decorator):
- continue
- if self.is_private_name(item.func.name, item.func.fullname):
- continue
- self.process_decorator(item)
- if not overload_chain:
- self.visit_func_def(item.func)
- if item.func.is_overload:
- overload_chain = True
- elif item.func.is_overload:
- self.visit_func_def(item.func)
- else:
- # skip the overload implementation and clear the decorator we just processed
- self.clear_decorators()
- def visit_func_def(self, o: FuncDef) -> None:
- if (
- self.is_private_name(o.name, o.fullname)
- or self.is_not_in_all(o.name)
- or (self.is_recorded_name(o.name) and not o.is_overload)
- ):
- self.clear_decorators()
- return
- if not self._indent and self._state not in (EMPTY, FUNC) and not o.is_awaitable_coroutine:
- self.add("\n")
- if not self.is_top_level():
- self_inits = find_self_initializers(o)
- for init, value in self_inits:
- if init in self.method_names:
- # Can't have both an attribute and a method/property with the same name.
- continue
- init_code = self.get_init(init, value)
- if init_code:
- self.add(init_code)
- # dump decorators, just before "def ..."
- for s in self._decorators:
- self.add(s)
- self.clear_decorators()
- self.add(f"{self._indent}{'async ' if o.is_coroutine else ''}def {o.name}(")
- self.record_name(o.name)
- args: list[str] = []
- for i, arg_ in enumerate(o.arguments):
- var = arg_.variable
- kind = arg_.kind
- name = var.name
- annotated_type = (
- o.unanalyzed_type.arg_types[i]
- if isinstance(o.unanalyzed_type, CallableType)
- else None
- )
- # I think the name check is incorrect: there are libraries which
- # name their 0th argument other than self/cls
- is_self_arg = i == 0 and name == "self"
- is_cls_arg = i == 0 and name == "cls"
- annotation = ""
- if annotated_type and not is_self_arg and not is_cls_arg:
- # Luckily, an argument explicitly annotated with "Any" has
- # type "UnboundType" and will not match.
- if not isinstance(get_proper_type(annotated_type), AnyType):
- annotation = f": {self.print_annotation(annotated_type)}"
- if kind.is_named() and not any(arg.startswith("*") for arg in args):
- args.append("*")
- if arg_.initializer:
- if not annotation:
- typename = self.get_str_type_of_node(arg_.initializer, True, False)
- if typename == "":
- annotation = "=..."
- else:
- annotation = f": {typename} = ..."
- else:
- annotation += " = ..."
- arg = name + annotation
- elif kind == ARG_STAR:
- arg = f"*{name}{annotation}"
- elif kind == ARG_STAR2:
- arg = f"**{name}{annotation}"
- else:
- arg = name + annotation
- args.append(arg)
- retname = None
- if o.name != "__init__" and isinstance(o.unanalyzed_type, CallableType):
- if isinstance(get_proper_type(o.unanalyzed_type.ret_type), AnyType):
- # Luckily, a return type explicitly annotated with "Any" has
- # type "UnboundType" and will enter the else branch.
- retname = None # implicit Any
- else:
- retname = self.print_annotation(o.unanalyzed_type.ret_type)
- elif o.abstract_status == IS_ABSTRACT or o.name in METHODS_WITH_RETURN_VALUE:
- # Always assume abstract methods return Any unless explicitly annotated. Also
- # some dunder methods should not have a None return type.
- retname = None # implicit Any
- elif o.name in KNOWN_MAGIC_METHODS_RETURN_TYPES:
- retname = KNOWN_MAGIC_METHODS_RETURN_TYPES[o.name]
- elif has_yield_expression(o) or has_yield_from_expression(o):
- self.add_typing_import("Generator")
- yield_name = "None"
- send_name = "None"
- return_name = "None"
- if has_yield_from_expression(o):
- self.add_typing_import("Incomplete")
- yield_name = send_name = self.typing_name("Incomplete")
- else:
- for expr, in_assignment in all_yield_expressions(o):
- if expr.expr is not None and not self.is_none_expr(expr.expr):
- self.add_typing_import("Incomplete")
- yield_name = self.typing_name("Incomplete")
- if in_assignment:
- self.add_typing_import("Incomplete")
- send_name = self.typing_name("Incomplete")
- if has_return_statement(o):
- self.add_typing_import("Incomplete")
- return_name = self.typing_name("Incomplete")
- generator_name = self.typing_name("Generator")
- retname = f"{generator_name}[{yield_name}, {send_name}, {return_name}]"
- elif not has_return_statement(o) and o.abstract_status == NOT_ABSTRACT:
- retname = "None"
- retfield = ""
- if retname is not None:
- retfield = " -> " + retname
- self.add(", ".join(args))
- self.add(f"){retfield}: ...\n")
- self._state = FUNC
- def is_none_expr(self, expr: Expression) -> bool:
- return isinstance(expr, NameExpr) and expr.name == "None"
- def visit_decorator(self, o: Decorator) -> None:
- if self.is_private_name(o.func.name, o.func.fullname):
- return
- self.process_decorator(o)
- self.visit_func_def(o.func)
- def process_decorator(self, o: Decorator) -> None:
- """Process a series of decorators.
- Only preserve certain special decorators such as @abstractmethod.
- """
- for decorator in o.original_decorators:
- if not isinstance(decorator, (NameExpr, MemberExpr)):
- continue
- qualname = get_qualified_name(decorator)
- fullname = self.get_fullname(decorator)
- if fullname in (
- "builtins.property",
- "builtins.staticmethod",
- "builtins.classmethod",
- "functools.cached_property",
- ):
- self.add_decorator(qualname, require_name=True)
- elif fullname in (
- "asyncio.coroutine",
- "asyncio.coroutines.coroutine",
- "types.coroutine",
- ):
- o.func.is_awaitable_coroutine = True
- self.add_decorator(qualname, require_name=True)
- elif fullname == "abc.abstractmethod":
- self.add_decorator(qualname, require_name=True)
- o.func.abstract_status = IS_ABSTRACT
- elif fullname in (
- "abc.abstractproperty",
- "abc.abstractstaticmethod",
- "abc.abstractclassmethod",
- ):
- abc_module = qualname.rpartition(".")[0]
- if not abc_module:
- self.import_tracker.add_import("abc")
- builtin_decorator_replacement = fullname[len("abc.abstract") :]
- self.add_decorator(builtin_decorator_replacement, require_name=False)
- self.add_decorator(f"{abc_module or 'abc'}.abstractmethod", require_name=True)
- o.func.abstract_status = IS_ABSTRACT
- elif fullname in OVERLOAD_NAMES:
- self.add_decorator(qualname, require_name=True)
- o.func.is_overload = True
- elif qualname.endswith(".setter"):
- self.add_decorator(qualname, require_name=False)
- def get_fullname(self, expr: Expression) -> str:
- """Return the full name resolving imports and import aliases."""
- if (
- self.analyzed
- and isinstance(expr, (NameExpr, MemberExpr))
- and expr.fullname
- and not (isinstance(expr.node, Var) and expr.node.is_suppressed_import)
- ):
- return expr.fullname
- name = get_qualified_name(expr)
- if "." not in name:
- real_module = self.import_tracker.module_for.get(name)
- real_short = self.import_tracker.reverse_alias.get(name, name)
- if real_module is None and real_short not in self.defined_names:
- real_module = "builtins" # not imported and not defined, must be a builtin
- else:
- name_module, real_short = name.split(".", 1)
- real_module = self.import_tracker.reverse_alias.get(name_module, name_module)
- resolved_name = real_short if real_module is None else f"{real_module}.{real_short}"
- return resolved_name
- def visit_class_def(self, o: ClassDef) -> None:
- self.method_names = find_method_names(o.defs.body)
- sep: int | None = None
- if not self._indent and self._state != EMPTY:
- sep = len(self._output)
- self.add("\n")
- self.add(f"{self._indent}class {o.name}")
- self.record_name(o.name)
- base_types = self.get_base_types(o)
- if base_types:
- for base in base_types:
- self.import_tracker.require_name(base)
- if isinstance(o.metaclass, (NameExpr, MemberExpr)):
- meta = o.metaclass.accept(AliasPrinter(self))
- base_types.append("metaclass=" + meta)
- elif self.analyzed and o.info.is_abstract and not o.info.is_protocol:
- base_types.append("metaclass=abc.ABCMeta")
- self.import_tracker.add_import("abc")
- self.import_tracker.require_name("abc")
- if base_types:
- self.add(f"({', '.join(base_types)})")
- self.add(":\n")
- n = len(self._output)
- self._indent += " "
- self._vars.append([])
- super().visit_class_def(o)
- self._indent = self._indent[:-4]
- self._vars.pop()
- self._vars[-1].append(o.name)
- if len(self._output) == n:
- if self._state == EMPTY_CLASS and sep is not None:
- self._output[sep] = ""
- self._output[-1] = self._output[-1][:-1] + " ...\n"
- self._state = EMPTY_CLASS
- else:
- self._state = CLASS
- self.method_names = set()
- def get_base_types(self, cdef: ClassDef) -> list[str]:
- """Get list of base classes for a class."""
- base_types: list[str] = []
- p = AliasPrinter(self)
- for base in cdef.base_type_exprs + cdef.removed_base_type_exprs:
- if isinstance(base, (NameExpr, MemberExpr)):
- if self.get_fullname(base) != "builtins.object":
- base_types.append(get_qualified_name(base))
- elif isinstance(base, IndexExpr):
- base_types.append(base.accept(p))
- elif isinstance(base, CallExpr):
- # namedtuple(typename, fields), NamedTuple(typename, fields) calls can
- # be used as a base class. The first argument is a string literal that
- # is usually the same as the class name.
- #
- # Note:
- # A call-based named tuple as a base class cannot be safely converted to
- # a class-based NamedTuple definition because class attributes defined
- # in the body of the class inheriting from the named tuple call are not
- # namedtuple fields at runtime.
- if self.is_namedtuple(base):
- nt_fields = self._get_namedtuple_fields(base)
- assert isinstance(base.args[0], StrExpr)
- typename = base.args[0].value
- if nt_fields is not None:
- fields_str = ", ".join(f"({f!r}, {t})" for f, t in nt_fields)
- namedtuple_name = self.typing_name("NamedTuple")
- base_types.append(f"{namedtuple_name}({typename!r}, [{fields_str}])")
- self.add_typing_import("NamedTuple")
- else:
- # Invalid namedtuple() call, cannot determine fields
- base_types.append(self.typing_name("Incomplete"))
- elif self.is_typed_namedtuple(base):
- base_types.append(base.accept(p))
- else:
- # At this point, we don't know what the base class is, so we
- # just use Incomplete as the base class.
- base_types.append(self.typing_name("Incomplete"))
- self.add_typing_import("Incomplete")
- for name, value in cdef.keywords.items():
- if name == "metaclass":
- continue # handled separately
- base_types.append(f"{name}={value.accept(p)}")
- return base_types
- def visit_block(self, o: Block) -> None:
- # Unreachable statements may be partially uninitialized and that may
- # cause trouble.
- if not o.is_unreachable:
- super().visit_block(o)
- def visit_assignment_stmt(self, o: AssignmentStmt) -> None:
- foundl = []
- for lvalue in o.lvalues:
- if isinstance(lvalue, NameExpr) and isinstance(o.rvalue, CallExpr):
- if self.is_namedtuple(o.rvalue) or self.is_typed_namedtuple(o.rvalue):
- self.process_namedtuple(lvalue, o.rvalue)
- foundl.append(False) # state is updated in process_namedtuple
- continue
- if self.is_typeddict(o.rvalue):
- self.process_typeddict(lvalue, o.rvalue)
- foundl.append(False) # state is updated in process_typeddict
- continue
- if (
- isinstance(lvalue, NameExpr)
- and not self.is_private_name(lvalue.name)
- # it is never an alias with explicit annotation
- and not o.unanalyzed_type
- and self.is_alias_expression(o.rvalue)
- ):
- self.process_typealias(lvalue, o.rvalue)
- continue
- if isinstance(lvalue, (TupleExpr, ListExpr)):
- items = lvalue.items
- if isinstance(o.unanalyzed_type, TupleType): # type: ignore[misc]
- annotations: Iterable[Type | None] = o.unanalyzed_type.items
- else:
- annotations = [None] * len(items)
- else:
- items = [lvalue]
- annotations = [o.unanalyzed_type]
- sep = False
- found = False
- for item, annotation in zip(items, annotations):
- if isinstance(item, NameExpr):
- init = self.get_init(item.name, o.rvalue, annotation)
- if init:
- found = True
- if not sep and not self._indent and self._state not in (EMPTY, VAR):
- init = "\n" + init
- sep = True
- self.add(init)
- self.record_name(item.name)
- foundl.append(found)
- if all(foundl):
- self._state = VAR
- def is_namedtuple(self, expr: CallExpr) -> bool:
- return self.get_fullname(expr.callee) == "collections.namedtuple"
- def is_typed_namedtuple(self, expr: CallExpr) -> bool:
- return self.get_fullname(expr.callee) in TYPED_NAMEDTUPLE_NAMES
- def _get_namedtuple_fields(self, call: CallExpr) -> list[tuple[str, str]] | None:
- if self.is_namedtuple(call):
- fields_arg = call.args[1]
- if isinstance(fields_arg, StrExpr):
- field_names = fields_arg.value.replace(",", " ").split()
- elif isinstance(fields_arg, (ListExpr, TupleExpr)):
- field_names = []
- for field in fields_arg.items:
- if not isinstance(field, StrExpr):
- return None
- field_names.append(field.value)
- else:
- return None # Invalid namedtuple fields type
- if field_names:
- self.add_typing_import("Incomplete")
- incomplete = self.typing_name("Incomplete")
- return [(field_name, incomplete) for field_name in field_names]
- elif self.is_typed_namedtuple(call):
- fields_arg = call.args[1]
- if not isinstance(fields_arg, (ListExpr, TupleExpr)):
- return None
- fields: list[tuple[str, str]] = []
- p = AliasPrinter(self)
- for field in fields_arg.items:
- if not (isinstance(field, TupleExpr) and len(field.items) == 2):
- return None
- field_name, field_type = field.items
- if not isinstance(field_name, StrExpr):
- return None
- fields.append((field_name.value, field_type.accept(p)))
- return fields
- else:
- return None # Not a named tuple call
- def process_namedtuple(self, lvalue: NameExpr, rvalue: CallExpr) -> None:
- if self._state == CLASS:
- self.add("\n")
- if not isinstance(rvalue.args[0], StrExpr):
- self.annotate_as_incomplete(lvalue)
- return
- fields = self._get_namedtuple_fields(rvalue)
- if fields is None:
- self.annotate_as_incomplete(lvalue)
- return
- self.add_typing_import("NamedTuple")
- bases = self.typing_name("NamedTuple")
- # TODO: Add support for generic NamedTuples. Requires `Generic` as base class.
- class_def = f"{self._indent}class {lvalue.name}({bases}):"
- if len(fields) == 0:
- self.add(f"{class_def} ...\n")
- self._state = EMPTY_CLASS
- else:
- if self._state not in (EMPTY, CLASS):
- self.add("\n")
- self.add(f"{class_def}\n")
- for f_name, f_type in fields:
- self.add(f"{self._indent} {f_name}: {f_type}\n")
- self._state = CLASS
- def is_typeddict(self, expr: CallExpr) -> bool:
- return self.get_fullname(expr.callee) in TPDICT_NAMES
- def process_typeddict(self, lvalue: NameExpr, rvalue: CallExpr) -> None:
- if self._state == CLASS:
- self.add("\n")
- if not isinstance(rvalue.args[0], StrExpr):
- self.annotate_as_incomplete(lvalue)
- return
- items: list[tuple[str, Expression]] = []
- total: Expression | None = None
- if len(rvalue.args) > 1 and rvalue.arg_kinds[1] == ARG_POS:
- if not isinstance(rvalue.args[1], DictExpr):
- self.annotate_as_incomplete(lvalue)
- return
- for attr_name, attr_type in rvalue.args[1].items:
- if not isinstance(attr_name, StrExpr):
- self.annotate_as_incomplete(lvalue)
- return
- items.append((attr_name.value, attr_type))
- if len(rvalue.args) > 2:
- if rvalue.arg_kinds[2] != ARG_NAMED or rvalue.arg_names[2] != "total":
- self.annotate_as_incomplete(lvalue)
- return
- total = rvalue.args[2]
- else:
- for arg_name, arg in zip(rvalue.arg_names[1:], rvalue.args[1:]):
- if not isinstance(arg_name, str):
- self.annotate_as_incomplete(lvalue)
- return
- if arg_name == "total":
- total = arg
- else:
- items.append((arg_name, arg))
- self.add_typing_import("TypedDict")
- p = AliasPrinter(self)
- if any(not key.isidentifier() or keyword.iskeyword(key) for key, _ in items):
- # Keep the call syntax if there are non-identifier or reserved keyword keys.
- self.add(f"{self._indent}{lvalue.name} = {rvalue.accept(p)}\n")
- self._state = VAR
- else:
- bases = self.typing_name("TypedDict")
- # TODO: Add support for generic TypedDicts. Requires `Generic` as base class.
- if total is not None:
- bases += f", total={total.accept(p)}"
- class_def = f"{self._indent}class {lvalue.name}({bases}):"
- if len(items) == 0:
- self.add(f"{class_def} ...\n")
- self._state = EMPTY_CLASS
- else:
- if self._state not in (EMPTY, CLASS):
- self.add("\n")
- self.add(f"{class_def}\n")
- for key, key_type in items:
- self.add(f"{self._indent} {key}: {key_type.accept(p)}\n")
- self._state = CLASS
- def annotate_as_incomplete(self, lvalue: NameExpr) -> None:
- self.add_typing_import("Incomplete")
- self.add(f"{self._indent}{lvalue.name}: {self.typing_name('Incomplete')}\n")
- self._state = VAR
- def is_alias_expression(self, expr: Expression, top_level: bool = True) -> bool:
- """Return True for things that look like target for an alias.
- Used to know if assignments look like type aliases, function alias,
- or module alias.
- """
- # Assignment of TypeVar(...) and other typevar-likes are passed through
- if isinstance(expr, CallExpr) and self.get_fullname(expr.callee) in (
- "typing.TypeVar",
- "typing_extensions.TypeVar",
- "typing.ParamSpec",
- "typing_extensions.ParamSpec",
- "typing.TypeVarTuple",
- "typing_extensions.TypeVarTuple",
- ):
- return True
- elif isinstance(expr, EllipsisExpr):
- return not top_level
- elif isinstance(expr, NameExpr):
- if expr.name in ("True", "False"):
- return False
- elif expr.name == "None":
- return not top_level
- else:
- return not self.is_private_name(expr.name)
- elif isinstance(expr, MemberExpr) and self.analyzed:
- # Also add function and module aliases.
- return (
- top_level
- and isinstance(expr.node, (FuncDef, Decorator, MypyFile))
- or isinstance(expr.node, TypeInfo)
- ) and not self.is_private_member(expr.node.fullname)
- elif (
- isinstance(expr, IndexExpr)
- and isinstance(expr.base, NameExpr)
- and not self.is_private_name(expr.base.name)
- ):
- if isinstance(expr.index, TupleExpr):
- indices = expr.index.items
- else:
- indices = [expr.index]
- if expr.base.name == "Callable" and len(indices) == 2:
- args, ret = indices
- if isinstance(args, EllipsisExpr):
- indices = [ret]
- elif isinstance(args, ListExpr):
- indices = args.items + [ret]
- else:
- return False
- return all(self.is_alias_expression(i, top_level=False) for i in indices)
- else:
- return False
- def process_typealias(self, lvalue: NameExpr, rvalue: Expression) -> None:
- p = AliasPrinter(self)
- self.add(f"{self._indent}{lvalue.name} = {rvalue.accept(p)}\n")
- self.record_name(lvalue.name)
- self._vars[-1].append(lvalue.name)
- def visit_if_stmt(self, o: IfStmt) -> None:
- # Ignore if __name__ == '__main__'.
- expr = o.expr[0]
- if (
- isinstance(expr, ComparisonExpr)
- and isinstance(expr.operands[0], NameExpr)
- and isinstance(expr.operands[1], StrExpr)
- and expr.operands[0].name == "__name__"
- and "__main__" in expr.operands[1].value
- ):
- return
- super().visit_if_stmt(o)
- def visit_import_all(self, o: ImportAll) -> None:
- self.add_import_line(f"from {'.' * o.relative}{o.id} import *\n")
- def visit_import_from(self, o: ImportFrom) -> None:
- exported_names: set[str] = set()
- import_names = []
- module, relative = translate_module_name(o.id, o.relative)
- if self.module:
- full_module, ok = mypy.util.correct_relative_import(
- self.module, relative, module, self.path.endswith(".__init__.py")
- )
- if not ok:
- full_module = module
- else:
- full_module = module
- if module == "__future__":
- return # Not preserved
- for name, as_name in o.names:
- if name == "six":
- # Vendored six -- translate into plain 'import six'.
- self.visit_import(Import([("six", None)]))
- continue
- exported = False
- if as_name is None and self.module and (self.module + "." + name) in EXTRA_EXPORTED:
- # Special case certain names that should be exported, against our general rules.
- exported = True
- is_private = self.is_private_name(name, full_module + "." + name)
- if (
- as_name is None
- and name not in self.referenced_names
- and (not self._all_ or name in IGNORED_DUNDERS)
- and not is_private
- and module not in ("abc", "asyncio") + TYPING_MODULE_NAMES
- ):
- # An imported name that is never referenced in the module is assumed to be
- # exported, unless there is an explicit __all__. Note that we need to special
- # case 'abc' since some references are deleted during semantic analysis.
- exported = True
- top_level = full_module.split(".")[0]
- if (
- as_name is None
- and not self.export_less
- and (not self._all_ or name in IGNORED_DUNDERS)
- and self.module
- and not is_private
- and top_level in (self.module.split(".")[0], "_" + self.module.split(".")[0])
- ):
- # Export imports from the same package, since we can't reliably tell whether they
- # are part of the public API.
- exported = True
- if exported:
- self.import_tracker.reexport(name)
- as_name = name
- import_names.append((name, as_name))
- self.import_tracker.add_import_from("." * relative + module, import_names)
- self._vars[-1].extend(alias or name for name, alias in import_names)
- for name, alias in import_names:
- self.record_name(alias or name)
- if self._all_:
- # Include "import from"s that import names defined in __all__.
- names = [
- name
- for name, alias in o.names
- if name in self._all_ and alias is None and name not in IGNORED_DUNDERS
- ]
- exported_names.update(names)
- def visit_import(self, o: Import) -> None:
- for id, as_id in o.ids:
- self.import_tracker.add_import(id, as_id)
- if as_id is None:
- target_name = id.split(".")[0]
- else:
- target_name = as_id
- self._vars[-1].append(target_name)
- self.record_name(target_name)
- def get_init(
- self, lvalue: str, rvalue: Expression, annotation: Type | None = None
- ) -> str | None:
- """Return initializer for a variable.
- Return None if we've generated one already or if the variable is internal.
- """
- if lvalue in self._vars[-1]:
- # We've generated an initializer already for this variable.
- return None
- # TODO: Only do this at module top level.
- if self.is_private_name(lvalue) or self.is_not_in_all(lvalue):
- return None
- self._vars[-1].append(lvalue)
- if annotation is not None:
- typename = self.print_annotation(annotation)
- if (
- isinstance(annotation, UnboundType)
- and not annotation.args
- and annotation.name == "Final"
- and self.import_tracker.module_for.get("Final") in TYPING_MODULE_NAMES
- ):
- # Final without type argument is invalid in stubs.
- final_arg = self.get_str_type_of_node(rvalue)
- typename += f"[{final_arg}]"
- else:
- typename = self.get_str_type_of_node(rvalue)
- return f"{self._indent}{lvalue}: {typename}\n"
- def add(self, string: str) -> None:
- """Add text to generated stub."""
- self._output.append(string)
- def add_decorator(self, name: str, require_name: bool = False) -> None:
- if require_name:
- self.import_tracker.require_name(name)
- if not self._indent and self._state not in (EMPTY, FUNC):
- self._decorators.append("\n")
- self._decorators.append(f"{self._indent}@{name}\n")
- def clear_decorators(self) -> None:
- self._decorators.clear()
- def typing_name(self, name: str) -> str:
- if name in self.defined_names:
- # Avoid name clash between name from typing and a name defined in stub.
- return "_" + name
- else:
- return name
- def add_typing_import(self, name: str) -> None:
- """Add a name to be imported for typing, unless it's imported already.
- The import will be internal to the stub.
- """
- name = self.typing_name(name)
- self.import_tracker.require_name(name)
- def add_import_line(self, line: str) -> None:
- """Add a line of text to the import section, unless it's already there."""
- if line not in self._import_lines:
- self._import_lines.append(line)
- def output(self) -> str:
- """Return the text for the stub."""
- imports = ""
- if self._import_lines:
- imports += "".join(self._import_lines)
- imports += "".join(self.import_tracker.import_lines())
- if imports and self._output:
- imports += "\n"
- return imports + "".join(self._output)
- def is_not_in_all(self, name: str) -> bool:
- if self.is_private_name(name):
- return False
- if self._all_:
- return self.is_top_level() and name not in self._all_
- return False
- def is_private_name(self, name: str, fullname: str | None = None) -> bool:
- if self._include_private:
- return False
- if fullname in EXTRA_EXPORTED:
- return False
- return name.startswith("_") and (not name.endswith("__") or name in IGNORED_DUNDERS)
- def is_private_member(self, fullname: str) -> bool:
- parts = fullname.split(".")
- return any(self.is_private_name(part) for part in parts)
- def get_str_type_of_node(
- self, rvalue: Expression, can_infer_optional: bool = False, can_be_any: bool = True
- ) -> str:
- rvalue = self.maybe_unwrap_unary_expr(rvalue)
- if isinstance(rvalue, IntExpr):
- return "int"
- if isinstance(rvalue, StrExpr):
- return "str"
- if isinstance(rvalue, BytesExpr):
- return "bytes"
- if isinstance(rvalue, FloatExpr):
- return "float"
- if isinstance(rvalue, ComplexExpr): # 1j
- return "complex"
- if isinstance(rvalue, OpExpr) and rvalue.op in ("-", "+"): # -1j + 1
- if isinstance(self.maybe_unwrap_unary_expr(rvalue.left), ComplexExpr) or isinstance(
- self.maybe_unwrap_unary_expr(rvalue.right), ComplexExpr
- ):
- return "complex"
- if isinstance(rvalue, NameExpr) and rvalue.name in ("True", "False"):
- return "bool"
- if can_infer_optional and isinstance(rvalue, NameExpr) and rvalue.name == "None":
- self.add_typing_import("Incomplete")
- return f"{self.typing_name('Incomplete')} | None"
- if can_be_any:
- self.add_typing_import("Incomplete")
- return self.typing_name("Incomplete")
- else:
- return ""
- def maybe_unwrap_unary_expr(self, expr: Expression) -> Expression:
- """Unwrap (possibly nested) unary expressions.
- But, some unary expressions can change the type of expression.
- While we want to preserve it. For example, `~True` is `int`.
- So, we only allow a subset of unary expressions to be unwrapped.
- """
- if not isinstance(expr, UnaryExpr):
- return expr
- # First, try to unwrap `[+-]+ (int|float|complex)` expr:
- math_ops = ("+", "-")
- if expr.op in math_ops:
- while isinstance(expr, UnaryExpr):
- if expr.op not in math_ops or not isinstance(
- expr.expr, (IntExpr, FloatExpr, ComplexExpr, UnaryExpr)
- ):
- break
- expr = expr.expr
- return expr
- # Next, try `not bool` expr:
- if expr.op == "not":
- while isinstance(expr, UnaryExpr):
- if expr.op != "not" or not isinstance(expr.expr, (NameExpr, UnaryExpr)):
- break
- if isinstance(expr.expr, NameExpr) and expr.expr.name not in ("True", "False"):
- break
- expr = expr.expr
- return expr
- # This is some other unary expr, we cannot do anything with it (yet?).
- return expr
- def print_annotation(self, t: Type) -> str:
- printer = AnnotationPrinter(self)
- return t.accept(printer)
- def is_top_level(self) -> bool:
- """Are we processing the top level of a file?"""
- return self._indent == ""
- def record_name(self, name: str) -> None:
- """Mark a name as defined.
- This only does anything if at the top level of a module.
- """
- if self.is_top_level():
- self._toplevel_names.append(name)
- def is_recorded_name(self, name: str) -> bool:
- """Has this name been recorded previously?"""
- return self.is_top_level() and name in self._toplevel_names
- def find_method_names(defs: list[Statement]) -> set[str]:
- # TODO: Traverse into nested definitions
- result = set()
- for defn in defs:
- if isinstance(defn, FuncDef):
- result.add(defn.name)
- elif isinstance(defn, Decorator):
- result.add(defn.func.name)
- elif isinstance(defn, OverloadedFuncDef):
- for item in defn.items:
- result.update(find_method_names([item]))
- return result
- class SelfTraverser(mypy.traverser.TraverserVisitor):
- def __init__(self) -> None:
- self.results: list[tuple[str, Expression]] = []
- def visit_assignment_stmt(self, o: AssignmentStmt) -> None:
- lvalue = o.lvalues[0]
- if (
- isinstance(lvalue, MemberExpr)
- and isinstance(lvalue.expr, NameExpr)
- and lvalue.expr.name == "self"
- ):
- self.results.append((lvalue.name, o.rvalue))
- def find_self_initializers(fdef: FuncBase) -> list[tuple[str, Expression]]:
- """Find attribute initializers in a method.
- Return a list of pairs (attribute name, r.h.s. expression).
- """
- traverser = SelfTraverser()
- fdef.accept(traverser)
- return traverser.results
- def get_qualified_name(o: Expression) -> str:
- if isinstance(o, NameExpr):
- return o.name
- elif isinstance(o, MemberExpr):
- return f"{get_qualified_name(o.expr)}.{o.name}"
- else:
- return ERROR_MARKER
- def remove_blacklisted_modules(modules: list[StubSource]) -> list[StubSource]:
- return [
- module for module in modules if module.path is None or not is_blacklisted_path(module.path)
- ]
- def is_blacklisted_path(path: str) -> bool:
- return any(substr in (normalize_path_separators(path) + "\n") for substr in BLACKLIST)
- def normalize_path_separators(path: str) -> str:
- if sys.platform == "win32":
- return path.replace("\\", "/")
- return path
- def collect_build_targets(
- options: Options, mypy_opts: MypyOptions
- ) -> tuple[list[StubSource], list[StubSource]]:
- """Collect files for which we need to generate stubs.
- Return list of Python modules and C modules.
- """
- if options.packages or options.modules:
- if options.no_import:
- py_modules = find_module_paths_using_search(
- options.modules, options.packages, options.search_path, options.pyversion
- )
- c_modules: list[StubSource] = []
- else:
- # Using imports is the default, since we can also find C modules.
- py_modules, c_modules = find_module_paths_using_imports(
- options.modules, options.packages, options.verbose, options.quiet
- )
- else:
- # Use mypy native source collection for files and directories.
- try:
- source_list = create_source_list(options.files, mypy_opts)
- except InvalidSourceList as e:
- raise SystemExit(str(e)) from e
- py_modules = [StubSource(m.module, m.path) for m in source_list]
- c_modules = []
- py_modules = remove_blacklisted_modules(py_modules)
- return py_modules, c_modules
- def find_module_paths_using_imports(
- modules: list[str], packages: list[str], verbose: bool, quiet: bool
- ) -> tuple[list[StubSource], list[StubSource]]:
- """Find path and runtime value of __all__ (if possible) for modules and packages.
- This function uses runtime Python imports to get the information.
- """
- with ModuleInspect() as inspect:
- py_modules: list[StubSource] = []
- c_modules: list[StubSource] = []
- found = list(walk_packages(inspect, packages, verbose))
- modules = modules + found
- modules = [
- mod for mod in modules if not is_non_library_module(mod)
- ] # We don't want to run any tests or scripts
- for mod in modules:
- try:
- result = find_module_path_and_all_py3(inspect, mod, verbose)
- except CantImport as e:
- tb = traceback.format_exc()
- if verbose:
- sys.stdout.write(tb)
- if not quiet:
- report_missing(mod, e.message, tb)
- continue
- if not result:
- c_modules.append(StubSource(mod))
- else:
- path, runtime_all = result
- py_modules.append(StubSource(mod, path, runtime_all))
- return py_modules, c_modules
- def is_non_library_module(module: str) -> bool:
- """Does module look like a test module or a script?"""
- if module.endswith(
- (
- ".tests",
- ".test",
- ".testing",
- "_tests",
- "_test_suite",
- "test_util",
- "test_utils",
- "test_base",
- ".__main__",
- ".conftest", # Used by pytest
- ".setup", # Typically an install script
- )
- ):
- return True
- if module.split(".")[-1].startswith("test_"):
- return True
- if (
- ".tests." in module
- or ".test." in module
- or ".testing." in module
- or ".SelfTest." in module
- ):
- return True
- return False
- def translate_module_name(module: str, relative: int) -> tuple[str, int]:
- for pkg in VENDOR_PACKAGES:
- for alt in "six.moves", "six":
- substr = f"{pkg}.{alt}"
- if module.endswith("." + substr) or (module == substr and relative):
- return alt, 0
- if "." + substr + "." in module:
- return alt + "." + module.partition("." + substr + ".")[2], 0
- return module, relative
- def find_module_paths_using_search(
- modules: list[str], packages: list[str], search_path: list[str], pyversion: tuple[int, int]
- ) -> list[StubSource]:
- """Find sources for modules and packages requested.
- This function just looks for source files at the file system level.
- This is used if user passes --no-import, and will not find C modules.
- Exit if some of the modules or packages can't be found.
- """
- result: list[StubSource] = []
- typeshed_path = default_lib_path(mypy.build.default_data_dir(), pyversion, None)
- search_paths = SearchPaths((".",) + tuple(search_path), (), (), tuple(typeshed_path))
- cache = FindModuleCache(search_paths, fscache=None, options=None)
- for module in modules:
- m_result = cache.find_module(module)
- if isinstance(m_result, ModuleNotFoundReason):
- fail_missing(module, m_result)
- module_path = None
- else:
- module_path = m_result
- result.append(StubSource(module, module_path))
- for package in packages:
- p_result = cache.find_modules_recursive(package)
- if p_result:
- fail_missing(package, ModuleNotFoundReason.NOT_FOUND)
- sources = [StubSource(m.module, m.path) for m in p_result]
- result.extend(sources)
- result = [m for m in result if not is_non_library_module(m.module)]
- return result
- def mypy_options(stubgen_options: Options) -> MypyOptions:
- """Generate mypy options using the flag passed by user."""
- options = MypyOptions()
- options.follow_imports = "skip"
- options.incremental = False
- options.ignore_errors = True
- options.semantic_analysis_only = True
- options.python_version = stubgen_options.pyversion
- options.show_traceback = True
- options.transform_source = remove_misplaced_type_comments
- options.preserve_asts = True
- # Override cache_dir if provided in the environment
- environ_cache_dir = os.getenv("MYPY_CACHE_DIR", "")
- if environ_cache_dir.strip():
- options.cache_dir = environ_cache_dir
- options.cache_dir = os.path.expanduser(options.cache_dir)
- return options
- def parse_source_file(mod: StubSource, mypy_options: MypyOptions) -> None:
- """Parse a source file.
- On success, store AST in the corresponding attribute of the stub source.
- If there are syntax errors, print them and exit.
- """
- assert mod.path is not None, "Not found module was not skipped"
- with open(mod.path, "rb") as f:
- data = f.read()
- source = mypy.util.decode_python_encoding(data)
- errors = Errors(mypy_options)
- mod.ast = mypy.parse.parse(
- source, fnam=mod.path, module=mod.module, errors=errors, options=mypy_options
- )
- mod.ast._fullname = mod.module
- if errors.is_blockers():
- # Syntax error!
- for m in errors.new_messages():
- sys.stderr.write(f"{m}\n")
- sys.exit(1)
- def generate_asts_for_modules(
- py_modules: list[StubSource], parse_only: bool, mypy_options: MypyOptions, verbose: bool
- ) -> None:
- """Use mypy to parse (and optionally analyze) source files."""
- if not py_modules:
- return # Nothing to do here, but there may be C modules
- if verbose:
- print(f"Processing {len(py_modules)} files...")
- if parse_only:
- for mod in py_modules:
- parse_source_file(mod, mypy_options)
- return
- # Perform full semantic analysis of the source set.
- try:
- res = build([module.source for module in py_modules], mypy_options)
- except CompileError as e:
- raise SystemExit(f"Critical error during semantic analysis: {e}") from e
- for mod in py_modules:
- mod.ast = res.graph[mod.module].tree
- # Use statically inferred __all__ if there is no runtime one.
- if mod.runtime_all is None:
- mod.runtime_all = res.manager.semantic_analyzer.export_map[mod.module]
- def generate_stub_from_ast(
- mod: StubSource,
- target: str,
- parse_only: bool = False,
- include_private: bool = False,
- export_less: bool = False,
- ) -> None:
- """Use analysed (or just parsed) AST to generate type stub for single file.
- If directory for target doesn't exist it will created. Existing stub
- will be overwritten.
- """
- gen = StubGenerator(
- mod.runtime_all,
- include_private=include_private,
- analyzed=not parse_only,
- export_less=export_less,
- )
- assert mod.ast is not None, "This function must be used only with analyzed modules"
- mod.ast.accept(gen)
- # Write output to file.
- subdir = os.path.dirname(target)
- if subdir and not os.path.isdir(subdir):
- os.makedirs(subdir)
- with open(target, "w") as file:
- file.write("".join(gen.output()))
- def get_sig_generators(options: Options) -> list[SignatureGenerator]:
- sig_generators: list[SignatureGenerator] = [
- DocstringSignatureGenerator(),
- FallbackSignatureGenerator(),
- ]
- if options.doc_dir:
- # Collect info from docs (if given). Always check these first.
- sigs, class_sigs = collect_docs_signatures(options.doc_dir)
- sig_generators.insert(0, ExternalSignatureGenerator(sigs, class_sigs))
- return sig_generators
- def collect_docs_signatures(doc_dir: str) -> tuple[dict[str, str], dict[str, str]]:
- """Gather all function and class signatures in the docs.
- Return a tuple (function signatures, class signatures).
- Currently only used for C modules.
- """
- all_sigs: list[Sig] = []
- all_class_sigs: list[Sig] = []
- for path in glob.glob(f"{doc_dir}/*.rst"):
- with open(path) as f:
- loc_sigs, loc_class_sigs = parse_all_signatures(f.readlines())
- all_sigs += loc_sigs
- all_class_sigs += loc_class_sigs
- sigs = dict(find_unique_signatures(all_sigs))
- class_sigs = dict(find_unique_signatures(all_class_sigs))
- return sigs, class_sigs
- def generate_stubs(options: Options) -> None:
- """Main entry point for the program."""
- mypy_opts = mypy_options(options)
- py_modules, c_modules = collect_build_targets(options, mypy_opts)
- sig_generators = get_sig_generators(options)
- # Use parsed sources to generate stubs for Python modules.
- generate_asts_for_modules(py_modules, options.parse_only, mypy_opts, options.verbose)
- files = []
- for mod in py_modules:
- assert mod.path is not None, "Not found module was not skipped"
- target = mod.module.replace(".", "/")
- if os.path.basename(mod.path) == "__init__.py":
- target += "/__init__.pyi"
- else:
- target += ".pyi"
- target = os.path.join(options.output_dir, target)
- files.append(target)
- with generate_guarded(mod.module, target, options.ignore_errors, options.verbose):
- generate_stub_from_ast(
- mod, target, options.parse_only, options.include_private, options.export_less
- )
- # Separately analyse C modules using different logic.
- all_modules = sorted(m.module for m in (py_modules + c_modules))
- for mod in c_modules:
- if any(py_mod.module.startswith(mod.module + ".") for py_mod in py_modules + c_modules):
- target = mod.module.replace(".", "/") + "/__init__.pyi"
- else:
- target = mod.module.replace(".", "/") + ".pyi"
- target = os.path.join(options.output_dir, target)
- files.append(target)
- with generate_guarded(mod.module, target, options.ignore_errors, options.verbose):
- generate_stub_for_c_module(
- mod.module, target, known_modules=all_modules, sig_generators=sig_generators
- )
- num_modules = len(py_modules) + len(c_modules)
- if not options.quiet and num_modules > 0:
- print("Processed %d modules" % num_modules)
- if len(files) == 1:
- print(f"Generated {files[0]}")
- else:
- print(f"Generated files under {common_dir_prefix(files)}" + os.sep)
- HEADER = """%(prog)s [-h] [more options, see -h]
- [-m MODULE] [-p PACKAGE] [files ...]"""
- DESCRIPTION = """
- Generate draft stubs for modules.
- Stubs are generated in directory ./out, to avoid overriding files with
- manual changes. This directory is assumed to exist.
- """
- def parse_options(args: list[str]) -> Options:
- parser = argparse.ArgumentParser(prog="stubgen", usage=HEADER, description=DESCRIPTION)
- parser.add_argument(
- "--ignore-errors",
- action="store_true",
- help="ignore errors when trying to generate stubs for modules",
- )
- parser.add_argument(
- "--no-import",
- action="store_true",
- help="don't import the modules, just parse and analyze them "
- "(doesn't work with C extension modules and might not "
- "respect __all__)",
- )
- parser.add_argument(
- "--parse-only",
- action="store_true",
- help="don't perform semantic analysis of sources, just parse them "
- "(only applies to Python modules, might affect quality of stubs)",
- )
- parser.add_argument(
- "--include-private",
- action="store_true",
- help="generate stubs for objects and members considered private "
- "(single leading underscore and no trailing underscores)",
- )
- parser.add_argument(
- "--export-less",
- action="store_true",
- help="don't implicitly export all names imported from other modules in the same package",
- )
- parser.add_argument("-v", "--verbose", action="store_true", help="show more verbose messages")
- parser.add_argument("-q", "--quiet", action="store_true", help="show fewer messages")
- parser.add_argument(
- "--doc-dir",
- metavar="PATH",
- default="",
- help="use .rst documentation in PATH (this may result in "
- "better stubs in some cases; consider setting this to "
- "DIR/Python-X.Y.Z/Doc/library)",
- )
- parser.add_argument(
- "--search-path",
- metavar="PATH",
- default="",
- help="specify module search directories, separated by ':' "
- "(currently only used if --no-import is given)",
- )
- parser.add_argument(
- "-o",
- "--output",
- metavar="PATH",
- dest="output_dir",
- default="out",
- help="change the output directory [default: %(default)s]",
- )
- parser.add_argument(
- "-m",
- "--module",
- action="append",
- metavar="MODULE",
- dest="modules",
- default=[],
- help="generate stub for module; can repeat for more modules",
- )
- parser.add_argument(
- "-p",
- "--package",
- action="append",
- metavar="PACKAGE",
- dest="packages",
- default=[],
- help="generate stubs for package recursively; can be repeated",
- )
- parser.add_argument(
- metavar="files",
- nargs="*",
- dest="files",
- help="generate stubs for given files or directories",
- )
- ns = parser.parse_args(args)
- pyversion = sys.version_info[:2]
- ns.interpreter = sys.executable
- if ns.modules + ns.packages and ns.files:
- parser.error("May only specify one of: modules/packages or files.")
- if ns.quiet and ns.verbose:
- parser.error("Cannot specify both quiet and verbose messages")
- # Create the output folder if it doesn't already exist.
- if not os.path.exists(ns.output_dir):
- os.makedirs(ns.output_dir)
- return Options(
- pyversion=pyversion,
- no_import=ns.no_import,
- doc_dir=ns.doc_dir,
- search_path=ns.search_path.split(":"),
- interpreter=ns.interpreter,
- ignore_errors=ns.ignore_errors,
- parse_only=ns.parse_only,
- include_private=ns.include_private,
- output_dir=ns.output_dir,
- modules=ns.modules,
- packages=ns.packages,
- files=ns.files,
- verbose=ns.verbose,
- quiet=ns.quiet,
- export_less=ns.export_less,
- )
- def main(args: list[str] | None = None) -> None:
- mypy.util.check_python_version("stubgen")
- # Make sure that the current directory is in sys.path so that
- # stubgen can be run on packages in the current directory.
- if not ("" in sys.path or "." in sys.path):
- sys.path.insert(0, "")
- options = parse_options(sys.argv[1:] if args is None else args)
- generate_stubs(options)
- if __name__ == "__main__":
- main()
|