Skip to content

Commit

Permalink
Add better types for asyncio.gather (#9678)
Browse files Browse the repository at this point in the history
  • Loading branch information
@Gobot1234
Gobot1234 authored Oct 4, 2023
1 parent c30c85d commit 25eb99c
Show file tree
Hide file tree
Showing 2 changed files with 73 additions and 11 deletions.
68 changes: 63 additions & 5 deletions stdlib/asyncio/tasks.pyi
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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]]: ...
Expand Down Expand Up @@ -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]: ...

Expand Down
16 changes: 10 additions & 6 deletions test_cases/stdlib/asyncio/check_gather.py
View file
Original file line number Diff line number Diff line change
@@ -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


Expand All @@ -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()))

0 comments on commit 25eb99c

Please sign in to comment.