Skip to content

Commit

Permalink
InspectionStubGenerator respects kw only args
Browse files Browse the repository at this point in the history
  • Loading branch information
ericmarkmartin committed Jun 28, 2024
1 parent de4e9d6 commit 0f22268
Show file tree
Hide file tree
Showing 2 changed files with 66 additions and 15 deletions.
53 changes: 38 additions & 15 deletions mypy/stubgenc.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -291,6 +291,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:
Expand All @@ -303,27 +305,48 @@ 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, get_default_value: Callable[[int, str], object | None]):
for i, arg in enumerate(args):
# Check if the argument has a default value
if default_value := get_default_value(i, arg):
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) -> object | 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) -> object | None:
if kwonlydefaults:
return kwonlydefaults.get(arg)
else:
return None

add_args(kwonlyargs, get_kw_default)

# Add **kwargs if present
if kwargs:
Expand Down
28 changes: 28 additions & 0 deletions mypy/test/teststubgen.py
Original file line number Diff line number Diff line change
Expand Up @@ -845,6 +845,34 @@ 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, *, keyword_only : str, keyword_only_with_default : int = 7):
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, *, keyword_only: str, keyword_only_with_default: int = ...): ..."
],
)

def test_generate_c_type_inheritance(self) -> None:
class TestClass(KeyError):
pass
Expand Down

0 comments on commit 0f22268

Please sign in to comment.