Improve stub generation so it works better with pybind11 wrappers (#5814)

1. Modify walk_packages() so when --recursive option is set, submodules
   in C extensions are also visited properly.

2. use inspect package to implement is_c_function(), is_c_method(),
   is_c_classmethod(), and is_c_type() so it is more accurate.  The
   inspect methods are OR'ed with the old check to prevent breaking
   old code.

3. Improve infer_sig_from_docstring() so it supports Python-annotation
   style type hints (like a: int) in docstrings.

4. Misc. cleanup.

Author: pkerichang

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)
@@ -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)
@@ -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:

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):