Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix union of generators having no return type #15168

Closed
wants to merge 15 commits into from
Closed
6 changes: 5 additions & 1 deletion mypy/checkexpr.py
Original file line number Diff line number Diff line change
Expand Up @@ -5095,7 +5095,11 @@ def visit_yield_from_expr(self, e: YieldFromExpr, allow_none_return: bool = Fals

# Determine the type of the entire yield from expression.
iter_type = get_proper_type(iter_type)
if isinstance(iter_type, Instance) and iter_type.type.fullname == "typing.Generator":
if (
isinstance(iter_type, Instance)
and iter_type.type.fullname == "typing.Generator"
or isinstance(iter_type, UnionType)
omarsilva1 marked this conversation as resolved.
Show resolved Hide resolved
):
expr_type = self.chk.get_generator_return_type(iter_type, False)
else:
# Non-Generators don't return anything from `yield from` expressions.
Expand Down
26 changes: 26 additions & 0 deletions test-data/unit/check-statements.test
Original file line number Diff line number Diff line change
Expand Up @@ -2195,6 +2195,32 @@ class B: pass
def foo(x: int) -> Union[Generator[A, None, None], Generator[B, None, None]]:
yield x # E: Incompatible types in "yield" (actual type "int", expected type "Union[A, B]")

[case testReturnTypeOfGeneratorUnion]
from typing import Generator, Union

class T: pass

def foo(arg: Union[Generator[int, None, T], Generator[str, None, T]]) -> Generator[Union[int, str], None, T]:
omarsilva1 marked this conversation as resolved.
Show resolved Hide resolved
return (yield from arg)

[case testYieldFromError]
from typing import Generator, Union

class T: pass
class R: pass

def foo(arg: Union[T, R]) -> Generator[Union[int, str], None, T]:
omarsilva1 marked this conversation as resolved.
Show resolved Hide resolved
return (yield from arg) # E: "yield from" cannot be applied to "Union[T, R]"

[case testUnionOfGeneratorWithIterableStr]
from typing import Generator, Union, Iterable

class T: pass

def foo(arg: Union[Generator[int, None, T], Iterable[str]]) -> Generator[Union[int, str], None, T]:
return (yield from arg)
[builtins fixtures/tuple.pyi]

[case testNoCrashOnStarRightHandSide]
x = *(1, 2, 3) # E: Can use starred expression only as assignment target
[builtins fixtures/tuple.pyi]
Expand Down