Fix union of generators having no return type

[omarsilva1]

Fixes #15141

This adds a check if iter_type is a union type so that expr_type is set to the generator return type instead of none.

[JelleZijlstra]

(the timing change is probably noise)

[hauntsaninja]

Thanks for working on this! I think ikonst's observation is a good one, e.g. what if we have a Union of Iterator and Generator or something.

I was playing around with this diff:

diff --git a/mypy/checker.py b/mypy/checker.py
index c1c31538b..6f1705cf7 100644
--- a/mypy/checker.py
+++ b/mypy/checker.py
@@ -952,8 +952,9 @@ class TypeChecker(NodeVisitor[None], CheckerPluginInterface):
             # 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 cd0ff1100..34dcbdef0 100644
--- a/mypy/checkexpr.py
+++ b/mypy/checkexpr.py
@@ -5177,17 +5177,7 @@ class ExpressionChecker(ExpressionVisitor[Type]):
 
         # 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, False)
 
         if not allow_none_return and isinstance(get_proper_type(expr_type), NoneType):
             self.chk.msg.does_not_return_value(None, e)

It causes two other tests to fail, but I think the new behaviour is better than the old behaviour. What do you think?

[omarsilva1]

Thanks for the feedback! I added your suggestions, I also think that this new behaviour is more natural than the previous one.

I only had one test failing locally, where Iterator[int] was the return type of a function. Should be fixed now.

(I see that the second test that fails has to do with mypyc (run-generators.test, specifically the second run_generator_and_throw call. I havent got that test suite locally configured yet)

[github-actions]

Diff from mypy_primer, showing the effect of this PR on open source code:

graphql-core (https://github.com/graphql-python/graphql-core): typechecking got 1.05x slower (363.9s -> 382.2s)
(Performance measurements are based on a single noisy sample)

rich (https://github.com/Textualize/rich)
+ rich/segment.py:608: error: No return value expected  [return-value]

[hauntsaninja]

Any luck with fixing the test?