Add option to include docstrings with stubgen

Add a --include-docstrings flag to stubgen
This was suggested in python#11965.
When using this flag, the .pyi files will include
docstrings for Python classes and functions and
for C extension functions.

Author: chylek

mypy/fastparse.py

@@ -1028,6 +1028,8 @@ def do_func_def(
             # FuncDef overrides set_line -- can't use self.set_line
             func_def.set_line(lineno, n.col_offset, end_line, end_column)
             retval = func_def
+        if ast3.get_docstring(n):
+            func_def.docstring = ast3.get_docstring(n, clean=False)
         self.class_and_function_stack.pop()
         return retval
 
@@ -1137,6 +1139,8 @@ def visit_ClassDef(self, n: ast3.ClassDef) -> ClassDef:
         else:
             cdef.line = n.lineno
             cdef.deco_line = n.decorator_list[0].lineno if n.decorator_list else None
+        if ast3.get_docstring(n):
+            cdef.docstring = ast3.get_docstring(n, clean=False)
         cdef.column = n.col_offset
         cdef.end_line = getattr(n, "end_lineno", None)
         cdef.end_column = getattr(n, "end_col_offset", None)

mypy/nodes.py

@@ -774,6 +774,7 @@ class FuncDef(FuncItem, SymbolNode, Statement):
         "is_abstract",
         "original_def",
         "deco_line",
+        "docstring",
     )
 
     # Note that all __init__ args must have default values
@@ -794,6 +795,7 @@ def __init__(
         self.original_def: Union[None, FuncDef, Var, Decorator] = None
         # Used for error reporting (to keep backwad compatibility with pre-3.8)
         self.deco_line: Optional[int] = None
+        self.docstring = None
 
     @property
     def name(self) -> str:
@@ -1070,6 +1072,7 @@ class ClassDef(Statement):
         "analyzed",
         "has_incompatible_baseclass",
         "deco_line",
+        "docstring",
     )
 
     name: str  # Name of the class without module prefix
