Improve type completeness

Author: bryanforbes

asyncpg/cluster.py

@@ -724,6 +724,8 @@ def start(self, wait: int = 60, *,
 
 
 class RunningCluster(Cluster):
+    conn_spec: _ConnectionSpec
+
     def __init__(self, **kwargs: str) -> None:
         self.conn_spec = typing.cast(_ConnectionSpec, kwargs)
 

asyncpg/connect_utils.py

@@ -729,6 +729,12 @@ def _parse_connect_arguments(*, dsn: typing.Optional[str],
 
 
 class TLSUpgradeProto(asyncio.Protocol):
+    on_data: 'asyncio.Future[bool]'
+    host: str
+    port: int
+    ssl_context: ssl_module.SSLContext
+    ssl_is_advisory: typing.Optional[bool]
+
     def __init__(self, loop: typing.Optional[asyncio.AbstractEventLoop],
                  host: str, port: int, ssl_context: ssl_module.SSLContext,
                  ssl_is_advisory: typing.Optional[bool]) -> None:

asyncpg/connection.py

@@ -29,12 +29,12 @@
 from . import protocol
 from . import serverversion
 from . import transaction
+from . import types
 from . import utils
+from .protocol import protocol as _cprotocol
 
 if typing.TYPE_CHECKING:
-    from .protocol import protocol as _cprotocol
     from . import pool_connection_proxy as _pool
-    from . import types
 
 
 _Connection = typing.TypeVar('_Connection', bound='Connection[typing.Any]')
@@ -336,7 +336,7 @@ def get_server_pid(self) -> int:
         """Return the PID of the Postgres server the connection is bound to."""
         return self._protocol.get_server_pid()
 
-    def get_server_version(self) -> 'types.ServerVersion':
+    def get_server_version(self) -> types.ServerVersion:
         """Return the version of the connected PostgreSQL server.
 
         The returned value is a named tuple similar to that in
@@ -352,7 +352,7 @@ def get_server_version(self) -> 'types.ServerVersion':
         """
         return self._server_version
 
-    def get_settings(self) -> '_cprotocol.ConnectionSettings':
+    def get_settings(self) -> _cprotocol.ConnectionSettings:
         """Return connection settings.
 
         :return: :class:`~asyncpg.ConnectionSettings`.
@@ -3176,8 +3176,8 @@ class ServerCapabilities(typing.NamedTuple):
 
 
 def _detect_server_capabilities(
-    server_version: 'types.ServerVersion',
-    connection_settings: '_cprotocol.ConnectionSettings'
+    server_version: types.ServerVersion,
+    connection_settings: _cprotocol.ConnectionSettings
 ) -> ServerCapabilities:
     if hasattr(connection_settings, 'padb_revision'):
         # Amazon Redshift detected.

asyncpg/exceptions/_base.py

@@ -56,11 +56,13 @@ class UnknownPostgresError(FatalPostgresError):
 
 class InterfaceMessage:
     args: typing.Tuple[typing.Any, ...]
+    detail: typing.Optional[str]
+    hint: typing.Optional[str]
 
     def __init__(self, *, detail: typing.Optional[str] = None,
                  hint: typing.Optional[str] = None) -> None:
-        self.detail: typing.Optional[str] = detail
-        self.hint: typing.Optional[str] = hint
+        self.detail = detail
+        self.hint = hint
 
     def __str__(self) -> str:
         msg: str = self.args[0]
@@ -118,13 +120,17 @@ class ProtocolError(InternalClientError):
 class OutdatedSchemaCacheError(InternalClientError):
     """A value decoding error caused by a schema change before row fetching."""
 
+    schema_name: typing.Optional[str]
+    data_type_name: typing.Optional[str]
+    position: typing.Optional[str]
+
     def __init__(self, msg: str, *, schema: typing.Optional[str] = None,
                  data_type: typing.Optional[str] = None,
                  position: typing.Optional[str] = None) -> None:
         super().__init__(msg)
-        self.schema_name: typing.Optional[str] = schema
-        self.data_type_name: typing.Optional[str] = data_type
-        self.position: typing.Optional[str] = position
+        self.schema_name = schema
+        self.data_type_name = data_type
+        self.position = position
 
 
 class PostgresLogMessage(PostgresMessage):

asyncpg/pool.py

@@ -19,7 +19,7 @@
 from . import protocol
 
 
-logger = logging.getLogger(__name__)
+_logger = logging.getLogger(__name__)
 
 _Connection = typing.TypeVar(
     '_Connection',
@@ -994,10 +994,10 @@ async def close(self) -> None:
             self._closing = False
 
     def _warn_on_long_close(self) -> None:
-        logger.warning('Pool.close() is taking over 60 seconds to complete. '
-                       'Check if you have any unreleased connections left. '
-                       'Use asyncio.wait_for() to set a timeout for '
-                       'Pool.close().')
+        _logger.warning('Pool.close() is taking over 60 seconds to complete. '
+                        'Check if you have any unreleased connections left. '
+                        'Use asyncio.wait_for() to set a timeout for '
+                        'Pool.close().')
 
     def terminate(self) -> None:
         """Terminate all connections in the pool."""
@@ -1060,14 +1060,19 @@ class PoolAcquireContext(typing.Generic[_Record]):
 
     __slots__ = ('timeout', 'connection', 'done', 'pool')
 
+    timeout: typing.Optional[float]
+    connection: typing.Optional[
+        pool_connection_proxy.PoolConnectionProxy[_Record]
+    ]
+    done: bool
+    pool: Pool[_Record]
+
     def __init__(
         self, pool: Pool[_Record], timeout: typing.Optional[float]
     ) -> None:
         self.pool = pool
         self.timeout = timeout
-        self.connection: typing.Optional[
-            pool_connection_proxy.PoolConnectionProxy[_Record]
-        ] = None
+        self.connection = None
         self.done = False
 
     async def __aenter__(self) -> pool_connection_proxy.PoolConnectionProxy[

asyncpg/pool_connection_proxy.py

@@ -79,7 +79,7 @@ class PoolConnectionProxy(connection._ConnectionProxy[_Record],
     __slots__ = ('_con', '_holder')
 
     def __init__(self, holder: 'pool.PoolConnectionHolder[_Record]',
-                 con: connection.Connection):
+                 con: connection.Connection[_Record]):
         self._con = con
         self._holder = holder
         con._set_proxy(self)

asyncpg/serverversion.py

@@ -10,7 +10,7 @@
 from .compat import TypedDict
 from .types import ServerVersion
 
-version_regex = re.compile(
+_version_regex = re.compile(
     r"(Postgre[^\s]*)?\s*"
     r"(?P<major>[0-9]+)\.?"
     r"((?P<minor>[0-9]+)\.?)?"
@@ -29,7 +29,7 @@ class _VersionDict(TypedDict):
 
 
 def split_server_version_string(version_string: str) -> ServerVersion:
-    version_match = version_regex.search(version_string)
+    version_match = _version_regex.search(version_string)
 
     if version_match is None:
         raise ValueError(

asyncpg/types.py

@@ -69,6 +69,7 @@ def __gt__(self, other: '_RangeValue') -> bool:
 
 
 _V = typing.TypeVar('_V', bound=_RangeValue)
+_R = typing.TypeVar('_R', bound='Range[typing.Any]')
 
 
 class Range(typing.Generic[_V]):
@@ -208,4 +209,4 @@ def __repr__(self) -> str:
 
         return '<Range {}>'.format(desc)
 
-    __str__ = __repr__
+    __str__: typing.Callable[['Range[typing.Any]'], str] = __repr__