feature/generic-connect-funcs

Author: glav-git

aio_pika/connection.py

@@ -23,6 +23,7 @@
 
 log = get_logger(__name__)
 T = TypeVar("T")
+ConnT = TypeVar("ConnT", bound=AbstractConnection)
 
 
 class Connection(AbstractConnection):
@@ -290,9 +291,9 @@ async def connect(
     ssl_context: Optional[SSLContext] = None,
     timeout: TimeoutType = None,
     client_properties: Optional[FieldTable] = None,
-    connection_class: Type[AbstractConnection] = Connection,
+    connection_class: Type[ConnT] = Connection,
     **kwargs: Any,
-) -> AbstractConnection:
+) -> ConnT:
     """ Make connection to the broker.
 
     Example:
@@ -373,7 +374,7 @@ async def main():
 
     """
 
-    connection: AbstractConnection = connection_class(
+    connection: ConnT = connection_class(
         make_url(
             url,
             host=host,

aio_pika/robust_connection.py

@@ -1,6 +1,6 @@
 import asyncio
 from ssl import SSLContext
-from typing import Any, Optional, Tuple, Type, Union
+from typing import Any, Optional, Tuple, Type, TypeVar, Union
 from weakref import WeakSet
 
 import aiormq.abc
@@ -20,6 +20,7 @@
 
 
 log = get_logger(__name__)
+ConnT = TypeVar("ConnT", bound=AbstractRobustConnection)
 
 
 class RobustConnection(Connection, AbstractRobustConnection):
@@ -245,9 +246,9 @@ async def connect_robust(
     ssl_context: Optional[SSLContext] = None,
     timeout: TimeoutType = None,
     client_properties: Optional[FieldTable] = None,
-    connection_class: Type[AbstractRobustConnection] = RobustConnection,
+    connection_class: Type[ConnT] = RobustConnection,
     **kwargs: Any,
-) -> AbstractRobustConnection:
+) -> ConnT:
 
     """Make connection to the broker.
 
@@ -330,7 +331,7 @@ async def main():
 
     """
 
-    connection: AbstractRobustConnection = connection_class(
+    connection: ConnT = connection_class(
         make_url(
             url,
             host=host,