From 72c413d2352da5ce1433ef241faca8f40fa1fe27 Mon Sep 17 00:00:00 2001 From: Sebastian Rittau Date: Thu, 12 Sep 2024 16:26:36 +0200 Subject: [PATCH] stubgen: Use `Generator` type var defaults (#17670) Fixes #17669 --- mypy/stubgen.py | 13 ++++++++++--- test-data/unit/stubgen.test | 34 +++++++++++++++++----------------- 2 files changed, 27 insertions(+), 20 deletions(-) diff --git a/mypy/stubgen.py b/mypy/stubgen.py index 14417f55545e..02c0c1e58ab5 100755 --- a/mypy/stubgen.py +++ b/mypy/stubgen.py @@ -588,8 +588,8 @@ def _get_func_return(self, o: FuncDef, ctx: FunctionContext) -> str | None: if has_yield_expression(o) or has_yield_from_expression(o): generator_name = self.add_name("collections.abc.Generator") yield_name = "None" - send_name = "None" - return_name = "None" + send_name: str | None = None + return_name: str | None = None if has_yield_from_expression(o): yield_name = send_name = self.add_name("_typeshed.Incomplete") else: @@ -600,7 +600,14 @@ def _get_func_return(self, o: FuncDef, ctx: FunctionContext) -> str | None: send_name = self.add_name("_typeshed.Incomplete") if has_return_statement(o): return_name = self.add_name("_typeshed.Incomplete") - return f"{generator_name}[{yield_name}, {send_name}, {return_name}]" + if return_name is not None: + if send_name is None: + send_name = "None" + return f"{generator_name}[{yield_name}, {send_name}, {return_name}]" + elif send_name is not None: + return f"{generator_name}[{yield_name}, {send_name}]" + else: + return f"{generator_name}[{yield_name}]" if not has_return_statement(o) and o.abstract_status == NOT_ABSTRACT: return "None" return None diff --git a/test-data/unit/stubgen.test b/test-data/unit/stubgen.test index fe0538159aa3..69781e9d2143 100644 --- a/test-data/unit/stubgen.test +++ b/test-data/unit/stubgen.test @@ -1640,11 +1640,11 @@ def all(): from _typeshed import Incomplete from collections.abc import Generator -def f() -> Generator[Incomplete, None, None]: ... -def g() -> Generator[None, Incomplete, None]: ... -def h1() -> Generator[None, None, None]: ... +def f() -> Generator[Incomplete]: ... +def g() -> Generator[None, Incomplete]: ... +def h1() -> Generator[None]: ... def h2() -> Generator[None, None, Incomplete]: ... -def h3() -> Generator[None, None, None]: ... +def h3() -> Generator[None]: ... def all() -> Generator[Incomplete, Incomplete, Incomplete]: ... [case testFunctionYieldsNone] @@ -1656,8 +1656,8 @@ def g(): [out] from collections.abc import Generator -def f() -> Generator[None, None, None]: ... -def g() -> Generator[None, None, None]: ... +def f() -> Generator[None]: ... +def g() -> Generator[None]: ... [case testGeneratorAlreadyDefined] class Generator: @@ -1671,7 +1671,7 @@ from collections.abc import Generator as _Generator class Generator: ... -def f() -> _Generator[Incomplete, None, None]: ... +def f() -> _Generator[Incomplete]: ... [case testGeneratorYieldFrom] def g1(): @@ -1692,10 +1692,10 @@ def g5(): from _typeshed import Incomplete from collections.abc import Generator -def g1() -> Generator[Incomplete, Incomplete, None]: ... -def g2() -> Generator[Incomplete, Incomplete, None]: ... -def g3() -> Generator[Incomplete, Incomplete, None]: ... -def g4() -> Generator[Incomplete, Incomplete, None]: ... +def g1() -> Generator[Incomplete, Incomplete]: ... +def g2() -> Generator[Incomplete, Incomplete]: ... +def g3() -> Generator[Incomplete, Incomplete]: ... +def g4() -> Generator[Incomplete, Incomplete]: ... def g5() -> Generator[Incomplete, Incomplete, Incomplete]: ... [case testGeneratorYieldAndYieldFrom] @@ -1728,13 +1728,13 @@ def g7(): from _typeshed import Incomplete from collections.abc import Generator -def g1() -> Generator[Incomplete, Incomplete, None]: ... -def g2() -> Generator[Incomplete, Incomplete, None]: ... -def g3() -> Generator[Incomplete, Incomplete, None]: ... -def g4() -> Generator[Incomplete, Incomplete, None]: ... -def g5() -> Generator[Incomplete, Incomplete, None]: ... +def g1() -> Generator[Incomplete, Incomplete]: ... +def g2() -> Generator[Incomplete, Incomplete]: ... +def g3() -> Generator[Incomplete, Incomplete]: ... +def g4() -> Generator[Incomplete, Incomplete]: ... +def g5() -> Generator[Incomplete, Incomplete]: ... def g6() -> Generator[Incomplete, Incomplete, Incomplete]: ... -def g7() -> Generator[Incomplete, Incomplete, None]: ... +def g7() -> Generator[Incomplete, Incomplete]: ... [case testCallable] from typing import Callable