Skip to content

Commit

Permalink
stubgen: Use Generator type var defaults (#17670)
Browse files Browse the repository at this point in the history
Fixes #17669
  • Loading branch information
srittau committed Sep 12, 2024
1 parent 0c10367 commit 72c413d
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 20 deletions.
13 changes: 10 additions & 3 deletions mypy/stubgen.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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
Expand Down
34 changes: 17 additions & 17 deletions test-data/unit/stubgen.test
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand All @@ -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:
Expand All @@ -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():
Expand All @@ -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]
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 72c413d

Please sign in to comment.