Check for docstrings as well

scripts/ci/python_check_signatures.py

@@ -8,19 +8,64 @@
 from __future__ import annotations
 
 import ast
+import difflib
 import importlib
 import inspect
 import sys
+import textwrap
 from inspect import Parameter, Signature
 from pathlib import Path
 from typing import Any
 
 import parso
+from colorama import Fore, Style, init as colorama_init
 
-TotalSignature = dict[str, Signature | dict[str, Signature]]
+colorama_init()
 
 
-def parse_function_signature(node: Any) -> Signature:
+def print_colored_diff(runtime, stub):
+    # Split the strings into lines
+    runtime_lines = runtime.splitlines()
+    stub_lines = stub.splitlines()
+
+    # Generate the diff
+    diff = difflib.unified_diff(runtime_lines, stub_lines, fromfile="runtime", tofile="stub")
+
+    # Print the diff output with colored lines
+    for line in diff:
+        if line.startswith("+"):
+            print(Fore.GREEN + line + Style.RESET_ALL)
+        elif line.startswith("-"):
+            print(Fore.RED + line + Style.RESET_ALL)
+        elif line.startswith("?"):
+            print(Fore.YELLOW + line + Style.RESET_ALL)
+        else:
+            print(line)
+
+
+class APIDef:
+    def __init__(self, name: str, signature: Signature, doc: str | None):
+        self.name = name
+        self.signature = signature
+        self.doc = inspect.cleandoc(doc) if doc else None
+
+    def __str__(self):
+        doclines = (self.doc or "").split("\n")
+        if len(doclines) == 1:
+            docstring = f'"""{doclines[0]}"""'
+        else:
+            docstring = '"""\n' + "\n".join(doclines) + '\n"""'
+        docstring = textwrap.indent(docstring, "    ")
+        return f"{self.name}{self.signature}:\n{docstring}"
+
+    def __eq__(self, other):
+        return self.name == other.name and self.signature == other.signature and self.doc == other.doc
+
+
+TotalSignature = dict[str, APIDef | dict[str, APIDef]]
+
+
+def parse_function_signature(node: Any) -> APIDef:
     """Convert a parso function definition node into a Python inspect.Signature object."""
     params = []
 
@@ -53,7 +98,15 @@ def parse_function_signature(node: Any) -> Signature:
 
         params.append(Parameter(name=param_name, kind=kind, default=default))
 
-    return Signature(parameters=params)
+    doc = None
+    for child in node.children:
+        if child.type == "suite":
+            first_child = child.children[1]
+            if first_child.type == "simple_stmt" and first_child.children[0].type == "string":
+                doc = first_child.children[0].value.strip('"""')
+
+    sig = Signature(parameters=params)
+    return APIDef(node.name.value, sig, doc)
 
 
 def load_stub_signatures(pyi_file: Path) -> TotalSignature:
@@ -94,9 +147,11 @@ def load_runtime_signatures(module_name: str) -> TotalSignature:
     # Get top-level functions and classes
     for name, obj in inspect.getmembers(module):
         if inspect.isfunction(obj):
-            signatures[name] = inspect.signature(obj)
+            api_def = APIDef(name, inspect.signature(obj), obj.__doc__)
+            signatures[name] = api_def
         elif inspect.isbuiltin(obj):
-            signatures[name] = inspect.signature(obj)
+            api_def = APIDef(name, inspect.signature(obj), obj.__doc__)
+            signatures[name] = api_def
         elif inspect.isclass(obj):
             class_def = {}
             # Get methods within the class
@@ -109,14 +164,20 @@ def load_runtime_signatures(module_name: str) -> TotalSignature:
                     class_def[method_name] = parse_function_signature(parsed)
                     continue
                 try:
-                    class_def[method_name] = inspect.signature(method_obj)
+                    api_def = APIDef(method_name, inspect.signature(method_obj), method_obj.__doc__)
+                    class_def[method_name] = api_def
                 except Exception:
                     pass
             # Get property getters
             for method_name, method_obj in inspect.getmembers(
                 obj, lambda o: o.__class__.__name__ == "getset_descriptor"
             ):
-                class_def[method_name] = Signature(parameters=[Parameter("self", Parameter.POSITIONAL_ONLY)])
+                api_def = APIDef(
+                    method_name,
+                    Signature(parameters=[Parameter("self", Parameter.POSITIONAL_ONLY)]),
+                    method_obj.__doc__,
+                )
+                class_def[method_name] = api_def
             signatures[name] = class_def
 
     return signatures
@@ -132,32 +193,40 @@ def compare_signatures(stub_signatures: TotalSignature, runtime_signatures: Tota
             if name in runtime_signatures:
                 runtime_class_signature = runtime_signatures.get(name)
                 if not isinstance(runtime_class_signature, dict):
+                    print()
                     print(f"{name} signature mismatch:")
-                    print(f"  Stub: {stub_signature}")
-                    print(f"  Runtime: {runtime_class_signature}")
-                    result += 1
+                    print("Stub expected class, but runtime provided function.")
                     continue
                 for method_name, stub_method_signature in stub_signature.items():
+                    if stub_method_signature.doc is None:
+                        print()
+                        print(f"{name}.{method_name} missing docstring")
+                        result += 1
                     runtime_method_signature = runtime_class_signature.get(method_name)
                     if runtime_method_signature != stub_method_signature:
+                        print()
                         print(f"{name}.{method_name}(…) signature mismatch:")
-                        print(f"    Stub: {stub_method_signature}")
-                        print(f"    Runtime: {runtime_method_signature}")
+                        print_colored_diff(str(runtime_method_signature), str(stub_method_signature))
                         result += 1
 
             else:
                 print(f"Class {name} not found in runtime")
                 result += 1
         else:
+            if stub_signature.doc is None:
+                print()
+                print(f"{name} missing docstring")
+                result += 1
             if name in runtime_signatures:
                 # Handle top-level functions
                 runtime_signature = runtime_signatures.get(name)
                 if runtime_signature != stub_signature:
+                    print()
                     print(f"{name}(…) signature mismatch:")
-                    print(f"  Stub: {stub_signature}")
-                    print(f"  Runtime: {runtime_signature}")
+                    print_colored_diff(str(runtime_signature), str(stub_signature))
                     result += 1
             else:
+                print()
                 print(f"Function {name} not found in runtime")
                 result += 1