Repository: python/mypy

Issue:
@overload on async callables accepting and returning union types leads to type erasure
**Bug Report**

 Async callable which accepts a parameter with a union type, returns a union type and has `@overload` on said parameter has its return type erased when awaited. I'm hoping this sentence makes any sense, but in any case (pun not intended) the code snippet should demonstrate the issue.

**To Reproduce**

Here's the absolute minimum reproduction I could muster:

```python
from typing import overload

class A: pass
class B: pass

@overload
async def foo(x: A) -> B: ...

@overload
async def foo(x: B) -> A: ...

async def foo(x: A | B) -> A | B:
    return A() if isinstance(x, B) else B()

async def breaks(x: A | B) -> None:
    y = await foo(x)
    reveal_type(y) # Reveals Any, should reveal A | B
```

The following includes non async version of this code which actually works.

Gist URL: https://gist.github.com/2f8288673587ff27f6fca5024ddd0543
Playground URL: https://mypy-play.net/?mypy=master&python=3.11&gist=2f8288673587ff27f6fca5024ddd0543&flags=strict

**Expected Behavior**

`await foo(x)` in the above example should have type `A | B`.

**Actual Behavior**

`note: Revealed type is "Any"` is the actual result.

**Your Environment**

Original bug was found on Python 3.10 and mypy 0.991, with `--strict` and a number of per-file settings. I can include relevant excerpt of our `pyproject.toml`, although it seems that this issue is not configuration-specific.

On the playground the following was tried:
- Mypy version used: 0.991 (also tried with latest `master` as well as a random sampling of versions between `0.670` and `0.990`
- Mypy command-line flags: none and `--strict`, in fact none of the flag combinations lead to correct result
- Python version used: 3.10, 3.11




Discussion:

Resolving patch:
diff --git a/mypy/checker.py b/mypy/checker.py
--- a/mypy/checker.py
+++ b/mypy/checker.py
@@ -843,6 +843,10 @@ def get_generator_yield_type(self, return_type: Type, is_coroutine: bool) -> Typ
 
         if isinstance(return_type, AnyType):
             return AnyType(TypeOfAny.from_another_any, source_any=return_type)
+        elif isinstance(return_type, UnionType):
+            return make_simplified_union(
+                [self.get_generator_yield_type(item, is_coroutine) for item in return_type.items]
+            )
         elif not self.is_generator_return_type(
             return_type, is_coroutine
         ) and not self.is_async_generator_return_type(return_type):
@@ -873,6 +877,10 @@ def get_generator_receive_type(self, return_type: Type, is_coroutine: bool) -> T
 
         if isinstance(return_type, AnyType):
             return AnyType(TypeOfAny.from_another_any, source_any=return_type)
+        elif isinstance(return_type, UnionType):
+            return make_simplified_union(
+                [self.get_generator_receive_type(item, is_coroutine) for item in return_type.items]
+            )
         elif not self.is_generator_return_type(
             return_type, is_coroutine
         ) and not self.is_async_generator_return_type(return_type):
@@ -912,6 +920,10 @@ def get_generator_return_type(self, return_type: Type, is_coroutine: bool) -> Ty
 
         if isinstance(return_type, AnyType):
             return AnyType(TypeOfAny.from_another_any, source_any=return_type)
+        elif isinstance(return_type, UnionType):
+            return make_simplified_union(
+                [self.get_generator_return_type(item, is_coroutine) for item in return_type.items]
+            )
         elif not self.is_generator_return_type(return_type, is_coroutine):
             # If the function doesn't have a proper Generator (or
             # Awaitable) return type, anything is permissible.


Test patch:
diff --git a/test-data/unit/check-async-await.test b/test-data/unit/check-async-await.test
--- a/test-data/unit/check-async-await.test
+++ b/test-data/unit/check-async-await.test
@@ -925,3 +925,21 @@ async def f() -> AsyncGenerator[int, None]:
 
 [builtins fixtures/async_await.pyi]
 [typing fixtures/typing-async.pyi]
+
+[case testAwaitUnion]
+from typing import overload, Union
+
+class A: ...
+class B: ...
+
+@overload
+async def foo(x: A) -> B: ...
+@overload
+async def foo(x: B) -> A: ...
+async def foo(x): ...
+
+async def bar(x: Union[A, B]) -> None:
+    reveal_type(await foo(x))  # N: Revealed type is "Union[__main__.B, __main__.A]"
+
+[builtins fixtures/async_await.pyi]
+[typing fixtures/typing-async.pyi]
diff --git a/test-data/unit/check-statements.test b/test-data/unit/check-statements.test
--- a/test-data/unit/check-statements.test
+++ b/test-data/unit/check-statements.test
@@ -2206,3 +2206,12 @@ def foo():
     x: int = "no"  # N: By default the bodies of untyped functions are not checked, consider using --check-untyped-defs
     y = "no"  # type: int  # N: By default the bodies of untyped functions are not checked, consider using --check-untyped-defs
     z: int  # N: By default the bodies of untyped functions are not checked, consider using --check-untyped-defs
+
+[case testGeneratorUnion]
+from typing import Generator, Union
+
+class A: pass
+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]")