Make PoolAcquireContext return a Connection type

[Yrlish]

__aenter__ in PoolAcquireContext should be defined to return a Connection type. This will help IDEs to properly help suggesting methods when using syntax:

with my_pool.acquire() as con: con # this is now defined as type Connection by the acquire method pass

instead of doing this:

with my_pool.acquire() as con: # type: Connection pass

Before:

After:

[Yrlish]

I forgot I have made an relevant issue 1.5 years ago #387.

[elprans]

There is a comprehensive effort to add type annotations in #577

[DanielNoord]

This type annotation would actually be incorrect. The correct patch is:

diff --git a/asyncpg/pool.py b/asyncpg/pool.py index e3898d5..0bea17a 100644 --- a/asyncpg/pool.py +++ b/asyncpg/pool.py @@ -10,6 +10,7 @@ import functools import inspect import logging import time +from typing import Optional import warnings from . import compat @@ -384,7 +385,7 @@ class Pool: self._holders = [] self._initialized = False self._initializing = False - self._queue = None + self._queue: Optional[asyncio.LifoQueue[PoolConnectionHolder]] = None self._connection_class = connection_class self._record_class = record_class @@ -842,11 +843,12 @@ class Pool: """ return PoolAcquireContext(self, timeout) - async def _acquire(self, timeout): - async def _acquire_impl(): - ch = await self._queue.get() # type: PoolConnectionHolder + async def _acquire(self, timeout: Optional[float]) -> PoolConnectionProxy: + async def _acquire_impl() -> PoolConnectionProxy: + assert self._queue is not None, "Pool is not initialized" + ch = await self._queue.get() try: - proxy = await ch.acquire() # type: PoolConnectionProxy + proxy = await ch.acquire() except (Exception, asyncio.CancelledError): self._queue.put_nowait(ch) raise @@ -1015,10 +1017,10 @@ class PoolAcquireContext: def __init__(self, pool, timeout): self.pool = pool self.timeout = timeout - self.connection = None + self.connection: Optional[PoolConnectionProxy] = None self.done = False - async def __aenter__(self): + async def __aenter__(self) -> PoolConnectionProxy: if self.connection is not None or self.done: raise exceptions.InterfaceError('a connection is already acquired') self.connection = await self.pool._acquire(self.timeout) 

Since _queue on Pool is Generic on PoolConnectionProxy __aenter__ should also reflect this.
I'd suggest to make a PR out of this patch after #1197 is merged (as it types PoolAcquireContext.pool correctly, which enables the inference of the type as I did above) and closing this PR.

[DanielNoord]

This is superseded by #1209