mosquito / aio-pika

AMQP 0.9 client designed for asyncio and humans.
https://aio-pika.readthedocs.org/
Apache License 2.0
1.23k stars 187 forks source link

Correct type annotations #564

Open Lancetnik opened 1 year ago

Lancetnik commented 1 year ago

@mosquito hello again! You know, the only thing annoying to me when I had working with the aio-pika is some function annotations disharmony. Like in the following case Here, the function returns RobustConnection by default, but the response objected is annotated by AbstractRobustConnection causes to some autocompletion troubles and making a DE worse.

So, this idea has taken over my mind and finally, I found a decision:

from typing import TypeVar, overload

ConnectionCls = TypeVar("ConnectionCls", bound=AbstractRobustConnection)

@overload
async def connect_robust(
    ...
    connection_class: Type[RobustConnection] = RobustConnection,
) -> RobustConnection: ...

@overload
async def connect_robust(
    ...
    connection_class: Type[ConnectionCls],
) -> ConnectionCls: ...

async def connect_robust(
    ...
    connection_class: Type[ConnectionCls] = RobustConnection,
) -> ConnectionCls: ...

If we add the 2 overload method annotations, it will be annotated correctly in both cases (with custom class arg and the default one). So, the package users will get a better experience this way.

What do you think about fixing this behavior in all package? I can make a PR, but it is a pretty big job to fix it everywhere and I am still busy with the Propan for now. If somebody can fix it earlier, it should be great.