From e4de4e32eeb87c33a35b5abe88676cf3f34f388f Mon Sep 17 00:00:00 2001 From: Eric Mark Martin Date: Sat, 29 Jun 2024 19:32:51 -0400 Subject: [PATCH] Include keyword only args when generating signatures in stubgenc (#17448) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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> --- mypy/stubgenc.py | 56 +++++++++++++++++++++++++++---------- mypy/test/teststubgen.py | 29 +++++++++++++++++++ test-data/unit/stubgen.test | 9 ++++++ 3 files changed, 79 insertions(+), 15 deletions(-) diff --git a/mypy/stubgenc.py b/mypy/stubgenc.py index 0aa6088a4e02..7ab500b4fe12 100755 --- 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 index 05a3809179bd..e65a16c8f395 100644 --- 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 index 5dcb0706a8cb..94d0edb2ae37 100644 --- 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