Improve support of @overload for await expressions

[rmk135]

Feature

The feature is to support @overload'ing for await expressions by return type only. Here is an example:

from typing import Callable, Any, TypeVar, Generic, Awaitable, overload


T = TypeVar('T')


class Provider(Generic[T]):
    def __init__(self, provides: Callable[..., T], *args: Any, **kwargs: Any) -> None: ...
    @overload
    def __call__(self, *args: Any, **kwargs: Any) -> T: ...
    @overload
    def __call__(self, *args: Any, **kwargs: Any) -> Awaitable[T]: ...


class Animal:
    ...


class Cat(Animal):
    ...


# Test 1: to check the return type
provider = Provider(Cat)
animal: Animal = provider(1, 2, 3)


# Test 2: to check the return type with "await"
async def _get_animal_async() -> None:
    animal: Animal = await provider(1, 2, 3)

With mypy 0.790 it produces next errors:

mypy_async.py:12: error: Overloaded function signature 2 will never be matched: signature 1's parameter type(s) are the same or broader
mypy_async.py:30: error: Incompatible types in "await" (actual type "Cat", expected type "Awaitable[Any]")

Pitch

I'm looking for this feature for the Dependency Injector.

The Provider class here is a generic wrapper. It returns actual instance or awaitable of this instance. The return type depends on the dependencies of the provider. If any of the underlaying providers returns awaitable it turns all the providers in the chain to the async mode:

provider1()  [async mode enabled]
│
├──> provider2()
│
├──> provider3()  [async mode enabled]
│    │
│    └──> provider4()  [returned Awaitable[T]]
│
└──> provider5()
     │
     └──> provider6()

result = await provider1()  # always returns awaitable in async mode

When async mode is enabled provider will always return awaitable.

I see this as a specific case of using @overload. It is intended to specify different return types depending on how the callable is called. While "how the callable is called" now is defined by the number of parameters, the use of callable in the await expression is just another kind of "how the callable is called".

PS: PyCharm already supports this example:

Appreciate community feedback.

[gvanrossum]

That overload doesn't work because there's no difference in argument list. What implementation did you have in mind? IOW how is your __call__ implementation going to know whether it should return an Awaitable or a result?

[hauntsaninja]

IIUC the proposal is to select overloads by return type as well (using some sort of bidirectional type inference), which is unsafe. This makes it pretty much a non-starter (there's a high bar to changing overload semantics, let alone to something unsafe).

I'm not familiar with PyCharm, but your screenshot + highlight just seems to show that it's using your type annotation to determine the type of animal.

The context on Gitter made this a little clearer (for me); I think OP basically wants unsafe union like semantics / to trust that the user (or some black magic) will call the function with the right name for the right provider.

[rmk135]

@gvanrossum @hauntsaninja

I'm working on the asynchronous dependency injection mode for the Dependency Injector.

Here is an example:

import asyncio
from typing import AsyncIterable

import aioredis
from dependency_injector import containers,  providers


async def init_redis_pool() -> AsyncIterable[aioredis.Redis]:
    redis = await aioredis.create_redis_pool('redis://localhost')

    yield redis

    redis.close()
    await redis.wait_closed()


class Service:
    def __init__(self, redis: aioredis.Redis):
        self._redis = redis


class Container(containers.DeclarativeContainer):

    redis_client = providers.Resource(init_redis_pool)

    service = providers.Factory(
        Service,
        redis=redis_client,
    )


async def main():
    container = Container()
    service = await container.service()


if __name__ == '__main__':
    asyncio.run(main())

Class Service has a dependency on redis connection pool. Redis connection pool initialization is asynchronous. It automatically turns container.service factory into async mode. When factory is set to async mode it will always return awaitable. Even if there are no async injections a synthetic future will be returned. That's an analogue of what happens here:

async def service() -> Service:
    return Service(
        redis=await init_redis_pool(),
    )

User can also explicitly enable the async mode by calling container.service.set_async_mode(True).

[rmk135]

@hauntsaninja

I'm not familiar with PyCharm, but your screenshot + highlight just seems to show that it's using your type annotation to determine the type of animal.

No. It breaks if the type is incompatible:

---

The context on Gitter made this a little clearer (for me); I think OP basically wants unsafe union like semantics / to trust that the user (or some black magic) will call the function with the right name for the right provider.

I don't feel that this feature introduces more magic than @overload example from the docs:

from typing import Union, overload

# Overload *variants* for 'mouse_event'.
# These variants give extra information to the type checker.
# They are ignored at runtime.

@overload
def mouse_event(x1: int, y1: int) -> ClickEvent: ...
@overload
def mouse_event(x1: int, y1: int, x2: int, y2: int) -> DragEvent: ...

# The actual *implementation* of 'mouse_event'.
# The implementation contains the actual runtime logic.
#
# It may or may not have type hints. If it does, mypy
# will check the body of the implementation against the
# type hints.
#
# Mypy will also check and make sure the signature is
# consistent with the provided variants.

def mouse_event(x1: int,
                y1: int,
                x2: Optional[int] = None,
                y2: Optional[int] = None) -> Union[ClickEvent, DragEvent]:
    if x2 is None and y2 is None:
        return ClickEvent(x1, y1)
    elif x2 is not None and y2 is not None:
        return DragEvent(x1, y1, x2, y2)
    else:
        raise TypeError("Bad arguments")

If we already do the type specification based on arguments, may be we also can do it for the return type in await expression:

@overload
def mouse_event(x1: int, y1: int) -> ClickEvent: ...
@overload
def mouse_event(x1: int, y1: int, x2: int, y2: int) -> DragEvent: ...
@overload
def mouse_event(x1: int, y1: int, x2: int, y2: int) -> Awaitable[DragEvent]: ...


event: DragEvent = mouse_event(1, 2)
event: DragEvent = mouse_event(1, 2, 3, 4)
event: DragEvent = await mouse_event(1, 2, 3, 4)

[rmk135]

@gvanrossum @hauntsaninja

I'm looking for your feedback to proceed with a feature I'm working on.

1. Do you consider this as a relevant use case (prioritize awaitable return type, if specified, for await expressions)?
2. If yes, when can you release it?
3. If no, could you recommend other ways to achieve what I'm doing (object with sync and async interfaces)? Currently my best guess is to just have separate methods for this. I'm curious if there is a better solution.

Respectfully,
Roman

[hauntsaninja]

No, this isn't something mypy is going to do. Separate methods (i.e. await provider.async_(1, 2, 3)) is an easy solution.