Updates for custom records and method updates

MagicStack · Aug 21, 2020 · 8332c03 · 8332c03
1 parent 86a92f9
commit 8332c03
Show file tree

Hide file tree

Showing 11 changed files with 949 additions and 190 deletions.
diff --git a/asyncpg/cluster.py b/asyncpg/cluster.py
@@ -130,11 +130,15 @@ def get_status(self) -> str:
 
     async def connect(self,
                       loop: typing.Optional[asyncio.AbstractEventLoop] = None,
-                      **kwargs: typing.Any) -> 'connection.Connection':
+                      **kwargs: typing.Any) \
+            -> 'connection.Connection[typing.Any]':
         conn_info = typing.cast(typing.Dict[str, typing.Any],
                                 self.get_connection_spec())
         conn_info.update(kwargs)
-        return await asyncpg.connect(loop=loop, **conn_info)
+        return typing.cast(
+            'connection.Connection[typing.Any]',
+            await asyncpg.connect(loop=loop, **conn_info)
+        )
 
     def init(self, **settings: str) -> str:
         """Initialize cluster."""

diff --git a/asyncpg/connect_utils.py b/asyncpg/connect_utils.py
@@ -29,6 +29,7 @@
 
 _Connection = typing.TypeVar('_Connection')
 _Protocol = typing.TypeVar('_Protocol', bound=asyncio.Protocol)
+_Record = typing.TypeVar('_Record', bound=protocol.Record)
 
 _TPTupleType = typing.Tuple[asyncio.WriteTransport, _Protocol]
 AddrType = typing.Union[typing.Tuple[str, int], str]
@@ -654,7 +655,7 @@ async def _connect_addr(
     params: _ConnectionParameters,
     config: _ClientConfiguration,
     connection_class: typing.Type[_Connection],
-    record_class: typing.Any
+    record_class: typing.Type[_Record]
 ) -> _Connection:
     assert loop is not None
 
@@ -680,7 +681,7 @@ async def _connect_addr(
         assert not params.ssl
         connector = typing.cast(
             typing.Coroutine[typing.Any, None,
-                             _TPTupleType[protocol.Protocol]],
+                             _TPTupleType['protocol.Protocol[_Record]']],
             loop.create_unix_connection(proto_factory, addr))
     elif params.ssl:
         connector = _create_ssl_connection(
@@ -689,7 +690,7 @@ async def _connect_addr(
     else:
         connector = typing.cast(
             typing.Coroutine[typing.Any, None,
-                             _TPTupleType[protocol.Protocol]],
+                             _TPTupleType['protocol.Protocol[_Record]']],
             loop.create_connection(proto_factory, *addr))
 
     connector_future = asyncio.ensure_future(connector)
@@ -721,7 +722,7 @@ async def _connect(
     loop: typing.Optional[asyncio.AbstractEventLoop],
     timeout: float,
     connection_class: typing.Type[_Connection],
-    record_class: typing.Any,
+    record_class: typing.Type[_Record],
     **kwargs: typing.Any
 ) -> _Connection:
     if loop is None: