add typing to security

libp2p/security/security_multistream.py

@@ -2,6 +2,18 @@ from abc import ABC
 from libp2p.protocol_muxer.multiselect_client import MultiselectClient
 from libp2p.protocol_muxer.multiselect import Multiselect
 
+from typing import TYPE_CHECKING, NewType, Dict
+
+if TYPE_CHECKING:
+    from libp2p.network.connection.raw_connection_interface import IRawConnection
+    from libp2p.peer.id import ID
+    from .typing import TSecurityDetails
+    from .secure_conn_interface import ISecureConn
+    from .secure_transport_interface import ISecureTransport
+
+
+TProtocol = NewType("TProtocol", str)
+
 # pylint: disable=W0105
 
 """
@@ -13,7 +25,11 @@ Relevant go repo: https://github.com/libp2p/go-conn-security/blob/master/interfa
 
 
 class SecurityMultistream(ABC):
-    def __init__(self):
+    transports: Dict[TProtocol, "ISecureTransport"]
+    multiselect: "Multiselect"
+    multiselect_client: "MultiselectClient"
+
+    def __init__(self) -> None:
         # Map protocol to secure transport
         self.transports = {}
 
@@ -23,7 +39,7 @@ class SecurityMultistream(ABC):
         # Create multiselect client
         self.multiselect_client = MultiselectClient()
 
-    def add_transport(self, protocol, transport):
+    def add_transport(self, protocol: TProtocol, transport: "ISecureTransport") -> None:
         # Associate protocol with transport
         self.transports[protocol] = transport
 
@@ -32,7 +48,7 @@ class SecurityMultistream(ABC):
         # we only care about selecting the protocol, not any handler function
         self.multiselect.add_handler(protocol, None)
 
-    async def secure_inbound(self, conn):
+    async def secure_inbound(self, conn: "IRawConnection") -> "ISecureConn":
         """
         Secure the connection, either locally or by communicating with opposing node via conn,
         for an inbound connection (i.e. we are not the initiator)
@@ -47,7 +63,9 @@ class SecurityMultistream(ABC):
 
         return secure_conn
 
-    async def secure_outbound(self, conn, peer_id):
+    async def secure_outbound(
+        self, conn: "IRawConnection", peer_id: "ID"
+    ) -> "ISecureConn":
         """
         Secure the connection, either locally or by communicating with opposing node via conn,
         for an inbound connection (i.e. we are the initiator)
@@ -62,7 +80,9 @@ class SecurityMultistream(ABC):
 
         return secure_conn
 
-    async def select_transport(self, conn, initiator):
+    async def select_transport(
+        self, conn: "IRawConnection", initiator: bool
+    ) -> "ISecureTransport":
         """
         Select a transport that both us and the node on the
         other end of conn support and agree on