Add typing to auxiliary classes around Pool (#1197)

Author: DanielNoord

asyncpg/pool.py

@@ -4,12 +4,16 @@
 # This module is part of asyncpg and is released under
 # the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0
 
+from __future__ import annotations
 
 import asyncio
+from collections.abc import Awaitable, Callable
 import functools
 import inspect
 import logging
 import time
+from types import TracebackType
+from typing import Any, Optional, Type
 import warnings
 
 from . import compat
@@ -23,7 +27,14 @@
 
 class PoolConnectionProxyMeta(type):
 
- def __new__(mcls, name, bases, dct, *, wrap=False):
+ def __new__(
+ mcls,
+ name: str,
+ bases: tuple[Type[Any], ...],
+ dct: dict[str, Any],
+ *,
+ wrap: bool = False,
+ ) -> PoolConnectionProxyMeta:
  if wrap:
  for attrname in dir(connection.Connection):
  if attrname.startswith('_') or attrname in dct:
@@ -44,8 +55,10 @@ def __new__(mcls, name, bases, dct, *, wrap=False):
  return super().__new__(mcls, name, bases, dct)
 
  @staticmethod
- def _wrap_connection_method(meth_name, iscoroutine):
- def call_con_method(self, *args, **kwargs):
+ def _wrap_connection_method(
+ meth_name: str, iscoroutine: bool
+ ) -> Callable[..., Any]:
+ def call_con_method(self: Any, *args: Any, **kwargs: Any) -> Any:
  # This method will be owned by PoolConnectionProxy class.
  if self._con is None:
  raise exceptions.InterfaceError(
@@ -68,25 +81,26 @@ class PoolConnectionProxy(connection._ConnectionProxy,
 
  __slots__ = ('_con', '_holder')
 
- def __init__(self, holder: 'PoolConnectionHolder',
- con: connection.Connection):
+ def __init__(
+ self, holder: PoolConnectionHolder, con: connection.Connection
+ ) -> None:
  self._con = con
  self._holder = holder
  con._set_proxy(self)
 
- def __getattr__(self, attr):
+ def __getattr__(self, attr: str) -> Any:
  # Proxy all unresolved attributes to the wrapped Connection object.
  return getattr(self._con, attr)
 
- def _detach(self) -> connection.Connection:
+ def _detach(self) -> Optional[connection.Connection]:
  if self._con is None:
  return
 
  con, self._con = self._con, None
  con._set_proxy(None)
  return con
 
- def __repr__(self):
+ def __repr__(self) -> str:
  if self._con is None:
  return '<{classname} [released]{id:#x}>'.format(
  classname=self.__class__.__name__, id=id(self))
@@ -103,27 +117,34 @@ class PoolConnectionHolder:
  '_inactive_callback', '_timeout',
  '_generation')
 
- def __init__(self, pool, *, max_queries, setup, max_inactive_time):
+ def __init__(
+ self,
+ pool: "Pool",
+ *,
+ max_queries: float,
+ setup: Optional[Callable[[PoolConnectionProxy], Awaitable[None]]],
+ max_inactive_time: float,
+ ) -> None:
 
  self._pool = pool
- self._con = None
- self._proxy = None
+ self._con: Optional[connection.Connection] = None
+ self._proxy: Optional[PoolConnectionProxy] = None
 
  self._max_queries = max_queries
  self._max_inactive_time = max_inactive_time
  self._setup = setup
- self._inactive_callback = None
- self._in_use = None # type: asyncio.Future
- self._timeout = None
- self._generation = None
+ self._inactive_callback: Optional[Callable] = None
+ self._in_use: Optional[asyncio.Future] = None
+ self._timeout: Optional[float] = None
+ self._generation: Optional[int] = None
 
- def is_connected(self):
+ def is_connected(self) -> bool:
  return self._con is not None and not self._con.is_closed()
 
- def is_idle(self):
+ def is_idle(self) -> bool:
  return not self._in_use
 
- async def connect(self):
+ async def connect(self) -> None:
  if self._con is not None:
  raise exceptions.InternalClientError(
  'PoolConnectionHolder.connect() called while another '
@@ -171,7 +192,7 @@ async def acquire(self) -> PoolConnectionProxy:
 
  return proxy
 
- async def release(self, timeout):
+ async def release(self, timeout: Optional[float]) -> None:
  if self._in_use is None:
  raise exceptions.InternalClientError(
  'PoolConnectionHolder.release() called on '
@@ -234,25 +255,25 @@ async def release(self, timeout):
  # Rearm the connection inactivity timer.
  self._setup_inactive_callback()
 
- async def wait_until_released(self):
+ async def wait_until_released(self) -> None:
  if self._in_use is None:
  return
  else:
  await self._in_use
 
- async def close(self):
+ async def close(self) -> None:
  if self._con is not None:
  # Connection.close() will call _release_on_close() to
  # finish holder cleanup.
  await self._con.close()
 
- def terminate(self):
+ def terminate(self) -> None:
  if self._con is not None:
  # Connection.terminate() will call _release_on_close() to
  # finish holder cleanup.
  self._con.terminate()
 
- def _setup_inactive_callback(self):
+ def _setup_inactive_callback(self) -> None:
  if self._inactive_callback is not None:
  raise exceptions.InternalClientError(
  'pool connection inactivity timer already exists')
@@ -261,12 +282,12 @@ def _setup_inactive_callback(self):
  self._inactive_callback = self._pool._loop.call_later(
  self._max_inactive_time, self._deactivate_inactive_connection)
 
- def _maybe_cancel_inactive_callback(self):
+ def _maybe_cancel_inactive_callback(self) -> None:
  if self._inactive_callback is not None:
  self._inactive_callback.cancel()
  self._inactive_callback = None
 
- def _deactivate_inactive_connection(self):
+ def _deactivate_inactive_connection(self) -> None:
  if self._in_use is not None:
  raise exceptions.InternalClientError(
  'attempting to deactivate an acquired connection')
@@ -280,12 +301,12 @@ def _deactivate_inactive_connection(self):
  # so terminate() above will not call the below.
  self._release_on_close()
 
- def _release_on_close(self):
+ def _release_on_close(self) -> None:
  self._maybe_cancel_inactive_callback()
  self._release()
  self._con = None
 
- def _release(self):
+ def _release(self) -> None:
  """Release this connection holder."""
  if self._in_use is None:
  # The holder is not checked out.
@@ -1012,7 +1033,7 @@ class PoolAcquireContext:
 
  __slots__ = ('timeout', 'connection', 'done', 'pool')
 
- def __init__(self, pool, timeout):
+ def __init__(self, pool: Pool, timeout: Optional[float]) -> None:
  self.pool = pool
  self.timeout = timeout
  self.connection = None
@@ -1024,7 +1045,12 @@ async def __aenter__(self):
  self.connection = await self.pool._acquire(self.timeout)
  return self.connection
 
- async def __aexit__(self, *exc):
+ async def __aexit__(
+ self,
+ exc_type: Optional[Type[BaseException]] = None,
+ exc_val: Optional[BaseException] = None,
+ exc_tb: Optional[TracebackType] = None,
+ ) -> None:
  self.done = True
  con = self.connection
  self.connection = None