Advertisement
Lulz-Tigre

Stubgen

May 30th, 2017
157
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 30.41 KB | None | 0 0
  1. """Generator of dynamically typed draft stubs for arbitrary modules.
  2. Basic usage:
  3.  $ mkdir out
  4.  $ stubgen urllib.parse
  5.  => Generate out/urllib/parse.pyi.
  6. For Python 2 mode, use --py2:
  7.  $ stubgen --py2 textwrap
  8. For C modules, you can get more precise function signatures by parsing .rst (Sphinx)
  9. documentation for extra information. For this, use the --docpath option:
  10.  $ scripts/stubgen --docpath <DIR>/Python-3.4.2/Doc/library curses
  11.  => Generate out/curses.py.
  12. Use "stubgen -h" for more help.
  13. Note: You should verify the generated stubs manually.
  14. TODO:
  15. - support stubs for C modules in Python 2 mode
  16. - support non-default Python interpreters in Python 3 mode
  17. - if using --no-import, look for __all__ in the AST
  18. - infer some return types, such as no return statement with value -> None
  19. - detect 'if PY2 / is_py2' etc. and either preserve those or only include Python 2 or 3 case
  20. - maybe export more imported names if there is no __all__ (this affects ssl.SSLError, for example)
  21.   - a quick and dirty heuristic would be to turn this on if a module has something like
  22.     'from x import y as _y'
  23. - we don't seem to always detect properties ('closed' in 'io', for example)
  24. """
  25.  
  26. import glob
  27. import importlib
  28. import json
  29. import os.path
  30. import pkgutil
  31. import subprocess
  32. import sys
  33. import textwrap
  34. import traceback
  35.  
  36. from typing import (
  37.     Any, List, Dict, Tuple, Iterable, Iterator, Optional, NamedTuple, Set, Union, cast
  38. )
  39.  
  40. import mypy.build
  41. import mypy.parse
  42. import mypy.errors
  43. import mypy.traverser
  44. from mypy import defaults
  45. from mypy.nodes import (
  46.     Expression, IntExpr, UnaryExpr, StrExpr, BytesExpr, NameExpr, FloatExpr, MemberExpr, TupleExpr,
  47.     ListExpr, ComparisonExpr, CallExpr, ClassDef, MypyFile, Decorator, AssignmentStmt,
  48.     IfStmt, ImportAll, ImportFrom, Import, FuncDef, FuncBase,
  49.     ARG_STAR, ARG_STAR2, ARG_NAMED, ARG_NAMED_OPT,
  50. )
  51. from mypy.stubgenc import parse_all_signatures, find_unique_signatures, generate_stub_for_c_module
  52. from mypy.stubutil import is_c_module, write_header
  53. from mypy.options import Options as MypyOptions
  54.  
  55.  
  56. Options = NamedTuple('Options', [('pyversion', Tuple[int, int]),
  57.                                  ('no_import', bool),
  58.                                  ('doc_dir', str),
  59.                                  ('search_path', List[str]),
  60.                                  ('interpreter', str),
  61.                                  ('modules', List[str]),
  62.                                  ('ignore_errors', bool),
  63.                                  ('recursive', bool),
  64.                                  ('include_private', bool),
  65.                                  ])
  66.  
  67.  
  68. class CantImport(Exception):
  69.     pass
  70.  
  71.  
  72. def generate_stub_for_module(module: str, output_dir: str, quiet: bool = False,
  73.                              add_header: bool = False, sigs: Dict[str, str] = {},
  74.                              class_sigs: Dict[str, str] = {},
  75.                              pyversion: Tuple[int, int] = defaults.PYTHON3_VERSION,
  76.                              no_import: bool = False,
  77.                              search_path: List[str] = [],
  78.                              interpreter: str = sys.executable,
  79.                              include_private: bool = False) -> None:
  80.     target = module.replace('.', '/')
  81.     try:
  82.         result = find_module_path_and_all(module=module,
  83.                                           pyversion=pyversion,
  84.                                           no_import=no_import,
  85.                                           search_path=search_path,
  86.                                           interpreter=interpreter)
  87.     except CantImport:
  88.         if not quiet:
  89.             traceback.print_exc()
  90.         print('Failed to import %s; skipping it' % module)
  91.         return
  92.  
  93.     if not result:
  94.         # C module
  95.         target = os.path.join(output_dir, target + '.pyi')
  96.         generate_stub_for_c_module(module_name=module,
  97.                                    target=target,
  98.                                    add_header=add_header,
  99.                                    sigs=sigs,
  100.                                    class_sigs=class_sigs)
  101.     else:
  102.         # Python module
  103.         module_path, module_all = result
  104.         if os.path.basename(module_path) == '__init__.py':
  105.             target += '/__init__.pyi'
  106.         else:
  107.             target += '.pyi'
  108.         target = os.path.join(output_dir, target)
  109.         generate_stub(module_path, output_dir, module_all,
  110.                       target=target, add_header=add_header, module=module,
  111.                       pyversion=pyversion, include_private=include_private)
  112.     if not quiet:
  113.         print('Created %s' % target)
  114.  
  115.  
  116. def find_module_path_and_all(module: str, pyversion: Tuple[int, int],
  117.                              no_import: bool,
  118.                              search_path: List[str],
  119.                              interpreter: str) -> Optional[Tuple[str,
  120.                                                                  Optional[List[str]]]]:
  121.     """Find module and determine __all__.
  122.    Return None if the module is a C module. Return (module_path, __all__) if
  123.    Python module. Raise an exception or exit if failed.
  124.    """
  125.     if not no_import:
  126.         if pyversion[0] == 2:
  127.             module_path, module_all = load_python_module_info(module, interpreter)
  128.         else:
  129.             # TODO: Support custom interpreters.
  130.             try:
  131.                 mod = importlib.import_module(module)
  132.             except Exception:
  133.                 raise CantImport(module)
  134.             if is_c_module(mod):
  135.                 return None
  136.             module_path = mod.__file__
  137.             module_all = getattr(mod, '__all__', None)
  138.     else:
  139.         # Find module by going through search path.
  140.         module_path = mypy.build.find_module(module, ['.'] + search_path)
  141.         if not module_path:
  142.             raise SystemExit(
  143.                 "Can't find module '{}' (consider using --search-path)".format(module))
  144.         module_all = None
  145.     return module_path, module_all
  146.  
  147.  
  148. def load_python_module_info(module: str, interpreter: str) -> Tuple[str, Optional[List[str]]]:
  149.     """Return tuple (module path, module __all__) for a Python 2 module.
  150.    The path refers to the .py/.py[co] file. The second tuple item is
  151.    None if the module doesn't define __all__.
  152.    Exit if the module can't be imported or if it's a C extension module.
  153.    """
  154.     cmd_template = '{interpreter} -c "%s"'.format(interpreter=interpreter)
  155.     code = ("import importlib, json; mod = importlib.import_module('%s'); "
  156.             "print(mod.__file__); print(json.dumps(getattr(mod, '__all__', None)))") % module
  157.     try:
  158.         output_bytes = subprocess.check_output(cmd_template % code, shell=True)
  159.     except subprocess.CalledProcessError:
  160.         print("Can't import module %s" % module, file=sys.stderr)
  161.         sys.exit(1)
  162.     output = output_bytes.decode('ascii').strip().splitlines()
  163.     module_path = output[0]
  164.     if not module_path.endswith(('.py', '.pyc', '.pyo')):
  165.         raise SystemExit('%s looks like a C module; they are not supported for Python 2' %
  166.                          module)
  167.     if module_path.endswith(('.pyc', '.pyo')):
  168.         module_path = module_path[:-1]
  169.     module_all = json.loads(output[1])
  170.     return module_path, module_all
  171.  
  172.  
  173. def generate_stub(path: str, output_dir: str, _all_: Optional[List[str]] = None,
  174.                   target: str = None, add_header: bool = False, module: str = None,
  175.                   pyversion: Tuple[int, int] = defaults.PYTHON3_VERSION,
  176.                   include_private: bool = False
  177.                   ) -> None:
  178.     with open(path, 'rb') as f:
  179.         source = f.read()
  180.     options = MypyOptions()
  181.     options.python_version = pyversion
  182.     try:
  183.         ast = mypy.parse.parse(source, fnam=path, errors=None, options=options)
  184.     except mypy.errors.CompileError as e:
  185.         # Syntax error!
  186.         for m in e.messages:
  187.             sys.stderr.write('%s\n' % m)
  188.         sys.exit(1)
  189.  
  190.     gen = StubGenerator(_all_, pyversion=pyversion, include_private=include_private)
  191.     ast.accept(gen)
  192.     if not target:
  193.         target = os.path.join(output_dir, os.path.basename(path))
  194.     subdir = os.path.dirname(target)
  195.     if subdir and not os.path.isdir(subdir):
  196.         os.makedirs(subdir)
  197.     with open(target, 'w') as file:
  198.         if add_header:
  199.             write_header(file, module, pyversion=pyversion)
  200.         file.write(''.join(gen.output()))
  201.  
  202.  
  203. # What was generated previously in the stub file. We keep track of these to generate
  204. # nicely formatted output (add empty line between non-empty classes, for example).
  205. EMPTY = 'EMPTY'
  206. FUNC = 'FUNC'
  207. CLASS = 'CLASS'
  208. EMPTY_CLASS = 'EMPTY_CLASS'
  209. VAR = 'VAR'
  210. NOT_IN_ALL = 'NOT_IN_ALL'
  211.  
  212.  
  213. class StubGenerator(mypy.traverser.TraverserVisitor):
  214.     def __init__(self, _all_: Optional[List[str]], pyversion: Tuple[int, int],
  215.                  include_private: bool = False) -> None:
  216.         self._all_ = _all_
  217.         self._output = []  # type: List[str]
  218.         self._import_lines = []  # type: List[str]
  219.         self._imports = []  # type: List[str]
  220.         self._indent = ''
  221.         self._vars = [[]]  # type: List[List[str]]
  222.         self._state = EMPTY
  223.         self._toplevel_names = []  # type: List[str]
  224.         self._classes = set()  # type: Set[str]
  225.         self._base_classes = []  # type: List[str]
  226.         self._pyversion = pyversion
  227.         self._include_private = include_private
  228.  
  229.     def visit_mypy_file(self, o: MypyFile) -> None:
  230.         self._classes = find_classes(o)
  231.         for node in o.defs:
  232.             if isinstance(node, ClassDef):
  233.                 self._base_classes.extend(self.get_base_types(node))
  234.         super().visit_mypy_file(o)
  235.         undefined_names = [name for name in self._all_ or []
  236.                            if name not in self._toplevel_names]
  237.         if undefined_names:
  238.             if self._state != EMPTY:
  239.                 self.add('\n')
  240.             self.add('# Names in __all__ with no definition:\n')
  241.             for name in sorted(undefined_names):
  242.                 self.add('#   %s\n' % name)
  243.  
  244.     def visit_func_def(self, o: FuncDef) -> None:
  245.         if self.is_private_name(o.name()):
  246.             return
  247.         if self.is_not_in_all(o.name()):
  248.             return
  249.         if self.is_recorded_name(o.name()):
  250.             return
  251.         if not self._indent and self._state not in (EMPTY, FUNC):
  252.             self.add('\n')
  253.         if not self.is_top_level():
  254.             self_inits = find_self_initializers(o)
  255.             for init, value in self_inits:
  256.                 init_code = self.get_init(init, value)
  257.                 if init_code:
  258.                     self.add(init_code)
  259.         self.add("%sdef %s(" % (self._indent, o.name()))
  260.         self.record_name(o.name())
  261.         args = []  # type: List[str]
  262.         for i, arg_ in enumerate(o.arguments):
  263.             var = arg_.variable
  264.             kind = arg_.kind
  265.             name = var.name()
  266.             init_stmt = arg_.initialization_statement
  267.             if init_stmt:
  268.                 if kind in (ARG_NAMED, ARG_NAMED_OPT) and '*' not in args:
  269.                     args.append('*')
  270.                 typename = self.get_str_type_of_node(init_stmt.rvalue, True)
  271.                 arg = '{}: {} = ...'.format(name, typename)
  272.             elif kind == ARG_STAR:
  273.                 arg = '*%s' % name
  274.             elif kind == ARG_STAR2:
  275.                 arg = '**%s' % name
  276.             else:
  277.                 arg = name
  278.             args.append(arg)
  279.         retname = None
  280.         if o.name() == '__init__':
  281.             retname = 'None'
  282.         retfield = ''
  283.         if retname is not None:
  284.             retfield = ' -> ' + retname
  285.  
  286.         self.add(', '.join(args))
  287.         self.add("){}: ...\n".format(retfield))
  288.         self._state = FUNC
  289.  
  290.     def visit_decorator(self, o: Decorator) -> None:
  291.         if self.is_private_name(o.func.name()):
  292.             return
  293.         for decorator in o.decorators:
  294.             if isinstance(decorator, NameExpr) and decorator.name in ('property',
  295.                                                                       'staticmethod',
  296.                                                                       'classmethod'):
  297.                 self.add('%s@%s\n' % (self._indent, decorator.name))
  298.             elif (isinstance(decorator, MemberExpr) and decorator.name == 'setter' and
  299.                   isinstance(decorator.expr, NameExpr)):
  300.                 self.add('%s@%s.setter\n' % (self._indent, decorator.expr.name))
  301.         super().visit_decorator(o)
  302.  
  303.     def visit_class_def(self, o: ClassDef) -> None:
  304.         sep = None  # type: Optional[int]
  305.         if not self._indent and self._state != EMPTY:
  306.             sep = len(self._output)
  307.             self.add('\n')
  308.         self.add('%sclass %s' % (self._indent, o.name))
  309.         self.record_name(o.name)
  310.         base_types = self.get_base_types(o)
  311.         if base_types:
  312.             self.add('(%s)' % ', '.join(base_types))
  313.         self.add(':\n')
  314.         n = len(self._output)
  315.         self._indent += '    '
  316.         self._vars.append([])
  317.         super().visit_class_def(o)
  318.         self._indent = self._indent[:-4]
  319.         self._vars.pop()
  320.         if len(self._output) == n:
  321.             if self._state == EMPTY_CLASS and sep is not None:
  322.                 self._output[sep] = ''
  323.             self._output[-1] = self._output[-1][:-1] + ' ...\n'
  324.             self._state = EMPTY_CLASS
  325.         else:
  326.             self._state = CLASS
  327.  
  328.     def get_base_types(self, cdef: ClassDef) -> List[str]:
  329.         base_types = []  # type: List[str]
  330.         for base in cdef.base_type_exprs:
  331.             if isinstance(base, NameExpr):
  332.                 if base.name != 'object':
  333.                     base_types.append(base.name)
  334.             elif isinstance(base, MemberExpr):
  335.                 modname = get_qualified_name(base.expr)
  336.                 base_types.append('%s.%s' % (modname, base.name))
  337.                 self.add_import_line('import %s\n' % modname)
  338.         return base_types
  339.  
  340.     def visit_assignment_stmt(self, o: AssignmentStmt) -> None:
  341.         foundl = []
  342.  
  343.         for lvalue in o.lvalues:
  344.             if isinstance(lvalue, NameExpr) and self.is_namedtuple(o.rvalue):
  345.                 assert isinstance(o.rvalue, CallExpr)
  346.                 self.process_namedtuple(lvalue, o.rvalue)
  347.                 continue
  348.             if isinstance(lvalue, TupleExpr):
  349.                 items = lvalue.items
  350.             elif isinstance(lvalue, ListExpr):
  351.                 items = lvalue.items
  352.             else:
  353.                 items = [lvalue]
  354.             sep = False
  355.             found = False
  356.             for item in items:
  357.                 if isinstance(item, NameExpr):
  358.                     init = self.get_init(item.name, o.rvalue)
  359.                     if init:
  360.                         found = True
  361.                         if not sep and not self._indent and \
  362.                            self._state not in (EMPTY, VAR):
  363.                             init = '\n' + init
  364.                             sep = True
  365.                         self.add(init)
  366.                         self.record_name(item.name)
  367.             foundl.append(found)
  368.  
  369.         if all(foundl):
  370.             self._state = VAR
  371.  
  372.     def is_namedtuple(self, expr: Expression) -> bool:
  373.         if not isinstance(expr, CallExpr):
  374.             return False
  375.         callee = expr.callee
  376.         return ((isinstance(callee, NameExpr) and callee.name.endswith('namedtuple')) or
  377.                 (isinstance(callee, MemberExpr) and callee.name == 'namedtuple'))
  378.  
  379.     def process_namedtuple(self, lvalue: NameExpr, rvalue: CallExpr) -> None:
  380.         self.add_import_line('from collections import namedtuple\n')
  381.         if self._state != EMPTY:
  382.             self.add('\n')
  383.         name = repr(getattr(rvalue.args[0], 'value', '<ERROR>'))
  384.         if isinstance(rvalue.args[1], StrExpr):
  385.             items = repr(rvalue.args[1].value)
  386.         elif isinstance(rvalue.args[1], ListExpr):
  387.             list_items = cast(List[StrExpr], rvalue.args[1].items)
  388.             items = '[%s]' % ', '.join(repr(item.value) for item in list_items)
  389.         else:
  390.             items = '<ERROR>'
  391.         self.add('%s = namedtuple(%s, %s)\n' % (lvalue.name, name, items))
  392.         self._classes.add(lvalue.name)
  393.         self._state = CLASS
  394.  
  395.     def visit_if_stmt(self, o: IfStmt) -> None:
  396.         # Ignore if __name__ == '__main__'.
  397.         expr = o.expr[0]
  398.         if (isinstance(expr, ComparisonExpr) and
  399.                 isinstance(expr.operands[0], NameExpr) and
  400.                 isinstance(expr.operands[1], StrExpr) and
  401.                 expr.operands[0].name == '__name__' and
  402.                 '__main__' in expr.operands[1].value):
  403.             return
  404.         super().visit_if_stmt(o)
  405.  
  406.     def visit_import_all(self, o: ImportAll) -> None:
  407.         self.add_import_line('from %s%s import *\n' % ('.' * o.relative, o.id))
  408.  
  409.     def visit_import_from(self, o: ImportFrom) -> None:
  410.         exported_names = set()  # type: Set[str]
  411.         if self._all_:
  412.             # Include import froms that import names defined in __all__.
  413.             names = [name for name, alias in o.names
  414.                      if name in self._all_ and alias is None]
  415.             exported_names.update(names)
  416.             self.import_and_export_names(o.id, o.relative, names)
  417.         else:
  418.             # Include import from targets that import from a submodule of a package.
  419.             if o.relative:
  420.                 sub_names = [name for name, alias in o.names
  421.                              if alias is None]
  422.                 exported_names.update(sub_names)
  423.                 self.import_and_export_names(o.id, o.relative, sub_names)
  424.         # Import names used as base classes.
  425.         base_names = [(name, alias) for name, alias in o.names
  426.                       if alias or name in self._base_classes and name not in exported_names]
  427.         if base_names:
  428.             imp_names = []  # type: List[str]
  429.             for name, alias in base_names:
  430.                 if alias is not None and alias != name:
  431.                     imp_names.append('%s as %s' % (name, alias))
  432.                 else:
  433.                     imp_names.append(name)
  434.             self.add_import_line('from %s%s import %s\n' % (
  435.                 '.' * o.relative, o.id, ', '.join(imp_names)))
  436.  
  437.     def import_and_export_names(self, module_id: str, relative: int, names: Iterable[str]) -> None:
  438.         """Import names from a module and export them (via from ... import x as x)."""
  439.         if names and module_id:
  440.             full_module_name = '%s%s' % ('.' * relative, module_id)
  441.             imported_names = ', '.join(['%s as %s' % (name, name) for name in names])
  442.             self.add_import_line('from %s import %s\n' % (full_module_name, imported_names))
  443.             for name in names:
  444.                 self.record_name(name)
  445.  
  446.     def visit_import(self, o: Import) -> None:
  447.         for id, as_id in o.ids:
  448.             if as_id is None:
  449.                 target_name = id.split('.')[0]
  450.             else:
  451.                 target_name = as_id
  452.             if self._all_ and target_name in self._all_ and (as_id is not None or
  453.                                                              '.' not in id):
  454.                 self.add_import_line('import %s as %s\n' % (id, target_name))
  455.                 self.record_name(target_name)
  456.  
  457.     def get_init(self, lvalue: str, rvalue: Expression) -> Optional[str]:
  458.         """Return initializer for a variable.
  459.        Return None if we've generated one already or if the variable is internal.
  460.        """
  461.         if lvalue in self._vars[-1]:
  462.             # We've generated an initializer already for this variable.
  463.             return None
  464.         # TODO: Only do this at module top level.
  465.         if self.is_private_name(lvalue) or self.is_not_in_all(lvalue):
  466.             return None
  467.         self._vars[-1].append(lvalue)
  468.         typename = self.get_str_type_of_node(rvalue)
  469.         return '%s%s = ...  # type: %s\n' % (self._indent, lvalue, typename)
  470.  
  471.     def add(self, string: str) -> None:
  472.         """Add text to generated stub."""
  473.         self._output.append(string)
  474.  
  475.     def add_typing_import(self, name: str) -> None:
  476.         """Add a name to be imported from typing, unless it's imported already.
  477.        The import will be internal to the stub.
  478.        """
  479.         if name not in self._imports:
  480.             self._imports.append(name)
  481.  
  482.     def add_import_line(self, line: str) -> None:
  483.         """Add a line of text to the import section, unless it's already there."""
  484.         if line not in self._import_lines:
  485.             self._import_lines.append(line)
  486.  
  487.     def output(self) -> str:
  488.         """Return the text for the stub."""
  489.         imports = ''
  490.         if self._imports:
  491.             imports += 'from typing import %s\n' % ", ".join(sorted(self._imports))
  492.         if self._import_lines:
  493.             imports += ''.join(self._import_lines)
  494.         if imports and self._output:
  495.             imports += '\n'
  496.         return imports + ''.join(self._output)
  497.  
  498.     def is_not_in_all(self, name: str) -> bool:
  499.         if self.is_private_name(name):
  500.             return False
  501.         if self._all_:
  502.             return self.is_top_level() and name not in self._all_
  503.         return False
  504.  
  505.     def is_private_name(self, name: str) -> bool:
  506.         if self._include_private:
  507.             return False
  508.         return name.startswith('_') and (not name.endswith('__')
  509.                                          or name in ('__all__',
  510.                                                      '__author__',
  511.                                                      '__version__',
  512.                                                      '__str__',
  513.                                                      '__repr__',
  514.                                                      '__getstate__',
  515.                                                      '__setstate__',
  516.                                                      '__slots__'))
  517.  
  518.     def get_str_type_of_node(self, rvalue: Expression,
  519.                              can_infer_optional: bool = False) -> str:
  520.         if isinstance(rvalue, IntExpr):
  521.             return 'int'
  522.         if isinstance(rvalue, StrExpr):
  523.             return 'str'
  524.         if isinstance(rvalue, BytesExpr):
  525.             return 'bytes'
  526.         if isinstance(rvalue, FloatExpr):
  527.             return 'float'
  528.         if isinstance(rvalue, UnaryExpr) and isinstance(rvalue.expr, IntExpr):
  529.             return 'int'
  530.         if isinstance(rvalue, NameExpr) and rvalue.name in ('True', 'False'):
  531.             return 'bool'
  532.         if can_infer_optional and \
  533.            isinstance(rvalue, NameExpr) and rvalue.name == 'None':
  534.             self.add_typing_import('Optional')
  535.             self.add_typing_import('Any')
  536.             return 'Optional[Any]'
  537.         self.add_typing_import('Any')
  538.         return 'Any'
  539.  
  540.     def is_top_level(self) -> bool:
  541.         """Are we processing the top level of a file?"""
  542.         return self._indent == ''
  543.  
  544.     def record_name(self, name: str) -> None:
  545.         """Mark a name as defined.
  546.        This only does anything if at the top level of a module.
  547.        """
  548.         if self.is_top_level():
  549.             self._toplevel_names.append(name)
  550.  
  551.     def is_recorded_name(self, name: str) -> bool:
  552.         """Has this name been recorded previously?"""
  553.         return self.is_top_level() and name in self._toplevel_names
  554.  
  555.  
  556. def find_self_initializers(fdef: FuncBase) -> List[Tuple[str, Expression]]:
  557.     results = []  # type: List[Tuple[str, Expression]]
  558.  
  559.     class SelfTraverser(mypy.traverser.TraverserVisitor):
  560.         def visit_assignment_stmt(self, o: AssignmentStmt) -> None:
  561.             lvalue = o.lvalues[0]
  562.             if (isinstance(lvalue, MemberExpr) and
  563.                     isinstance(lvalue.expr, NameExpr) and
  564.                     lvalue.expr.name == 'self'):
  565.                 results.append((lvalue.name, o.rvalue))
  566.  
  567.     fdef.accept(SelfTraverser())
  568.     return results
  569.  
  570.  
  571. def find_classes(node: MypyFile) -> Set[str]:
  572.     results = set()  # type: Set[str]
  573.  
  574.     class ClassTraverser(mypy.traverser.TraverserVisitor):
  575.         def visit_class_def(self, o: ClassDef) -> None:
  576.             results.add(o.name)
  577.  
  578.     node.accept(ClassTraverser())
  579.     return results
  580.  
  581.  
  582. def get_qualified_name(o: Expression) -> str:
  583.     if isinstance(o, NameExpr):
  584.         return o.name
  585.     elif isinstance(o, MemberExpr):
  586.         return '%s.%s' % (get_qualified_name(o.expr), o.name)
  587.     else:
  588.         return '<ERROR>'
  589.  
  590.  
  591. def walk_packages(packages: List[str]) -> Iterator[str]:
  592.     for package_name in packages:
  593.         package = __import__(package_name)
  594.         yield package.__name__
  595.         for importer, qualified_name, ispkg in pkgutil.walk_packages(package.__path__,
  596.                                                                      prefix=package.__name__ + ".",
  597.                                                                      onerror=lambda r: None):
  598.             yield qualified_name
  599.  
  600.  
  601. def main() -> None:
  602.     options = parse_options(sys.argv[1:])
  603.     if not os.path.isdir('out'):
  604.         raise SystemExit('Directory "out" does not exist')
  605.     if options.recursive and options.no_import:
  606.         raise SystemExit('recursive stub generation without importing is not currently supported')
  607.     sigs = {}  # type: Any
  608.     class_sigs = {}  # type: Any
  609.     if options.doc_dir:
  610.         all_sigs = []  # type: Any
  611.         all_class_sigs = []  # type: Any
  612.         for path in glob.glob('%s/*.rst' % options.doc_dir):
  613.             with open(path) as f:
  614.                 func_sigs, class_sigs = parse_all_signatures(f.readlines())
  615.             all_sigs += func_sigs
  616.             all_class_sigs += class_sigs
  617.         sigs = dict(find_unique_signatures(all_sigs))
  618.         class_sigs = dict(find_unique_signatures(all_class_sigs))
  619.     for module in (options.modules if not options.recursive else walk_packages(options.modules)):
  620.         try:
  621.             generate_stub_for_module(module, 'out',
  622.                                      add_header=True,
  623.                                      sigs=sigs,
  624.                                      class_sigs=class_sigs,
  625.                                      pyversion=options.pyversion,
  626.                                      no_import=options.no_import,
  627.                                      search_path=options.search_path,
  628.                                      interpreter=options.interpreter,
  629.                                      include_private=options.include_private)
  630.         except Exception as e:
  631.             if not options.ignore_errors:
  632.                 raise e
  633.             else:
  634.                 print("Stub generation failed for", module, file=sys.stderr)
  635.  
  636.  
  637. def parse_options(args: List[str]) -> Options:
  638.     pyversion = defaults.PYTHON3_VERSION
  639.     no_import = False
  640.     recursive = False
  641.     ignore_errors = False
  642.     doc_dir = ''
  643.     search_path = []  # type: List[str]
  644.     interpreter = ''
  645.     include_private = False
  646.     while args and args[0].startswith('-'):
  647.         if args[0] == '--doc-dir':
  648.             doc_dir = args[1]
  649.             args = args[1:]
  650.         elif args[0] == '--search-path':
  651.             if not args[1]:
  652.                 usage()
  653.             search_path = args[1].split(':')
  654.             args = args[1:]
  655.         elif args[0] == '-p':
  656.             interpreter = args[1]
  657.             args = args[1:]
  658.         elif args[0] == '--recursive':
  659.             recursive = True
  660.         elif args[0] == '--ignore-errors':
  661.             ignore_errors = True
  662.         elif args[0] == '--py2':
  663.             pyversion = defaults.PYTHON2_VERSION
  664.         elif args[0] == '--no-import':
  665.             no_import = True
  666.         elif args[0] == '--include-private':
  667.             include_private = True
  668.         elif args[0] in ('-h', '--help'):
  669.             usage()
  670.         else:
  671.             raise SystemExit('Unrecognized option %s' % args[0])
  672.         args = args[1:]
  673.     if not args:
  674.         usage()
  675.     if not interpreter:
  676.         interpreter = sys.executable if pyversion[0] == 3 else default_python2_interpreter()
  677.     return Options(pyversion=pyversion,
  678.                    no_import=no_import,
  679.                    doc_dir=doc_dir,
  680.                    search_path=search_path,
  681.                    interpreter=interpreter,
  682.                    modules=args,
  683.                    ignore_errors=ignore_errors,
  684.                    recursive=recursive,
  685.                    include_private=include_private)
  686.  
  687.  
  688. def default_python2_interpreter() -> str:
  689.     # TODO: Make this do something reasonable in Windows.
  690.     for candidate in ('/usr/bin/python2', '/usr/bin/python'):
  691.         if not os.path.exists(candidate):
  692.             continue
  693.         output = subprocess.check_output([candidate, '--version'],
  694.                                          stderr=subprocess.STDOUT).strip()
  695.         if b'Python 2' in output:
  696.             return candidate
  697.     raise SystemExit("Can't find a Python 2 interpreter -- please use the -p option")
  698.  
  699.  
  700. def usage() -> None:
  701.     usage = textwrap.dedent("""\
  702.        usage: stubgen [--py2] [--no-import] [--doc-dir PATH]
  703.                       [--search-path PATH] [-p PATH] MODULE ...
  704.        Generate draft stubs for modules.
  705.        Stubs are generated in directory ./out, to avoid overriding files with
  706.        manual changes.  This directory is assumed to exist.
  707.        Options:
  708.          --py2           run in Python 2 mode (default: Python 3 mode)
  709.          --recursive     traverse listed modules to generate inner package modules as well
  710.          --ignore-errors ignore errors when trying to generate stubs for modules
  711.          --no-import     don't import the modules, just parse and analyze them
  712.                          (doesn't work with C extension modules and doesn't
  713.                          respect __all__)
  714.          --include-private
  715.                          generate stubs for objects and members considered private
  716.                          (single leading undescore and no trailing underscores)
  717.          --doc-dir PATH  use .rst documentation in PATH (this may result in
  718.                          better stubs in some cases; consider setting this to
  719.                          DIR/Python-X.Y.Z/Doc/library)
  720.          --search-path PATH
  721.                          specify module search directories, separated by ':'
  722.                          (currently only used if --no-import is given)
  723.          -p PATH         use Python interpreter at PATH (only works for
  724.                          Python 2 right now)
  725.          -h, --help      print this help message and exit
  726.    """.rstrip())
  727.  
  728.     raise SystemExit(usage)
  729.  
  730.  
  731. if __name__ == '__main__':
  732.     main()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement