Add better types for asyncio.gather (#9678)

Author: Gobot1234

stdlib/asyncio/tasks.pyi

@@ -68,6 +68,7 @@ _T2 = TypeVar("_T2")
 _T3 = TypeVar("_T3")
 _T4 = TypeVar("_T4")
 _T5 = TypeVar("_T5")
+_T6 = TypeVar("_T6")
 _FT = TypeVar("_FT", bound=Future[Any])
 _FutureLike: TypeAlias = Future[_T] | Generator[Any, None, _T] | Awaitable[_T]
 _TaskYieldType: TypeAlias = Future[object] | None
@@ -131,6 +132,19 @@ if sys.version_info >= (3, 10):
         return_exceptions: Literal[False] = False,
     ) -> Future[tuple[_T1, _T2, _T3, _T4, _T5]]: ...
     @overload
+    def gather(  # type: ignore[misc]
+        __coro_or_future1: _FutureLike[_T1],
+        __coro_or_future2: _FutureLike[_T2],
+        __coro_or_future3: _FutureLike[_T3],
+        __coro_or_future4: _FutureLike[_T4],
+        __coro_or_future5: _FutureLike[_T5],
+        __coro_or_future6: _FutureLike[_T6],
+        *,
+        return_exceptions: Literal[False] = False,
+    ) -> Future[tuple[_T1, _T2, _T3, _T4, _T5, _T6]]: ...
+    @overload
+    def gather(*coros_or_futures: _FutureLike[_T], return_exceptions: Literal[False] = False) -> Future[list[_T]]: ...  # type: ignore[misc]
+    @overload
     def gather(__coro_or_future1: _FutureLike[_T1], *, return_exceptions: bool) -> Future[tuple[_T1 | BaseException]]: ...  # type: ignore[misc]
     @overload
     def gather(  # type: ignore[misc]
@@ -166,7 +180,27 @@ if sys.version_info >= (3, 10):
         tuple[_T1 | BaseException, _T2 | BaseException, _T3 | BaseException, _T4 | BaseException, _T5 | BaseException]
     ]: ...
     @overload
-    def gather(*coros_or_futures: _FutureLike[Any], return_exceptions: bool = False) -> Future[list[Any]]: ...
+    def gather(  # type: ignore[misc]
+        __coro_or_future1: _FutureLike[_T1],
+        __coro_or_future2: _FutureLike[_T2],
+        __coro_or_future3: _FutureLike[_T3],
+        __coro_or_future4: _FutureLike[_T4],
+        __coro_or_future5: _FutureLike[_T5],
+        __coro_or_future6: _FutureLike[_T6],
+        *,
+        return_exceptions: bool,
+    ) -> Future[
+        tuple[
+            _T1 | BaseException,
+            _T2 | BaseException,
+            _T3 | BaseException,
+            _T4 | BaseException,
+            _T5 | BaseException,
+            _T6 | BaseException,
+        ]
+    ]: ...
+    @overload
+    def gather(*coros_or_futures: _FutureLike[_T], return_exceptions: bool) -> Future[list[_T | BaseException]]: ...
 
 else:
     @overload
@@ -212,6 +246,22 @@ else:
         return_exceptions: Literal[False] = False,
     ) -> Future[tuple[_T1, _T2, _T3, _T4, _T5]]: ...
     @overload
+    def gather(  # type: ignore[misc]
+        __coro_or_future1: _FutureLike[_T1],
+        __coro_or_future2: _FutureLike[_T2],
+        __coro_or_future3: _FutureLike[_T3],
+        __coro_or_future4: _FutureLike[_T4],
+        __coro_or_future5: _FutureLike[_T5],
+        __coro_or_future6: _FutureLike[_T6],
+        *,
+        loop: AbstractEventLoop | None = None,
+        return_exceptions: Literal[False] = False,
+    ) -> Future[tuple[_T1, _T2, _T3, _T4, _T5, _T6]]: ...
+    @overload
+    def gather(  # type: ignore[misc]
+        *coros_or_futures: _FutureLike[_T], loop: AbstractEventLoop | None = None, return_exceptions: Literal[False] = False
+    ) -> Future[list[_T]]: ...
+    @overload
     def gather(  # type: ignore[misc]
         __coro_or_future1: _FutureLike[_T1], *, loop: AbstractEventLoop | None = None, return_exceptions: bool
     ) -> Future[tuple[_T1 | BaseException]]: ...
@@ -249,16 +299,24 @@ else:
         __coro_or_future3: _FutureLike[_T3],
         __coro_or_future4: _FutureLike[_T4],
         __coro_or_future5: _FutureLike[_T5],
+        __coro_or_future6: _FutureLike[_T6],
         *,
         loop: AbstractEventLoop | None = None,
         return_exceptions: bool,
     ) -> Future[
-        tuple[_T1 | BaseException, _T2 | BaseException, _T3 | BaseException, _T4 | BaseException, _T5 | BaseException]
+        tuple[
+            _T1 | BaseException,
+            _T2 | BaseException,
+            _T3 | BaseException,
+            _T4 | BaseException,
+            _T5 | BaseException,
+            _T6 | BaseException,
+        ]
     ]: ...
     @overload
-    def gather(
-        *coros_or_futures: _FutureLike[Any], loop: AbstractEventLoop | None = None, return_exceptions: bool = False
-    ) -> Future[list[Any]]: ...
+    def gather(  # type: ignore[misc]
+        *coros_or_futures: _FutureLike[_T], loop: AbstractEventLoop | None = None, return_exceptions: bool
+    ) -> Future[list[_T | BaseException]]: ...
 
 def run_coroutine_threadsafe(coro: _FutureLike[_T], loop: AbstractEventLoop) -> concurrent.futures.Future[_T]: ...
 

test_cases/stdlib/asyncio/check_gather.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import asyncio
-from typing import Any, Awaitable, List, Tuple, Union
+from typing import Awaitable, List, Tuple, Union
 from typing_extensions import assert_type
 
 
@@ -21,14 +21,18 @@ async def test_gather(awaitable1: Awaitable[int], awaitable2: Awaitable[str]) ->
     assert_type(b, Tuple[Union[int, BaseException], Union[str, BaseException]])
 
     c = await asyncio.gather(awaitable1, awaitable2, awaitable1, awaitable1, awaitable1, awaitable1)
-    assert_type(c, List[Any])
+    assert_type(c, Tuple[int, str, int, int, int, int])
+
+    d = await asyncio.gather(awaitable1, awaitable1, awaitable1, awaitable1, awaitable1, awaitable1, awaitable1)
+    assert_type(d, List[int])
 
     awaitables_list: list[Awaitable[int]] = [awaitable1]
-    d = await asyncio.gather(*awaitables_list)
-    assert_type(d, List[Any])
+    e = await asyncio.gather(*awaitables_list)
+    assert_type(e, List[int])
 
-    e = await asyncio.gather()
-    assert_type(e, List[Any])
+    # this case isn't reliable between typecheckers, no one would ever call it with no args anyway
+    # f = await asyncio.gather()
+    # assert_type(f, list[Any])
 
 
 asyncio.run(test_gather(coro1(), coro2()))