Include keyword only args when generating signatures in stubgenc (python#17448)

ericmarkmartin · pre-commit-ci[bot] · web-flow · commit e4de4e32eeb8 · 2024-06-30T00:32:51.000+01:00
Currently, signatures generated for callable by the `InspectionStubGenerator` won’t include keywords only arguments or their defaults. This change includes them in the generated signatures.  --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
diff --git a/mypy/stubgenc.py b/mypy/stubgenc.py
@@ -12,7 +12,7 @@
 import keyword
 import os.path
 from types import FunctionType, ModuleType
-from typing import Any, Mapping
+from typing import Any, Callable, Mapping
 
 from mypy.fastparse import parse_type_comment
 from mypy.moduleinspect import is_c_module
@@ -292,6 +292,8 @@ def get_default_function_sig(self, func: object, ctx: FunctionContext) -> Functi
         varargs = argspec.varargs
         kwargs = argspec.varkw
         annotations = argspec.annotations
+        kwonlyargs = argspec.kwonlyargs
+        kwonlydefaults = argspec.kwonlydefaults
 
         def get_annotation(key: str) -> str | None:
             if key not in annotations:
@@ -304,27 +306,51 @@ def get_annotation(key: str) -> str | None:
             return argtype
 
         arglist: list[ArgSig] = []
+
         # Add the arguments to the signature
-        for i, arg in enumerate(args):
-            # Check if the argument has a default value
-            if defaults and i >= len(args) - len(defaults):
-                default_value = defaults[i - (len(args) - len(defaults))]
-                if arg in annotations:
-                    argtype = annotations[arg]
+        def add_args(
+            args: list[str], get_default_value: Callable[[int, str], object | None]
+        ) -> None:
+            for i, arg in enumerate(args):
+                # Check if the argument has a default value
+                default_value = get_default_value(i, arg)
+                if default_value is not None:
+                    if arg in annotations:
+                        argtype = annotations[arg]
+                    else:
+                        argtype = self.get_type_annotation(default_value)
+                        if argtype == "None":
+                            # None is not a useful annotation, but we can infer that the arg
+                            # is optional
+                            incomplete = self.add_name("_typeshed.Incomplete")
+                            argtype = f"{incomplete} | None"
+
+                    arglist.append(ArgSig(arg, argtype, default=True))
                 else:
-                    argtype = self.get_type_annotation(default_value)
-                    if argtype == "None":
-                        # None is not a useful annotation, but we can infer that the arg
-                        # is optional
-                        incomplete = self.add_name("_typeshed.Incomplete")
-                        argtype = f"{incomplete} | None"
-                arglist.append(ArgSig(arg, argtype, default=True))
+                    arglist.append(ArgSig(arg, get_annotation(arg), default=False))
+
+        def get_pos_default(i: int, _arg: str) -> Any | None:
+            if defaults and i >= len(args) - len(defaults):
+                return defaults[i - (len(args) - len(defaults))]
             else:
-                arglist.append(ArgSig(arg, get_annotation(arg), default=False))
+                return None
+
+        add_args(args, get_pos_default)
 
         # Add *args if present
         if varargs:
             arglist.append(ArgSig(f"*{varargs}", get_annotation(varargs)))
+        # if we have keyword only args, then wee need to add "*"
+        elif kwonlyargs:
+            arglist.append(ArgSig("*"))
+
+        def get_kw_default(_i: int, arg: str) -> Any | None:
+            if kwonlydefaults:
+                return kwonlydefaults.get(arg)
+            else:
+                return None
+
+        add_args(kwonlyargs, get_kw_default)
 
         # Add **kwargs if present
         if kwargs:
diff --git a/mypy/test/teststubgen.py b/mypy/test/teststubgen.py
@@ -845,6 +845,35 @@ class TestClassVariableCls:
         assert_equal(gen.get_imports().splitlines(), ["from typing import ClassVar"])
         assert_equal(output, ["class C:", "    x: ClassVar[int] = ..."])
 
+    def test_non_c_generate_signature_with_kw_only_args(self) -> None:
+        class TestClass:
+            def test(
+                self, arg0: str, *, keyword_only: str, keyword_only_with_default: int = 7
+            ) -> None:
+                pass
+
+        output: list[str] = []
+        mod = ModuleType(TestClass.__module__, "")
+        gen = InspectionStubGenerator(mod.__name__, known_modules=[mod.__name__], module=mod)
+        gen.is_c_module = False
+        gen.generate_function_stub(
+            "test",
+            TestClass.test,
+            output=output,
+            class_info=ClassInfo(
+                self_var="self",
+                cls=TestClass,
+                name="TestClass",
+                docstring=getattr(TestClass, "__doc__", None),
+            ),
+        )
+        assert_equal(
+            output,
+            [
+                "def test(self, arg0: str, *, keyword_only: str, keyword_only_with_default: int = ...) -> None: ..."
+            ],
+        )
+
     def test_generate_c_type_inheritance(self) -> None:
         class TestClass(KeyError):
             pass
diff --git a/test-data/unit/stubgen.test b/test-data/unit/stubgen.test
@@ -361,6 +361,15 @@ def g(x, *, y=1, z=2): ...
 def f(x, *, y: int = 1) -> None: ...
 def g(x, *, y: int = 1, z: int = 2) -> None: ...
 
+[case testKeywordOnlyArg_inspect]
+def f(x, *, y=1): ...
+def g(x, *, y=1, z=2): ...
+def h(x, *, y, z=2): ...
+[out]
+def f(x, *, y: int = ...): ...
+def g(x, *, y: int = ..., z: int = ...): ...
+def h(x, *, y, z: int = ...): ...
+
 [case testProperty]
 class A:
     @property
