Skip to content

Commit 72c413d

Browse files
authored
stubgen: Use Generator type var defaults (#17670)
Fixes #17669
1 parent 0c10367 commit 72c413d

File tree

2 files changed

+27
-20
lines changed

2 files changed

+27
-20
lines changed

mypy/stubgen.py

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -588,8 +588,8 @@ def _get_func_return(self, o: FuncDef, ctx: FunctionContext) -> str | None:
588588
if has_yield_expression(o) or has_yield_from_expression(o):
589589
generator_name = self.add_name("collections.abc.Generator")
590590
yield_name = "None"
591-
send_name = "None"
592-
return_name = "None"
591+
send_name: str | None = None
592+
return_name: str | None = None
593593
if has_yield_from_expression(o):
594594
yield_name = send_name = self.add_name("_typeshed.Incomplete")
595595
else:
@@ -600,7 +600,14 @@ def _get_func_return(self, o: FuncDef, ctx: FunctionContext) -> str | None:
600600
send_name = self.add_name("_typeshed.Incomplete")
601601
if has_return_statement(o):
602602
return_name = self.add_name("_typeshed.Incomplete")
603-
return f"{generator_name}[{yield_name}, {send_name}, {return_name}]"
603+
if return_name is not None:
604+
if send_name is None:
605+
send_name = "None"
606+
return f"{generator_name}[{yield_name}, {send_name}, {return_name}]"
607+
elif send_name is not None:
608+
return f"{generator_name}[{yield_name}, {send_name}]"
609+
else:
610+
return f"{generator_name}[{yield_name}]"
604611
if not has_return_statement(o) and o.abstract_status == NOT_ABSTRACT:
605612
return "None"
606613
return None

test-data/unit/stubgen.test

Lines changed: 17 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1640,11 +1640,11 @@ def all():
16401640
from _typeshed import Incomplete
16411641
from collections.abc import Generator
16421642

1643-
def f() -> Generator[Incomplete, None, None]: ...
1644-
def g() -> Generator[None, Incomplete, None]: ...
1645-
def h1() -> Generator[None, None, None]: ...
1643+
def f() -> Generator[Incomplete]: ...
1644+
def g() -> Generator[None, Incomplete]: ...
1645+
def h1() -> Generator[None]: ...
16461646
def h2() -> Generator[None, None, Incomplete]: ...
1647-
def h3() -> Generator[None, None, None]: ...
1647+
def h3() -> Generator[None]: ...
16481648
def all() -> Generator[Incomplete, Incomplete, Incomplete]: ...
16491649

16501650
[case testFunctionYieldsNone]
@@ -1656,8 +1656,8 @@ def g():
16561656
[out]
16571657
from collections.abc import Generator
16581658

1659-
def f() -> Generator[None, None, None]: ...
1660-
def g() -> Generator[None, None, None]: ...
1659+
def f() -> Generator[None]: ...
1660+
def g() -> Generator[None]: ...
16611661

16621662
[case testGeneratorAlreadyDefined]
16631663
class Generator:
@@ -1671,7 +1671,7 @@ from collections.abc import Generator as _Generator
16711671

16721672
class Generator: ...
16731673

1674-
def f() -> _Generator[Incomplete, None, None]: ...
1674+
def f() -> _Generator[Incomplete]: ...
16751675

16761676
[case testGeneratorYieldFrom]
16771677
def g1():
@@ -1692,10 +1692,10 @@ def g5():
16921692
from _typeshed import Incomplete
16931693
from collections.abc import Generator
16941694

1695-
def g1() -> Generator[Incomplete, Incomplete, None]: ...
1696-
def g2() -> Generator[Incomplete, Incomplete, None]: ...
1697-
def g3() -> Generator[Incomplete, Incomplete, None]: ...
1698-
def g4() -> Generator[Incomplete, Incomplete, None]: ...
1695+
def g1() -> Generator[Incomplete, Incomplete]: ...
1696+
def g2() -> Generator[Incomplete, Incomplete]: ...
1697+
def g3() -> Generator[Incomplete, Incomplete]: ...
1698+
def g4() -> Generator[Incomplete, Incomplete]: ...
16991699
def g5() -> Generator[Incomplete, Incomplete, Incomplete]: ...
17001700

17011701
[case testGeneratorYieldAndYieldFrom]
@@ -1728,13 +1728,13 @@ def g7():
17281728
from _typeshed import Incomplete
17291729
from collections.abc import Generator
17301730

1731-
def g1() -> Generator[Incomplete, Incomplete, None]: ...
1732-
def g2() -> Generator[Incomplete, Incomplete, None]: ...
1733-
def g3() -> Generator[Incomplete, Incomplete, None]: ...
1734-
def g4() -> Generator[Incomplete, Incomplete, None]: ...
1735-
def g5() -> Generator[Incomplete, Incomplete, None]: ...
1731+
def g1() -> Generator[Incomplete, Incomplete]: ...
1732+
def g2() -> Generator[Incomplete, Incomplete]: ...
1733+
def g3() -> Generator[Incomplete, Incomplete]: ...
1734+
def g4() -> Generator[Incomplete, Incomplete]: ...
1735+
def g5() -> Generator[Incomplete, Incomplete]: ...
17361736
def g6() -> Generator[Incomplete, Incomplete, Incomplete]: ...
1737-
def g7() -> Generator[Incomplete, Incomplete, None]: ...
1737+
def g7() -> Generator[Incomplete, Incomplete]: ...
17381738

17391739
[case testCallable]
17401740
from typing import Callable

0 commit comments

Comments
 (0)