Skip to content

Make PoolAcquireContext return a Connection type #578

New issue

Have a question about this project? # for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “#”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? # to your account

Open
wants to merge 1 commit into
base: master
Choose a base branch
from

Conversation

Yrlish
Copy link

@Yrlish Yrlish commented May 20, 2020

__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:
image

After:
image

`__aenter__` should be defined to return a `Connection` type. This will help IDEs to properly help suggesting methods when using with `pool.acquire() as con`
@Yrlish
Copy link
Author

Yrlish commented May 20, 2020

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

@elprans
Copy link
Member

elprans commented Jul 18, 2020

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

@DanielNoord
Copy link
Contributor

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
Copy link
Contributor

This is superseded by #1209

# for free to join this conversation on GitHub. Already have an account? # to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants