From 394d17b758bae6c95cbe91b84c5cccf0f4d73c28 Mon Sep 17 00:00:00 2001 From: Shantanu <12621235+hauntsaninja@users.noreply.github.com> Date: Wed, 20 Mar 2024 19:08:56 -0700 Subject: [PATCH] Improve yield from inference for unions of generators (#16717) Fixes #15141, closes #15168 --- mypy/checker.py | 5 +-- mypy/checkexpr.py | 12 +------ mypyc/test-data/run-generators.test | 9 +++--- test-data/unit/check-statements.test | 47 +++++++++++++++++++++++++++- 4 files changed, 55 insertions(+), 18 deletions(-) diff --git a/mypy/checker.py b/mypy/checker.py index 941dc06f1c71..5d243195d50f 100644 --- a/mypy/checker.py +++ b/mypy/checker.py @@ -989,8 +989,9 @@ def get_generator_return_type(self, return_type: Type, is_coroutine: bool) -> Ty # AwaitableGenerator, Generator: tr is args[2]. return return_type.args[2] else: - # Supertype of Generator (Iterator, Iterable, object): tr is any. - return AnyType(TypeOfAny.special_form) + # We have a supertype of Generator (Iterator, Iterable, object) + # Treat `Iterator[X]` as a shorthand for `Generator[X, Any, None]`. + return NoneType() def visit_func_def(self, defn: FuncDef) -> None: if not self.recurse_into_functions: diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index 37a90ce55b9e..e7567eafb8fe 100644 --- a/mypy/checkexpr.py +++ b/mypy/checkexpr.py @@ -5963,17 +5963,7 @@ 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": - expr_type = self.chk.get_generator_return_type(iter_type, False) - else: - # Non-Generators don't return anything from `yield from` expressions. - # However special-case Any (which might be produced by an error). - actual_item_type = get_proper_type(actual_item_type) - if isinstance(actual_item_type, AnyType): - expr_type = AnyType(TypeOfAny.from_another_any, source_any=actual_item_type) - else: - # Treat `Iterator[X]` as a shorthand for `Generator[X, None, Any]`. - expr_type = NoneType() + expr_type = self.chk.get_generator_return_type(iter_type, is_coroutine=False) if not allow_none_return and isinstance(get_proper_type(expr_type), NoneType): self.chk.msg.does_not_return_value(None, e) diff --git a/mypyc/test-data/run-generators.test b/mypyc/test-data/run-generators.test index bcf9da1846ae..7e9804c49582 100644 --- a/mypyc/test-data/run-generators.test +++ b/mypyc/test-data/run-generators.test @@ -246,12 +246,12 @@ assert run_generator(another_triple()()) == ((1,), None) assert run_generator(outer()) == ((0, 1, 2, 3, 4), None) [case testYieldThrow] -from typing import Generator, Iterable, Any +from typing import Generator, Iterable, Any, Union from traceback import print_tb from contextlib import contextmanager import wrapsys -def generator() -> Iterable[int]: +def generator() -> Generator[int, None, Union[int, None]]: try: yield 1 yield 2 @@ -264,6 +264,7 @@ def generator() -> Iterable[int]: else: print('caught exception without value') return 0 + return None def no_except() -> Iterable[int]: yield 1 @@ -355,11 +356,11 @@ with ctx_manager() as c: raise Exception File "native.py", line 10, in generator yield 3 - File "native.py", line 30, in wrapper + File "native.py", line 31, in wrapper return (yield from x) File "native.py", line 9, in generator yield 2 - File "native.py", line 30, in wrapper + File "native.py", line 31, in wrapper return (yield from x) caught exception without value caught exception with value some string diff --git a/test-data/unit/check-statements.test b/test-data/unit/check-statements.test index f5b47e7ab97f..71cc80719779 100644 --- a/test-data/unit/check-statements.test +++ b/test-data/unit/check-statements.test @@ -85,7 +85,7 @@ def f() -> Generator[int, None, None]: from typing import Iterator def f() -> Iterator[int]: yield 1 - return "foo" + return "foo" # E: No return value expected [out] @@ -2231,6 +2231,51 @@ 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 testYieldFromUnionOfGenerators] +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]: + return (yield from arg) + +[case testYieldFromInvalidUnionReturn] +from typing import Generator, Union + +class A: pass +class B: pass + +def foo(arg: Union[A, B]) -> Generator[Union[int, str], None, A]: + return (yield from arg) # E: "yield from" can't be applied to "Union[A, B]" + +[case testYieldFromUnionOfGeneratorWithIterableStr] +from typing import Generator, Union, Iterable, Optional + +def foo(arg: Union[Generator[int, None, bytes], Iterable[str]]) -> Generator[Union[int, str], None, Optional[bytes]]: + return (yield from arg) + +def bar(arg: Generator[str, None, str]) -> Generator[str, None, str]: + return foo(arg) # E: Incompatible return value type (got "Generator[Union[int, str], None, Optional[bytes]]", expected "Generator[str, None, str]") + +def launder(arg: Iterable[str]) -> Generator[Union[int, str], None, Optional[bytes]]: + return foo(arg) + +def baz(arg: Generator[str, None, str]) -> Generator[Union[int, str], None, Optional[bytes]]: + # this is unsound, the Generator return type will actually be str + return launder(arg) +[builtins fixtures/tuple.pyi] + +[case testYieldIteratorReturn] +from typing import Iterator + +def get_strings(foo: bool) -> Iterator[str]: + if foo: + return ["foo1", "foo2"] # E: No return value expected + else: + yield "bar1" + yield "bar2" +[builtins fixtures/tuple.pyi] + [case testNoCrashOnStarRightHandSide] x = *(1, 2, 3) # E: can't use starred expression here [builtins fixtures/tuple.pyi]