From 1a197f7e8206dc963ac5ce11b06506a32869ae5c Mon Sep 17 00:00:00 2001 From: Michael Carlstrom Date: Mon, 22 Jul 2024 11:34:09 -0400 Subject: [PATCH] added support for inline functions --- mypy/stubdoc.py | 5 +-- mypy/stubgenc.py | 8 ++++- mypy/stubutil.py | 5 +++ mypy/test/teststubgen.py | 67 +++++++++++++++++++++++++++------------- 4 files changed, 60 insertions(+), 25 deletions(-) diff --git a/mypy/stubdoc.py b/mypy/stubdoc.py index 928d024514f3..f305715c4c25 100644 --- a/mypy/stubdoc.py +++ b/mypy/stubdoc.py @@ -108,6 +108,7 @@ def format_sig( is_async: bool = False, any_val: str | None = None, docstring: str | None = None, + generic: str = "" ) -> str: args: list[str] = [] for arg in self.args: @@ -141,8 +142,8 @@ def format_sig( retfield = " -> " + ret_type prefix = "async " if is_async else "" - sig = "{indent}{prefix}def {name}({args}){ret}:".format( - indent=indent, prefix=prefix, name=self.name, args=", ".join(args), ret=retfield + sig = "{indent}{prefix}def {name}{generic}({args}){ret}:".format( + indent=indent, prefix=prefix, name=self.name, args=", ".join(args), ret=retfield, generic=generic ) if docstring: suffix = f"\n{indent} {mypy.util.quote_docstring(docstring)}" diff --git a/mypy/stubgenc.py b/mypy/stubgenc.py index 68b81a25aa22..09f709b8057c 100755 --- a/mypy/stubgenc.py +++ b/mypy/stubgenc.py @@ -647,7 +647,13 @@ def generate_function_stub( if docstring: docstring = self._indent_docstring(docstring) - output.extend(self.format_func_def(inferred, decorators=decorators, docstring=docstring)) + + if hasattr(obj, "__type_params__"): + generic = generate_inline_generic(obj.__type_params__) + else: + generic = "" + + output.extend(self.format_func_def(inferred, decorators=decorators, docstring=docstring, generic=generic)) self._fix_iter(ctx, inferred, output) def _indent_docstring(self, docstring: str) -> str: diff --git a/mypy/stubutil.py b/mypy/stubutil.py index c24960b87328..fe63acbe2421 100644 --- a/mypy/stubutil.py +++ b/mypy/stubutil.py @@ -766,6 +766,7 @@ def format_func_def( is_coroutine: bool = False, decorators: list[str] | None = None, docstring: str | None = None, + generic: str = "", ) -> list[str]: lines: list[str] = [] if decorators is None: @@ -781,6 +782,7 @@ def format_func_def( indent=self._indent, is_async=is_coroutine, docstring=docstring if self._include_docstrings else None, + generic=generic ) ) return lines @@ -847,6 +849,9 @@ def should_reexport(self, name: str, full_module: str, name_is_alias: bool) -> b def generate_inline_generic(type_params: tuple[TypeVar | ParamSpec | TypeVarTuple, ...]) -> str: """Generate stub for inline generic from __type_params__""" + if len(type_params) == 0: + return "" + generic_arg_list: list[str] = [] for type_param in type_params: diff --git a/mypy/test/teststubgen.py b/mypy/test/teststubgen.py index 8aa2edbb6b98..75af33d8d213 100644 --- a/mypy/test/teststubgen.py +++ b/mypy/test/teststubgen.py @@ -614,6 +614,29 @@ def test_common_dir_prefix_win(self) -> None: assert common_dir_prefix([r"foo\bar/x.pyi"]) == r"foo\bar" assert common_dir_prefix([r"foo/bar/x.pyi"]) == r"foo\bar" + def test_generate_inline_generic(self) -> None: + assert generate_inline_generic(()) == "" + T = TypeVar("T") + assert generate_inline_generic((T,)) == "[T]" + TBound = TypeVar("TBound", bound=int) + assert generate_inline_generic((TBound,)) == "[TBound: int]" + TBoundTuple = TypeVar("TBoundTuple", int, str) + assert generate_inline_generic((TBoundTuple,)) == "[TBoundTuple: (int, str)]" + P = ParamSpec("P") + p_tuple = cast(tuple[ParamSpec], (P,)) + assert generate_inline_generic(p_tuple) == "[**P]" + U = TypeVarTuple("U") + u_tuple = cast(tuple[TypeVarTuple], (U,)) + assert generate_inline_generic(u_tuple) == "[*U]" + all_tuple = cast( + tuple[TypeVar, TypeVar, TypeVar, ParamSpec, TypeVarTuple], + (T, TBound, TBoundTuple, P, U), + ) + assert ( + generate_inline_generic(all_tuple) + == "[T, TBound: int, TBoundTuple: (int, str), **P, *U]" + ) + class StubgenHelpersSuite(unittest.TestCase): def test_is_blacklisted_path(self) -> None: @@ -826,28 +849,6 @@ def test_infer_cast_sig(self) -> None: for op in ("float", "bool", "bytes", "int"): assert_equal(infer_method_ret_type(f"__{op}__"), op) - def test_generate_inline_generic(self) -> None: - T = TypeVar("T") - assert generate_inline_generic((T,)) == "[T]" - TBound = TypeVar("TBound", bound=int) - assert generate_inline_generic((TBound,)) == "[TBound: int]" - TBoundTuple = TypeVar("TBoundTuple", int, str) - assert generate_inline_generic((TBoundTuple,)) == "[TBoundTuple: (int, str)]" - P = ParamSpec("P") - p_tuple = cast(tuple[ParamSpec], (P,)) - assert generate_inline_generic(p_tuple) == "[**P]" - U = TypeVarTuple("U") - u_tuple = cast(tuple[TypeVarTuple], (U,)) - assert generate_inline_generic(u_tuple) == "[*U]" - all_tuple = cast( - tuple[TypeVar, TypeVar, TypeVar, ParamSpec, TypeVarTuple], - (T, TBound, TBoundTuple, P, U), - ) - assert ( - generate_inline_generic(all_tuple) - == "[T, TBound: int, TBoundTuple: (int, str), **P, *U]" - ) - def test_generate_class_stub_no_crash_for_object(self) -> None: output: list[str] = [] mod = ModuleType("module", "") # any module is fine @@ -942,6 +943,28 @@ class TestClass: gen.generate_class_stub("C", TestClass, output) assert_equal(output, ["class C[T]: ..."]) + def test_inline_generic_function(self) -> None: + T = TypeVar("T", bound=int) + class TestClass: + def test(self, arg0: T) -> T: + """ + test(self, arg0: T) -> T + """ + return arg0 + + test.__type_params__ = (T, ) + + output: list[str] = [] + mod = ModuleType(TestClass.__module__, "") + gen = InspectionStubGenerator(mod.__name__, known_modules=[mod.__name__], module=mod) + gen.generate_function_stub( + "test", + TestClass.test, + output=output, + class_info=ClassInfo(self_var="self", cls=TestClass, name="TestClass"), + ) + assert_equal(output, ["def test[T: int](self, arg0: T) -> T: ..."]) + def test_generate_c_type_inheritance_builtin_type(self) -> None: class TestClass(type): pass