Typehint *args for variable length heterogenous return values like asyncio.gather()

I'm having some trouble with the typesystem here. I need something like asyncio.gather() where the return value is the return values of the coroutines in the same order as given in the args.

async def f1() -> int: ...
async def f2() -> str: ...
async def f3() -> bool: ...

a, b, c = await asyncio.gather(f1(), f2(), f3())

This works:

a is of type int
b is of type str
c is of type bool

I have a function that accepts multiple partial[Coroutine[Any, Any, T]].

type PartialCoroutine[T] = partial[Coroutine[Any, Any, T]]

async def run runPartialCoroutines[T](*args: PartialCoroutine[T]) -> list[T]: ...

If I run my function it only uses the return value of the first function

async def f1() -> int: ...
async def f2() -> str: ...
async def f3() -> bool: ...

a, b, c = await runPartialCoroutines(f1(), f2(), f3())

This doesn't work:

a is of type int
b is of type int
c is of type int

I understand why it does this, but I can't seem to find a way to make it work. I've checked the typehint of asyncio.gather() but I have a feeling that the typecheckers have some custom logic built in for this case. Is it possible to do this at all?

Solution

The answer isn't pretty. Basically, typeshed provides a bunch of overloads for up to 6 arguments, so that static type checkers can make the correct inferences.

@overload
def gather(coro_or_future1: _FutureLike[_T1], /, *, return_exceptions: Literal[False] = False) -> Future[tuple[_T1]]: ...  # type: ignore[overload-overlap]
@overload
def gather(  # type: ignore[overload-overlap]
    coro_or_future1: _FutureLike[_T1], coro_or_future2: _FutureLike[_T2], /, *, return_exceptions: Literal[False] = False
) -> Future[tuple[_T1, _T2]]: ...
@overload
def gather(  # type: ignore[overload-overlap]
    coro_or_future1: _FutureLike[_T1],
    coro_or_future2: _FutureLike[_T2],
    coro_or_future3: _FutureLike[_T3],
    /,
    *,
    return_exceptions: Literal[False] = False,
) -> Future[tuple[_T1, _T2, _T3]]: ...
...