Improve stub generation so it works better with pybind11 wrappers.

mypy/stubgen.py

@@ -41,6 +41,7 @@
 import os
 import os.path
 import pkgutil
+import inspect
 import subprocess
 import sys
 import textwrap
@@ -211,7 +212,6 @@ def generate_stub(path: str,
  pyversion: Tuple[int, int] = defaults.PYTHON3_VERSION,
  include_private: bool = False
  ) -> None:
-
  with open(path, 'rb') as f:
  data = f.read()
  source = mypy.util.decode_python_encoding(data, pyversion)
@@ -254,7 +254,7 @@ def __init__(self, stubgen: 'StubGenerator') -> None:
  super().__init__()
  self.stubgen = stubgen
 
- def visit_unbound_type(self, t: UnboundType)-> str:
+ def visit_unbound_type(self, t: UnboundType) -> str:
  s = t.name
  base = s.split('.')[0]
  self.stubgen.import_tracker.require_name(base)
@@ -348,7 +348,7 @@ def add_import_from(self, module: str, names: List[Tuple[str, Optional[str]]]) -
  if alias:
  self.reverse_alias[alias] = name
 
- def add_import(self, module: str, alias: Optional[str]=None) -> None:
+ def add_import(self, module: str, alias: Optional[str] = None) -> None:
  name = module.split('.')[0]
  self.module_for[alias or name] = None
  self.direct_imports[name] = module
@@ -593,7 +593,7 @@ def visit_assignment_stmt(self, o: AssignmentStmt) -> None:
  if init:
  found = True
  if not sep and not self._indent and \
- self._state not in (EMPTY, VAR):
+  self._state not in (EMPTY, VAR):
  init = '\n' + init
  sep = True
  self.add(init)
@@ -625,7 +625,7 @@ def process_namedtuple(self, lvalue: NameExpr, rvalue: CallExpr) -> None:
  self.add('%s = namedtuple(%s, %s)\n' % (lvalue.name, name, items))
  self._state = CLASS
 
- def is_type_expression(self, expr: Expression, top_level: bool=True) -> bool:
+ def is_type_expression(self, expr: Expression, top_level: bool = True) -> bool:
  """Return True for things that look like type expressions
 
  Used to know if assignments look like typealiases
@@ -795,7 +795,7 @@ def get_str_type_of_node(self, rvalue: Expression,
  if isinstance(rvalue, NameExpr) and rvalue.name in ('True', 'False'):
  return 'bool'
  if can_infer_optional and \
- isinstance(rvalue, NameExpr) and rvalue.name == 'None':
+  isinstance(rvalue, NameExpr) and rvalue.name == 'None':
  self.add_typing_import('Optional')
  self.add_typing_import('Any')
  return 'Optional[Any]'
@@ -850,7 +850,6 @@ def visit_return_stmt(self, o: ReturnStmt) -> None:
 
 
 def has_return_statement(fdef: FuncBase) -> bool:
-
  seeker = ReturnSeeker()
  fdef.accept(seeker)
  return seeker.found
@@ -866,17 +865,35 @@ def get_qualified_name(o: Expression) -> str:
 
 
 def walk_packages(packages: List[str]) -> Iterator[str]:
+ """Iterates through all packages and sub-packages in the given list.
+
+ Python packages have a __path__ attribute defined, which pkgutil uses to determine
+ the package hierarchy. However, packages in C extensions do not have this attribute,
+ so we have to roll out our own.
+ """
  for package_name in packages:
  package = importlib.import_module(package_name)
  yield package.__name__
+ # get the path of the object (needed by pkgutil)
  path = getattr(package, '__path__', None)
  if path is None:
+ # object has no path; this means it's either a module inside a package
+ # (and thus no sub-packages), or it could be a C extension package.
+ if is_c_module(package):
+ # This is a C extension module, now get the list of all sub-packages
+ # using the inspect module
+ subpackages = [package.__name__ + "." + name
+ for name, val in inspect.getmembers(package)
+ if inspect.ismodule(val)]
+ # recursively iterate through the subpackages
+ for submodule in walk_packages(subpackages):
+ yield submodule
  # It's a module inside a package. There's nothing else to walk/yield.
-  continue
- for importer, qualified_name, ispkg in pkgutil.walk_packages(path,
-  prefix=package.__name__ + ".",
-  onerror=lambda r: None):
- yield qualified_name
+ else:
+  all_packages = pkgutil.walk_packages(path, prefix=package.__name__ + ".",
+ onerror=lambda r: None)
+ for importer, qualified_name, ispkg in all_packages:
+  yield qualified_name
 
 
 def main() -> None:
@@ -994,7 +1011,7 @@ def default_python2_interpreter() -> str:
  raise SystemExit("Can't find a Python 2 interpreter -- please use the -p option")
 
 
-def usage(exit_nonzero: bool=True) -> None:
+def usage(exit_nonzero: bool = True) -> None:
  usage = textwrap.dedent("""\
  usage: stubgen [--py2] [--no-import] [--doc-dir PATH]
  [--search-path PATH] [-p PATH] [-o PATH]

