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:

withmy_pool.acquire() ascon: con# this is now defined as type Connection by the acquire methodpass

instead of doing this:

withmy_pool.acquire() ascon: # type: Connectionpass

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