@@ -1111,6 +1114,7 @@ def __init__(
         self.has_incompatible_baseclass = False
         # Used for error reporting (to keep backwad compatibility with pre-3.8)
         self.deco_line: Optional[int] = None
+        self.docstring: str = None
 
     def accept(self, visitor: StatementVisitor[T]) -> T:
         return visitor.visit_class_def(self)

mypy/stubgen.py

@@ -206,6 +206,7 @@ def __init__(
         verbose: bool,
         quiet: bool,
         export_less: bool,
+        include_docstrings: bool,
     ) -> None:
         # See parse_options for descriptions of the flags.
         self.pyversion = pyversion
@@ -224,6 +225,7 @@ def __init__(
         self.verbose = verbose
         self.quiet = quiet
         self.export_less = export_less
+        self.include_docstrings = include_docstrings
 
 
 class StubSource:
@@ -572,6 +574,7 @@ def __init__(
         include_private: bool = False,
         analyzed: bool = False,
         export_less: bool = False,
+        include_docstrings: bool = False,
     ) -> None:
         # Best known value of __all__.
         self._all_ = _all_
@@ -587,6 +590,7 @@ def __init__(
         self._toplevel_names: List[str] = []
         self._pyversion = pyversion
         self._include_private = include_private
+        self._include_docstrings = include_docstrings
         self.import_tracker = ImportTracker()
         # Was the tree semantically analysed before?
         self.analyzed = analyzed
@@ -754,7 +758,11 @@ def visit_func_def(
             retfield = " -> " + retname
 
         self.add(", ".join(args))
-        self.add(f"){retfield}: ...\n")
+        self.add(f"){retfield}:")
+        if self._include_docstrings and o.docstring:
+            self.add(f'\n{self._indent}    """{o.docstring}"""\n{self._indent}   ')
+
+        self.add(" ...\n")
         self._state = FUNC
 
     def is_none_expr(self, expr: Expression) -> bool:
@@ -926,8 +934,10 @@ def visit_class_def(self, o: ClassDef) -> None:
         if base_types:
             self.add(f"({', '.join(base_types)})")
         self.add(":\n")
-        n = len(self._output)
         self._indent += "    "
+        if o.docstring:
+            self.add(f'{self._indent}"""{o.docstring}"""\n')
+        n = len(self._output)
         self._vars.append([])
         super().visit_class_def(o)
         self._indent = self._indent[:-4]
@@ -1605,6 +1615,7 @@ def generate_stub_from_ast(
     pyversion: Tuple[int, int] = defaults.PYTHON3_VERSION,
     include_private: bool = False,
     export_less: bool = False,
+    include_docstrings: bool = False,
 ) -> None:
     """Use analysed (or just parsed) AST to generate type stub for single file.
 
@@ -1617,6 +1628,7 @@ def generate_stub_from_ast(
         include_private=include_private,
         analyzed=not parse_only,
         export_less=export_less,
+        include_docstrings=include_docstrings,
     )
     assert mod.ast is not None, "This function must be used only with analyzed modules"
     mod.ast.accept(gen)
@@ -1677,6 +1689,7 @@ def generate_stubs(options: Options) -> None:
                 options.pyversion,
                 options.include_private,
                 options.export_less,
+                options.include_docstrings,
             )
 
     # Separately analyse C modules using different logic.
@@ -1688,7 +1701,7 @@ def generate_stubs(options: Options) -> None:
         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, sigs=sigs, class_sigs=class_sigs)
+            generate_stub_for_c_module(mod.module, target, sigs=sigs, class_sigs=class_sigs, include_docstrings=options.include_docstrings)
     num_modules = len(py_modules) + len(c_modules)
     if not options.quiet and num_modules > 0:
         print("Processed %d modules" % num_modules)
@@ -1743,6 +1756,13 @@ def parse_options(args: List[str]) -> Options:
             "don't implicitly export all names imported from other modules " "in the same package"
         ),
     )
+    parser.add_argument(
+        "--include-docstrings",
+        action="store_true",
+        help=(
+            "include existing docstrings with the stubs"
+        ),
+    )
     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(
@@ -1823,6 +1843,7 @@ def parse_options(args: List[str]) -> Options:
         verbose=ns.verbose,
         quiet=ns.quiet,
         export_less=ns.export_less,
+        include_docstrings=ns.include_docstrings,
     )
 
 

mypy/stubgenc.py

@@ -44,6 +44,7 @@ def generate_stub_for_c_module(
     target: str,
     sigs: Optional[Dict[str, str]] = None,
     class_sigs: Optional[Dict[str, str]] = None,
+    include_docstrings: bool = False,
 ) -> None:
     """Generate stub for C module.
 
@@ -64,15 +65,15 @@ def generate_stub_for_c_module(
     items = sorted(module.__dict__.items(), key=lambda x: x[0])
     for name, obj in items:
         if is_c_function(obj):
-            generate_c_function_stub(module, name, obj, functions, imports=imports, sigs=sigs)
+            generate_c_function_stub(module, name, obj, functions, imports=imports, sigs=sigs, include_docstrings=include_docstrings)
             done.add(name)
     types: List[str] = []
     for name, obj in items:
         if name.startswith("__") and name.endswith("__"):
             continue
         if is_c_type(obj):
             generate_c_type_stub(
-                module, name, obj, types, imports=imports, sigs=sigs, class_sigs=class_sigs
+                module, name, obj, types, imports=imports, sigs=sigs, class_sigs=class_sigs, include_docstrings=include_docstrings
             )
             done.add(name)
     variables = []
@@ -156,10 +157,11 @@ def generate_c_function_stub(
     sigs: Optional[Dict[str, str]] = None,
     class_name: Optional[str] = None,
     class_sigs: Optional[Dict[str, str]] = None,
+    include_docstrings: bool = False
 ) -> None:
     """Generate stub for a single function or method.
 
-    The result (always a single line) will be appended to 'output'.
+    The result will be appended to 'output'.
     If necessary, any required names will be added to 'imports'.
     The 'class_name' is used to find signature of __init__ or __new__ in
     'class_sigs'.
@@ -170,7 +172,7 @@ def generate_c_function_stub(
         class_sigs = {}
 
     ret_type = "None" if name == "__init__" and class_name else "Any"
-
+    docstr = None
     if (
         name in ("__new__", "__init__")
         and name not in sigs
@@ -236,13 +238,23 @@ def generate_c_function_stub(
 
             if is_overloaded:
                 output.append("@overload")
-            output.append(
-                "def {function}({args}) -> {ret}: ...".format(
-                    function=name,
-                    args=", ".join(sig),
-                    ret=strip_or_import(signature.ret_type, module, imports),
+            if include_docstrings and docstr:
+                output.append(
+                    "def {function}({args}) -> {ret}:\n\"\"\"{docstr}\"\"\"\n...".format(
+                        function=name,
+                        args=", ".join(sig),
+                        ret=strip_or_import(signature.ret_type, module, imports),
+                        docstr=docstr,
+                    )
+                )
+            else:
+                output.append(
+                    "def {function}({args}) -> {ret}: ...".format(
+                        function=name,
+                        args=", ".join(sig),
+                        ret=strip_or_import(signature.ret_type, module, imports),
+                    )
                 )
-            )
 
 
 def strip_or_import(typ: str, module: ModuleType, imports: List[str]) -> str:
@@ -339,6 +351,7 @@ def generate_c_type_stub(
     imports: List[str],
     sigs: Optional[Dict[str, str]] = None,
     class_sigs: Optional[Dict[str, str]] = None,
+    include_docstrings: bool = False
 ) -> None:
     """Generate stub for a single class using runtime introspection.
 
@@ -382,6 +395,7 @@ def generate_c_type_stub(
                     sigs=sigs,
                     class_name=class_name,
                     class_sigs=class_sigs,
+                    include_docstrings=include_docstrings,
                 )
         elif is_c_property(value):
             done.add(attr)
@@ -397,7 +411,7 @@ def generate_c_type_stub(
             )
         elif is_c_type(value):
             generate_c_type_stub(
-                module, attr, value, types, imports=imports, sigs=sigs, class_sigs=class_sigs
+                module, attr, value, types, imports=imports, sigs=sigs, class_sigs=class_sigs, include_docstrings=include_docstrings
             )
             done.add(attr)
 

test-data/unit/stubgen.test

@@ -2688,3 +2688,20 @@ def f():
     return 0
 [out]
 def f(): ...
+
+[case testIncludeDocstrings]
+# flags:  --include-docstrings
+class A:
+    """class docstring"""
+    def func():
+        """func docstring"""
+        ...
+    def nodoc():
+        ...
+[out]
+class A:
+    """class docstring"""
+    def func() -> None:
+        """func docstring"""
+        ...
+    def nodoc() -> None: ...