mypy/stubgenc.py

@@ -4,12 +4,16 @@
 """
 
 import importlib
+import inspect
 import os.path
 import re
-from typing import List, Dict, Tuple, Optional, Mapping, Any
+from typing import List, Dict, Tuple, Optional, Mapping, Any, Set
 from types import ModuleType
 
-from mypy.stubutil import is_c_module, write_header, infer_sig_from_docstring
+from mypy.stubutil import (
+ is_c_module, write_header, infer_sig_from_docstring,
+ infer_prop_type_from_docstring
+)
 
 
 def generate_stub_for_c_module(module_name: str,
@@ -41,7 +45,7 @@ def generate_stub_for_c_module(module_name: str,
  for name, obj in items:
  if name.startswith('__') and name.endswith('__'):
  continue
- if name not in done:
+ if name not in done and not inspect.ismodule(obj):
  type_str = type(obj).__name__
  if type_str not in ('int', 'str', 'bytes', 'float', 'bool'):
  type_str = 'Any'
@@ -67,7 +71,7 @@ def generate_stub_for_c_module(module_name: str,
 
 def add_typing_import(output: List[str]) -> List[str]:
  names = []
- for name in ['Any']:
+ for name in ['Any', 'Union', 'Tuple', 'Optional', 'List', 'Dict']:
  if any(re.search(r'\b%s\b' % name, line) for line in output):
  names.append(name)
  if names:
@@ -77,22 +81,30 @@ def add_typing_import(output: List[str]) -> List[str]:
 
 
 def is_c_function(obj: object) -> bool:
- return type(obj) is type(ord)
+ return inspect.isbuiltin(obj) or type(obj) is type(ord)
 
 
 def is_c_method(obj: object) -> bool:
- return type(obj) in (type(str.index),
- type(str.__add__),
- type(str.__new__))
+ return inspect.ismethoddescriptor(obj) or type(obj) in (type(str.index),
+  type(str.__add__),
+  type(str.__new__))
 
 
 def is_c_classmethod(obj: object) -> bool:
- type_str = type(obj).__name__
- return type_str == 'classmethod_descriptor'
+ return inspect.isbuiltin(obj) or type(obj).__name__ in ('classmethod',
+ 'classmethod_descriptor')
+
+
+def is_c_property(obj: object) -> bool:
+ return inspect.isdatadescriptor(obj) and hasattr(obj, 'fget')
+
+
+def is_c_property_readonly(prop: object) -> bool:
+ return getattr(prop, 'fset') is None
 
 
 def is_c_type(obj: object) -> bool:
- return type(obj) is type(int)
+ return inspect.isclass(obj) or type(obj) is type(int)
 
 
 def generate_c_function_stub(module: ModuleType,
@@ -104,6 +116,8 @@ def generate_c_function_stub(module: ModuleType,
  class_name: Optional[str] = None,
  class_sigs: Dict[str, str] = {},
  ) -> None:
+ ret_type = 'Any'
+
  if self_var:
  self_arg = '%s, ' % self_var
  else:
@@ -115,19 +129,37 @@ def generate_c_function_stub(module: ModuleType,
  docstr = getattr(obj, '__doc__', None)
  inferred = infer_sig_from_docstring(docstr, name)
  if inferred:
- sig = inferred
+ sig, ret_type = inferred
  else:
  if class_name and name not in sigs:
  sig = infer_method_sig(name)
  else:
  sig = sigs.get(name, '(*args, **kwargs)')
+ # strip away parenthesis
  sig = sig[1:-1]
  if sig:
- if sig.split(',', 1)[0] == self_var:
- self_arg = ''
+ if self_var:
+ # remove annotation on self from signature if present
+ groups = sig.split(',', 1)
+ if groups[0] == self_var or groups[0].startswith(self_var + ':'):
+ self_arg = ''
+ sig = '{},{}'.format(self_var, groups[1]) if len(groups) > 1 else self_var
  else:
  self_arg = self_arg.replace(', ', '')
- output.append('def %s(%s%s): ...' % (name, self_arg, sig))
+ output.append('def %s(%s%s) -> %s: ...' % (name, self_arg, sig, ret_type))
+
+
+def generate_c_property_stub(name: str, obj: object, output: List[str], readonly: bool) -> None:
+ docstr = getattr(obj, '__doc__', None)
+ inferred = infer_prop_type_from_docstring(docstr)
+ if not inferred:
+ inferred = 'Any'
+
+ output.append('@property')
+ output.append('def {}(self) -> {}: ...'.format(name, inferred))
+ if not readonly:
+ output.append('@{}.setter'.format(name))
+ output.append('def {}(self, val: {}) -> None: ...'.format(name, inferred))
 
 
 def generate_c_type_stub(module: ModuleType,
@@ -141,8 +173,9 @@ def generate_c_type_stub(module: ModuleType,
  # (it could be a mappingproxy!), which makes mypyc mad, so obfuscate it.
  obj_dict = getattr(obj, '__dict__') # type: Mapping[str, Any]
  items = sorted(obj_dict.items(), key=lambda x: method_name_sort_key(x[0]))
- methods = []
- done = set()
+ methods = [] # type: List[str]
+ properties = [] # type: List[str]
+ done = set() # type: Set[str]
  for attr, value in items:
  if is_c_method(value) or is_c_classmethod(value):
  done.add(attr)
@@ -162,6 +195,10 @@ def generate_c_type_stub(module: ModuleType,
  attr = '__init__'
  generate_c_function_stub(module, attr, value, methods, self_var, sigs=sigs,
  class_name=class_name, class_sigs=class_sigs)
+ elif is_c_property(value):
+ done.add(attr)
+ generate_c_property_stub(attr, value, properties, is_c_property_readonly(value))
+
  variables = []
  for attr, value in items:
  if is_skipped_attribute(attr):
@@ -183,14 +220,16 @@ def generate_c_type_stub(module: ModuleType,
  bases_str = '(%s)' % ', '.join(base.__name__ for base in bases)
  else:
  bases_str = ''
- if not methods and not variables:
+ if not methods and not variables and not properties:
  output.append('class %s%s: ...' % (class_name, bases_str))
  else:
  output.append('class %s%s:' % (class_name, bases_str))
  for variable in variables:
  output.append(' %s' % variable)
  for method in methods:
  output.append(' %s' % method)
+ for prop in properties:
+ output.append(' %s' % prop)
 
 
 def method_name_sort_key(name: str) -> Tuple[int, str]:

mypy/stubutil.py

@@ -106,12 +106,44 @@ def write_header(file: IO[str], module_name: Optional[str] = None,
  '# NOTE: This dynamically typed stub was automatically generated by stubgen.\n\n')
 
 
-def infer_sig_from_docstring(docstr: str, name: str) -> Optional[str]:
+def infer_sig_from_docstring(docstr: str, name: str) -> Optional[Tuple[str, str]]:
  if not docstr:
  return None
  docstr = docstr.lstrip()
- m = re.match(r'%s(\([a-zA-Z0-9_=, ]*\))' % name, docstr)
+ # look for function signature, which is any string of the format
+ # <function_name>(<signature>) -> <return type>
+ # or perhaps without the return type
+
+ # in the signature, we allow the following characters:
+ # colon/equal: to match default values, like "a: int=1"
+ # comma/space/brackets: for type hints like "a: Tuple[int, float]"
+ # dot: for classes annotating using full path, like "a: foo.bar.baz"
+ # to capture return type,
+ sig_str = r'\([a-zA-Z0-9_=:, \[\]\.]*\)'
+ sig_match = r'%s(%s)' % (name, sig_str)
+ # first, try to capture return type; we just match until end of line
+ m = re.match(sig_match + ' -> ([a-zA-Z].*)$', docstr, re.MULTILINE)
  if m:
- return m.group(1)
- else:
+ # strip potential white spaces at the right of return type
+ return m.group(1), m.group(2).rstrip()
+
+ # try to not match return type
+ m = re.match(sig_match, docstr)
+ if m:
+ return m.group(1), 'Any'
+ return None
+
+
+def infer_prop_type_from_docstring(docstr: str) -> Optional[str]:
+ if not docstr:
  return None
+
+ # check for Google/Numpy style docstring type annotation
+ # the docstring has the format "<type>: <descriptions>"
+ # in the type string, we allow the following characters
+ # dot: because something classes are annotated using full path,
+ # brackets: to allow type hints like List[int]
+ # comma/space: things like Tuple[int, int]
+ test_str = r'^([a-zA-Z0-9_, \.\[\]]*): '
+ m = re.match(test_str, docstr)
+ return m.group(1) if m else None

mypy/test/teststubgen.py

@@ -17,7 +17,7 @@
 from mypy.stubgenc import generate_c_type_stub, infer_method_sig
 from mypy.stubutil import (
  parse_signature, parse_all_signatures, build_signature, find_unique_signatures,
- infer_sig_from_docstring
+ infer_sig_from_docstring, infer_prop_type_from_docstring
 )
 
 
@@ -103,12 +103,31 @@ def test_find_unique_signatures(self) -> None:
  ('func3', '(arg, arg2)')])
 
  def test_infer_sig_from_docstring(self) -> None:
- assert_equal(infer_sig_from_docstring('\nfunc(x) - y', 'func'), '(x)')
- assert_equal(infer_sig_from_docstring('\nfunc(x, Y_a=None)', 'func'), '(x, Y_a=None)')
+ assert_equal(infer_sig_from_docstring('\nfunc(x) - y', 'func'), ('(x)', 'Any'))
+ assert_equal(infer_sig_from_docstring('\nfunc(x, Y_a=None)', 'func'),
+ ('(x, Y_a=None)', 'Any'))
  assert_equal(infer_sig_from_docstring('\nafunc(x) - y', 'func'), None)
  assert_equal(infer_sig_from_docstring('\nfunc(x, y', 'func'), None)
  assert_equal(infer_sig_from_docstring('\nfunc(x=z(y))', 'func'), None)
  assert_equal(infer_sig_from_docstring('\nfunc x', 'func'), None)
+ # try to infer signature from type annotation
+ assert_equal(infer_sig_from_docstring('\nfunc(x: int)', 'func'), ('(x: int)', 'Any'))
+ assert_equal(infer_sig_from_docstring('\nfunc(x: int=3)', 'func'), ('(x: int=3)', 'Any'))
+ assert_equal(infer_sig_from_docstring('\nfunc(x: int=3) -> int', 'func'),
+ ('(x: int=3)', 'int'))
+ assert_equal(infer_sig_from_docstring('\nfunc(x: int=3) -> int \n', 'func'),
+ ('(x: int=3)', 'int'))
+ assert_equal(infer_sig_from_docstring('\nfunc(x: Tuple[int, str]) -> str', 'func'),
+ ('(x: Tuple[int, str])', 'str'))
+ assert_equal(infer_sig_from_docstring('\nfunc(x: foo.bar)', 'func'),
+ ('(x: foo.bar)', 'Any'))
+
+ def infer_prop_type_from_docstring(self) -> None:
+ assert_equal(infer_prop_type_from_docstring('str: A string.'), 'str')
+ assert_equal(infer_prop_type_from_docstring('Optional[int]: An int.'), 'Optional[int]')
+ assert_equal(infer_prop_type_from_docstring('Tuple[int, int]: A tuple.'),
+ 'Tuple[int, int]')
+ assert_equal(infer_prop_type_from_docstring('\nstr: A string.'), None)
 
 
 class StubgenPythonSuite(DataSuite):