Do not wrap AsyncIterator and Coroutine with extra Coroutine in async functions

mypy/checker.py

@@ -1003,7 +1003,9 @@ def get_coroutine_return_type(self, return_type: Type) -> Type:
             return AnyType(TypeOfAny.from_another_any, source_any=return_type)
         assert isinstance(return_type, Instance), "Should only be called on coroutine functions."
         # Note: return type is the 3rd type parameter of Coroutine.
-        return return_type.args[2]
+        if len(return_type.args) >= 3:
+            return return_type.args[2]
+        return AnyType(TypeOfAny.from_error)
 
     def get_generator_return_type(self, return_type: Type, is_coroutine: bool) -> Type:
         """Given the declared return type of a generator (t), return the type it returns (tr)."""
@@ -1355,8 +1357,13 @@ def check_func_def(
                     self.binder.pop_frame(True, 0)
 
             if not unreachable:
-                if defn.is_generator or is_named_instance(
-                    self.return_types[-1], "typing.AwaitableGenerator"
+                if (
+                    defn.is_generator
+                    or is_named_instance(self.return_types[-1], "typing.AwaitableGenerator")
+                    or (
+                        (t := getattr(self.return_types[-1], "type", None))
+                        and t.has_base("typing.AsyncIterator")
+                    )
                 ):
                     return_type = self.get_generator_return_type(
                         self.return_types[-1], defn.is_coroutine

mypy/semanal.py

@@ -979,9 +979,13 @@ def analyze_func_def(self, defn: FuncDef) -> None:
             and isinstance(defn.type, CallableType)
             and self.wrapped_coro_return_types.get(defn) != defn.type
         ):
-            if defn.is_async_generator:
-                # Async generator types are handled elsewhere
+            ret_typeinfo = getattr(defn.type.ret_type, "type", None)
+            if ret_typeinfo and ret_typeinfo.has_base("typing.AsyncIterator"):
+                defn.is_coroutine = True
+                defn.is_async_generator = True
                 pass
+            elif ret_typeinfo and ret_typeinfo.has_base("typing.Awaitable"):
+                defn.is_coroutine = True
             else:
                 # A coroutine defined as `async def foo(...) -> T: ...`
                 # has external return type `Coroutine[Any, Any, T]`.

test-data/unit/check-async-await.test

@@ -177,8 +177,6 @@ async def f() -> None:
 [builtins fixtures/async_await.pyi]
 [typing fixtures/typing-async.pyi]
 [out]
-main:7: error: "Coroutine[Any, Any, AsyncGenerator[str, None]]" has no attribute "__aiter__" (not async iterable)
-main:7: note: Maybe you forgot to use "await"?
 
 [case testAsyncForErrorCanBeIgnored]
 
@@ -683,14 +681,13 @@ from typing import AsyncGenerator, Any
 async def f() -> AsyncGenerator[Any, Any]:
     yield None
 
-async def h() -> Any:
+async def h() -> Any:  # E: The return type of an async generator function should be "AsyncGenerator" or one of its supertypes
     yield 0
 
 async def g():  # E: Function is missing a return type annotation
     yield 0
 [builtins fixtures/async_await.pyi]
 [typing fixtures/typing-async.pyi]
-[out]
 
 [case testAsyncOverloadedFunction]
 from typing import overload
@@ -915,7 +912,7 @@ async def f() -> AsyncGenerator[int, None]:
     async def g(x: int) -> int:
         return x
 
-    return (await g(x) for x in [1, 2, 3])
+    return (await g(x) for x in [1, 2, 3])  # E: "return" with value in async generator is not allowed
 
 [builtins fixtures/async_await.pyi]
 [typing fixtures/typing-async.pyi]
@@ -1054,10 +1051,10 @@ class P(Protocol):
         raise BaseException
 
 class Launcher(P):
-    def launch(self) -> AsyncIterator[int]:  # E: Return type "AsyncIterator[int]" of "launch" incompatible with return type "Coroutine[Any, Any, AsyncIterator[int]]" in supertype "P" \
-                                             # N: Consider declaring "launch" in supertype "P" without "async" \
-                                             # N: See https://mypy.readthedocs.io/en/stable/more_types.html#asynchronous-iterators
+    def launch(self) -> AsyncIterator[int]:
         raise BaseException
 
+
+
 [builtins fixtures/async_await.pyi]
 [typing fixtures/typing-async.pyi]

test-data/unit/check-functions.test

@@ -3382,3 +3382,49 @@ for factory in (
     reveal_type(factory)  # N: Revealed type is "def () -> Union[builtins.str, None]"
     var = factory()
 [builtins fixtures/tuple.pyi]
+
+[case testAsyncIteratorOverload]
+from typing import AsyncIterator, Callable, Iterator, TypeVar, Union, overload
+
+T = TypeVar("T")
+U = TypeVar("U")
+
+def deco(fn: Callable[[T], AsyncIterator[U]]) -> Callable[[T], U]:
+    ...
+
+@overload
+@deco
+async def test_async(val: int) -> AsyncIterator[int]: ...
+
+@overload
+@deco
+async def test_async(val: str) -> AsyncIterator[str]: ...
+
+@deco
+async def test_async(val: Union[int, str]) -> AsyncIterator[Union[int, str]]:
+    yield val
+
+[typing fixtures/typing-async.pyi]
+
+[case testAsyncIteratorOverloadMixedAsync]
+from typing import AsyncIterator, Callable, Iterator, TypeVar, Union, overload
+
+T = TypeVar("T")
+U = TypeVar("U")
+
+def deco(fn: Callable[[T], AsyncIterator[U]]) -> Callable[[T], U]:
+    ...
+
+@overload
+@deco
+def test_async(val: int) -> AsyncIterator[int]: ...
+
+@overload
+@deco
+def test_async(val: str) -> AsyncIterator[str]: ...
+
+@deco
+async def test_async(val: Union[int, str]) -> AsyncIterator[Union[int, str]]:
+    yield val
+
+[typing fixtures/typing-async.pyi]