From 53134979cd4442b434637eef81557d4c0b2a5d00 Mon Sep 17 00:00:00 2001 From: Ali Hamdan Date: Mon, 28 Oct 2024 07:50:25 +0100 Subject: [PATCH] stubgen: add support for PEPs 695 and 696 syntax (#18054) closes #17997 --- mypy/stubdoc.py | 5 +- mypy/stubgen.py | 20 +++++++- mypy/stubutil.py | 26 +++++++++++ test-data/unit/stubgen.test | 93 +++++++++++++++++++++++++++++++++++++ 4 files changed, 139 insertions(+), 5 deletions(-) diff --git a/mypy/stubdoc.py b/mypy/stubdoc.py index 928d024514f3..434de0ea3bcb 100644 --- a/mypy/stubdoc.py +++ b/mypy/stubdoc.py @@ -76,6 +76,7 @@ class FunctionSig(NamedTuple): name: str args: list[ArgSig] ret_type: str | None + type_args: str = "" # TODO implement in stubgenc and remove the default def is_special_method(self) -> bool: return bool( @@ -141,9 +142,7 @@ 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 = f"{indent}{prefix}def {self.name}{self.type_args}({', '.join(args)}){retfield}:" if docstring: suffix = f"\n{indent} {mypy.util.quote_docstring(docstring)}" else: diff --git a/mypy/stubgen.py b/mypy/stubgen.py index 02c0c1e58ab5..684e85d58758 100755 --- a/mypy/stubgen.py +++ b/mypy/stubgen.py @@ -106,6 +106,7 @@ StrExpr, TempNode, TupleExpr, + TypeAliasStmt, TypeInfo, UnaryExpr, Var, @@ -398,6 +399,9 @@ def visit_assignment_stmt(self, o: AssignmentStmt) -> None: for name in get_assigned_names(o.lvalues): self.names.add(name) + def visit_type_alias_stmt(self, o: TypeAliasStmt) -> None: + self.names.add(o.name.name) + def find_referenced_names(file: MypyFile) -> set[str]: finder = ReferenceFinder() @@ -507,7 +511,8 @@ def visit_overloaded_func_def(self, o: OverloadedFuncDef) -> None: def get_default_function_sig(self, func_def: FuncDef, ctx: FunctionContext) -> FunctionSig: args = self._get_func_args(func_def, ctx) retname = self._get_func_return(func_def, ctx) - return FunctionSig(func_def.name, args, retname) + type_args = self.format_type_args(func_def) + return FunctionSig(func_def.name, args, retname, type_args) def _get_func_args(self, o: FuncDef, ctx: FunctionContext) -> list[ArgSig]: args: list[ArgSig] = [] @@ -765,7 +770,8 @@ def visit_class_def(self, o: ClassDef) -> None: self.import_tracker.add_import("abc") self.import_tracker.require_name("abc") bases = f"({', '.join(base_types)})" if base_types else "" - self.add(f"{self._indent}class {o.name}{bases}:\n") + type_args = self.format_type_args(o) + self.add(f"{self._indent}class {o.name}{type_args}{bases}:\n") self.indent() if self._include_docstrings and o.docstring: docstring = mypy.util.quote_docstring(o.docstring) @@ -1101,6 +1107,16 @@ def process_typealias(self, lvalue: NameExpr, rvalue: Expression) -> None: self.record_name(lvalue.name) self._vars[-1].append(lvalue.name) + def visit_type_alias_stmt(self, o: TypeAliasStmt) -> None: + """Type aliases defined with the `type` keyword (PEP 695).""" + p = AliasPrinter(self) + name = o.name.name + rvalue = o.value.expr() + type_args = self.format_type_args(o) + self.add(f"{self._indent}type {name}{type_args} = {rvalue.accept(p)}\n") + self.record_name(name) + self._vars[-1].append(name) + def visit_if_stmt(self, o: IfStmt) -> None: # Ignore if __name__ == '__main__'. expr = o.expr[0] diff --git a/mypy/stubutil.py b/mypy/stubutil.py index 04b36e149957..e824c3f1a8e2 100644 --- a/mypy/stubutil.py +++ b/mypy/stubutil.py @@ -16,6 +16,7 @@ import mypy.options from mypy.modulefinder import ModuleNotFoundReason from mypy.moduleinspect import InspectError, ModuleInspect +from mypy.nodes import PARAM_SPEC_KIND, TYPE_VAR_TUPLE_KIND, ClassDef, FuncDef, TypeAliasStmt from mypy.stubdoc import ArgSig, FunctionSig from mypy.types import AnyType, NoneType, Type, TypeList, TypeStrVisitor, UnboundType, UnionType @@ -777,6 +778,31 @@ def format_func_def( ) return lines + def format_type_args(self, o: TypeAliasStmt | FuncDef | ClassDef) -> str: + if not o.type_args: + return "" + p = AnnotationPrinter(self) + type_args_list: list[str] = [] + for type_arg in o.type_args: + if type_arg.kind == PARAM_SPEC_KIND: + prefix = "**" + elif type_arg.kind == TYPE_VAR_TUPLE_KIND: + prefix = "*" + else: + prefix = "" + if type_arg.upper_bound: + bound_or_values = f": {type_arg.upper_bound.accept(p)}" + elif type_arg.values: + bound_or_values = f": ({', '.join(v.accept(p) for v in type_arg.values)})" + else: + bound_or_values = "" + if type_arg.default: + default = f" = {type_arg.default.accept(p)}" + else: + default = "" + type_args_list.append(f"{prefix}{type_arg.name}{bound_or_values}{default}") + return "[" + ", ".join(type_args_list) + "]" + def print_annotation( self, t: Type, diff --git a/test-data/unit/stubgen.test b/test-data/unit/stubgen.test index 69781e9d2143..e64c9c66d65d 100644 --- a/test-data/unit/stubgen.test +++ b/test-data/unit/stubgen.test @@ -4415,3 +4415,96 @@ class Test(Whatever, a=1, b='b', c=True, d=1.5, e=None, f=1j, g=b'123'): ... class Test(Whatever, keyword=SomeName * 2, attr=SomeName.attr): ... [out] class Test(Whatever, keyword=SomeName * 2, attr=SomeName.attr): ... + +[case testPEP695GenericClass] +# flags: --python-version=3.12 + +class C[T]: ... +class C1[T1](int): ... +class C2[T2: int]: ... +class C3[T3: str | bytes]: ... +class C4[T4: (str, bytes)]: ... + +class Outer: + class Inner[T]: ... + +[out] +class C[T]: ... +class C1[T1](int): ... +class C2[T2: int]: ... +class C3[T3: str | bytes]: ... +class C4[T4: (str, bytes)]: ... + +class Outer: + class Inner[T]: ... + +[case testPEP695GenericFunction] +# flags: --python-version=3.12 + +def f1[T1](): ... +def f2[T2: int](): ... +def f3[T3: str | bytes](): ... +def f4[T4: (str, bytes)](): ... + +class C: + def m[T](self, x: T) -> T: ... + +[out] +def f1[T1]() -> None: ... +def f2[T2: int]() -> None: ... +def f3[T3: str | bytes]() -> None: ... +def f4[T4: (str, bytes)]() -> None: ... + +class C: + def m[T](self, x: T) -> T: ... + +[case testPEP695TypeAlias] +# flags: --python-version=3.12 + +type Alias = int | str +type Alias1[T1] = list[T1] | set[T1] +type Alias2[T2: int] = list[T2] | set[T2] +type Alias3[T3: str | bytes] = list[T3] | set[T3] +type Alias4[T4: (str, bytes)] = list[T4] | set[T4] + +class C: + type IndentedAlias[T] = list[T] + +[out] +type Alias = int | str +type Alias1[T1] = list[T1] | set[T1] +type Alias2[T2: int] = list[T2] | set[T2] +type Alias3[T3: str | bytes] = list[T3] | set[T3] +type Alias4[T4: (str, bytes)] = list[T4] | set[T4] +class C: + type IndentedAlias[T] = list[T] + +[case testPEP695Syntax_semanal] +# flags: --python-version=3.12 + +class C[T]: ... +def f[S](): ... +type A[R] = list[R] + +[out] +class C[T]: ... + +def f[S]() -> None: ... +type A[R] = list[R] + +[case testPEP696Syntax] +# flags: --python-version=3.13 + +type Alias1[T1 = int] = list[T1] | set[T1] +type Alias2[T2: int | float = int] = list[T2] | set[T2] +class C3[T3 = int]: ... +class C4[T4: int | float = int](list[T4]): ... +def f5[T5 = int](): ... + +[out] +type Alias1[T1 = int] = list[T1] | set[T1] +type Alias2[T2: int | float = int] = list[T2] | set[T2] +class C3[T3 = int]: ... +class C4[T4: int | float = int](list[T4]): ... + +def f5[T5 = int]() -> None: ...