Skip to content
New issue

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

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

Already on GitHub? Sign in to your account

Make PoolAcquireContext return a Connection type #578

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

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

Successfully merging this pull request may close these issues.

3 participants