stubgenc: add support for including class and property docstrings

Author: chadrik

mypy/stubdoc.py

@@ -78,6 +78,7 @@ class FunctionSig(NamedTuple):
     args: list[ArgSig]
     ret_type: str | None
     type_args: str = ""  # TODO implement in stubgenc and remove the default
+    docstring: str | None = None
 
     def is_special_method(self) -> bool:
         return bool(
@@ -110,6 +111,7 @@ def format_sig(
         is_async: bool = False,
         any_val: str | None = None,
         docstring: str | None = None,
+        include_docstrings: bool = False,
     ) -> str:
         args: list[str] = []
         for arg in self.args:
@@ -144,8 +146,10 @@ def format_sig(
 
         prefix = "async " if is_async else ""
         sig = f"{indent}{prefix}def {self.name}{self.type_args}({', '.join(args)}){retfield}:"
-        if docstring:
-            suffix = f"\n{indent}    {mypy.util.quote_docstring(docstring)}"
+        # if an overload has a docstring use it, otherwise fall back to the generic docstring
+        doc = (self.docstring or docstring) if include_docstrings else None
+        if doc:
+            suffix = f"\n{indent}    {mypy.util.quote_docstring(doc)}"
         else:
             suffix = " ..."
         return f"{sig}{suffix}"

mypy/stubgenc.py

@@ -37,6 +37,7 @@
     infer_method_arg_types,
     infer_method_ret_type,
 )
+from mypy.util import quote_docstring
 
 
 class ExternalSignatureGenerator(SignatureGenerator):
@@ -645,8 +646,7 @@ def generate_function_stub(
                 if inferred[0].args and inferred[0].args[0].name == "cls":
                     decorators.append("@classmethod")
 
-        if docstring:
-            docstring = self._indent_docstring(docstring)
+        docstring = self._indent_docstring(ctx.docstring) if ctx.docstring else None
         output.extend(self.format_func_def(inferred, decorators=decorators, docstring=docstring))
         self._fix_iter(ctx, inferred, output)
 
@@ -750,9 +750,16 @@ def generate_property_stub(
             )
         else:  # regular property
             if readonly:
+                docstring = self._indent_docstring(ctx.docstring) if ctx.docstring else None
                 ro_properties.append(f"{self._indent}@property")
                 sig = FunctionSig(name, [ArgSig("self")], inferred_type)
-                ro_properties.append(sig.format_sig(indent=self._indent))
+                ro_properties.append(
+                    sig.format_sig(
+                        indent=self._indent,
+                        docstring=docstring,
+                        include_docstrings=self._include_docstrings,
+                    )
+                )
             else:
                 if inferred_type is None:
                     inferred_type = self.add_name("_typeshed.Incomplete")
@@ -867,8 +874,16 @@ def generate_class_stub(
             bases_str = "(%s)" % ", ".join(bases)
         else:
             bases_str = ""
-        if types or static_properties or rw_properties or methods or ro_properties:
+
+        if class_info.docstring and self._include_docstrings:
+            doc = f"    {self._indent}{quote_docstring(class_info.docstring)}"
+            docstring = doc.splitlines(keepends=False)
+        else:
+            docstring = []
+
+        if docstring or types or static_properties or rw_properties or methods or ro_properties:
             output.append(f"{self._indent}class {class_name}{bases_str}:")
+            output.extend(docstring)
             for line in types:
                 if (
                     output
@@ -878,14 +893,10 @@ def generate_class_stub(
                 ):
                     output.append("")
                 output.append(line)
-            for line in static_properties:
-                output.append(line)
-            for line in rw_properties:
-                output.append(line)
-            for line in methods:
-                output.append(line)
-            for line in ro_properties:
-                output.append(line)
+            output.extend(static_properties)
+            output.extend(rw_properties)
+            output.extend(methods)
+            output.extend(ro_properties)
         else:
             output.append(f"{self._indent}class {class_name}{bases_str}: ...")
 

mypy/stubutil.py

@@ -803,7 +803,8 @@ def format_func_def(
                 signature.format_sig(
                     indent=self._indent,
                     is_async=is_coroutine,
-                    docstring=docstring if self._include_docstrings else None,
+                    docstring=docstring,
+                    include_docstrings=self._include_docstrings,
                 )
             )
         return lines