add typing to protocol_muxer

libp2p/host/basic_host.py

@@ -8,6 +8,7 @@ from libp2p.peer.id import ID
 from libp2p.peer.peerinfo import PeerInfo
 from libp2p.peer.peerstore_interface import IPeerStore
 from libp2p.routing.kademlia.kademlia_peer_router import KadmeliaPeerRouter
+from libp2p.typing import StreamHandlerFn
 
 from .host_interface import IHost
 
@@ -17,9 +18,6 @@ from .host_interface import IHost
 # telling it to listen on the given listen addresses.
 
 
-StreamHandlerFn = Callable[[INetStream], Awaitable[None]]
-
-
 class BasicHost(IHost):
 
     _network: INetwork

libp2p/host/host_interface.py

@@ -7,8 +7,7 @@ from libp2p.network.network_interface import INetwork
 from libp2p.network.stream.net_stream_interface import INetStream
 from libp2p.peer.id import ID
 from libp2p.peer.peerinfo import PeerInfo
-
-StreamHandlerFn = Callable[[INetStream], Awaitable[None]]
+from libp2p.typing import StreamHandlerFn
 
 
 class IHost(ABC):

libp2p/network/network_interface.py

@@ -7,6 +7,7 @@ from libp2p.peer.id import ID
 from libp2p.peer.peerstore import PeerStore
 from libp2p.stream_muxer.abc import IMuxedConn
 from libp2p.transport.listener_interface import IListener
+from libp2p.typing import StreamHandlerFn
 
 from .stream.net_stream_interface import INetStream
 
@@ -14,9 +15,6 @@ if TYPE_CHECKING:
     from .notifee_interface import INotifee
 
 
-StreamHandlerFn = Callable[[INetStream], Awaitable[None]]
-
-
 class INetwork(ABC):
 
     peerstore: PeerStore

libp2p/network/swarm.py

@@ -12,6 +12,7 @@ from libp2p.stream_muxer.abc import IMuxedConn, IMuxedStream
 from libp2p.transport.listener_interface import IListener
 from libp2p.transport.transport_interface import ITransport
 from libp2p.transport.upgrader import TransportUpgrader
+from libp2p.typing import StreamHandlerFn
 
 from .connection.raw_connection import RawConnection
 from .network_interface import INetwork
@@ -20,8 +21,6 @@ from .stream.net_stream import NetStream
 from .stream.net_stream_interface import INetStream
 from .typing import GenericProtocolHandlerFn
 
-StreamHandlerFn = Callable[[INetStream], Awaitable[None]]
-
 
 class Swarm(INetwork):
 

libp2p/protocol_muxer/multiselect.py

@@ -1,4 +1,10 @@
+from typing import Dict, Tuple, TypeVar
+
+from libp2p.stream_muxer.abc import IMuxedStream
+from libp2p.typing import StreamHandlerFn, TProtocol
+
 from .multiselect_communicator import MultiselectCommunicator
+from .multiselect_communicator_interface import IMultiselectCommunicator
 from .multiselect_muxer_interface import IMultiselectMuxer
 
 MULTISELECT_PROTOCOL_ID = "/multistream/1.0.0"
@@ -12,10 +18,12 @@ class Multiselect(IMultiselectMuxer):
     a specific protocol and handler pair to use for communication
     """
 
-    def __init__(self):
+    handlers: Dict[TProtocol, StreamHandlerFn]
+
+    def __init__(self) -> None:
         self.handlers = {}
 
-    def add_handler(self, protocol, handler):
+    def add_handler(self, protocol: TProtocol, handler: StreamHandlerFn) -> None:
         """
         Store the handler with the given protocol
         :param protocol: protocol name
@@ -23,7 +31,7 @@ class Multiselect(IMultiselectMuxer):
         """
         self.handlers[protocol] = handler
 
-    async def negotiate(self, stream):
+    async def negotiate(self, stream: IMuxedStream) -> Tuple[TProtocol, StreamHandlerFn]:
         """
         Negotiate performs protocol selection
         :param stream: stream to negotiate on
@@ -46,7 +54,7 @@ class Multiselect(IMultiselectMuxer):
                 # TODO: handle ls command
                 pass
             else:
-                protocol = command
+                protocol = TProtocol(command)
                 if protocol in self.handlers:
                     # Tell counterparty we have decided on a protocol
                     await communicator.write(protocol)
@@ -56,7 +64,7 @@ class Multiselect(IMultiselectMuxer):
                 # Tell counterparty this protocol was not found
                 await communicator.write(PROTOCOL_NOT_FOUND_MSG)
 
-    async def handshake(self, communicator):
+    async def handshake(self, communicator: IMultiselectCommunicator) -> None:
         """
         Perform handshake to agree on multiselect protocol
         :param communicator: communicator to use
@@ -78,7 +86,7 @@ class Multiselect(IMultiselectMuxer):
         # Handshake succeeded if this point is reached
 
 
-def validate_handshake(handshake_contents):
+def validate_handshake(handshake_contents: str) -> bool:
     """
     Determine if handshake is valid and should be confirmed
     :param handshake_contents: contents of handshake message

libp2p/protocol_muxer/multiselect_client.py

@@ -1,5 +1,11 @@
+from typing import Sequence
+
+from libp2p.stream_muxer.abc import IMuxedStream
+from libp2p.typing import TProtocol
+
 from .multiselect_client_interface import IMultiselectClient
 from .multiselect_communicator import MultiselectCommunicator
+from .multiselect_communicator_interface import IMultiselectCommunicator
 
 MULTISELECT_PROTOCOL_ID = "/multistream/1.0.0"
 PROTOCOL_NOT_FOUND_MSG = "na"
@@ -11,10 +17,7 @@ class MultiselectClient(IMultiselectClient):
     module in order to select a protocol id to communicate over
     """
 
-    def __init__(self):
-        pass
-
-    async def handshake(self, communicator):
+    async def handshake(self, communicator: IMultiselectCommunicator) -> None:
         """
         Ensure that the client and multiselect
         are both using the same multiselect protocol
@@ -36,7 +39,7 @@ class MultiselectClient(IMultiselectClient):
 
         # Handshake succeeded if this point is reached
 
-    async def select_protocol_or_fail(self, protocol, stream):
+    async def select_protocol_or_fail(self, protocol: TProtocol, stream: IMuxedStream) -> TProtocol:
         """
         Send message to multiselect selecting protocol
         and fail if multiselect does not return same protocol
@@ -55,7 +58,9 @@ class MultiselectClient(IMultiselectClient):
 
         return selected_protocol
 
-    async def select_one_of(self, protocols, stream):
+    async def select_one_of(
+        self, protocols: Sequence[TProtocol], stream: IMuxedStream
+    ) -> TProtocol:
         """
         For each protocol, send message to multiselect selecting protocol
         and fail if multiselect does not return same protocol. Returns first
@@ -83,7 +88,9 @@ class MultiselectClient(IMultiselectClient):
         # No protocols were found, so return no protocols supported error
         raise MultiselectClientError("protocols not supported")
 
-    async def try_select(self, communicator, protocol):
+    async def try_select(
+        self, communicator: IMultiselectCommunicator, protocol: TProtocol
+    ) -> TProtocol:
         """
         Try to select the given protocol or raise exception if fails
         :param communicator: communicator to use to communicate with counterparty
@@ -106,7 +113,7 @@ class MultiselectClient(IMultiselectClient):
         raise MultiselectClientError("unrecognized response: " + response)
 
 
-def validate_handshake(handshake_contents):
+def validate_handshake(handshake_contents: str) -> bool:
     """
     Determine if handshake is valid and should be confirmed
     :param handshake_contents: contents of handshake message

libp2p/protocol_muxer/multiselect_client_interface.py

@@ -1,4 +1,8 @@
 from abc import ABC, abstractmethod
+from typing import Sequence
+
+from libp2p.stream_muxer.abc import IMuxedStream
+from libp2p.typing import TProtocol
 
 
 class IMultiselectClient(ABC):
@@ -8,7 +12,7 @@ class IMultiselectClient(ABC):
     """
 
     @abstractmethod
-    def select_protocol_or_fail(self, protocol, stream):
+    async def select_protocol_or_fail(self, protocol: TProtocol, stream: IMuxedStream) -> TProtocol:
         """
         Send message to multiselect selecting protocol
         and fail if multiselect does not return same protocol
@@ -18,7 +22,9 @@ class IMultiselectClient(ABC):
         """
 
     @abstractmethod
-    def select_one_of(self, protocols, stream):
+    async def select_one_of(
+        self, protocols: Sequence[TProtocol], stream: IMuxedStream
+    ) -> TProtocol:
         """
         For each protocol, send message to multiselect selecting protocol
         and fail if multiselect does not return same protocol. Returns first

libp2p/protocol_muxer/multiselect_communicator.py

@@ -1,3 +1,5 @@
+from libp2p.stream_muxer.abc import IMuxedStream
+
 from .multiselect_communicator_interface import IMultiselectCommunicator
 
 
@@ -8,7 +10,9 @@ class MultiselectCommunicator(IMultiselectCommunicator):
     which is necessary for them to work
     """
 
-    def __init__(self, reader_writer):
+    reader_writer: IMuxedStream
+
+    def __init__(self, reader_writer: IMuxedStream) -> None:
         """
         MultistreamCommunicator expects a reader_writer object that has
         an async read and an async write function (this could be a stream,
@@ -16,14 +20,14 @@ class MultiselectCommunicator(IMultiselectCommunicator):
         """
         self.reader_writer = reader_writer
 
-    async def write(self, msg_str):
+    async def write(self, msg_str: str) -> None:
         """
         Write message to reader_writer
         :param msg_str: message to write
         """
         await self.reader_writer.write(msg_str.encode())
 
-    async def read_stream_until_eof(self):
+    async def read_stream_until_eof(self) -> str:
         """
         Reads message from reader_writer until EOF
         """

libp2p/protocol_muxer/multiselect_communicator_interface.py

@@ -9,14 +9,14 @@ class IMultiselectCommunicator(ABC):
     """
 
     @abstractmethod
-    def write(self, msg_str):
+    async def write(self, msg_str: str) -> None:
         """
         Write message to stream
         :param msg_str: message to write
         """
 
     @abstractmethod
-    def read_stream_until_eof(self):
+    async def read_stream_until_eof(self) -> str:
         """
         Reads message from stream until EOF
         """

libp2p/protocol_muxer/multiselect_muxer_interface.py

@@ -1,4 +1,8 @@
 from abc import ABC, abstractmethod
+from typing import Dict, Tuple, TypeVar
+
+from libp2p.stream_muxer.abc import IMuxedStream
+from libp2p.typing import StreamHandlerFn, TProtocol
 
 
 class IMultiselectMuxer(ABC):
@@ -8,8 +12,10 @@ class IMultiselectMuxer(ABC):
     a specific protocol and handler pair to use for communication
     """
 
+    handlers: Dict[TProtocol, StreamHandlerFn]
+
     @abstractmethod
-    def add_handler(self, protocol, handler):
+    def add_handler(self, protocol: TProtocol, handler: StreamHandlerFn) -> None:
         """
         Store the handler with the given protocol
         :param protocol: protocol name
@@ -17,7 +23,7 @@ class IMultiselectMuxer(ABC):
         """
 
     @abstractmethod
-    def negotiate(self, stream):
+    async def negotiate(self, stream: IMuxedStream) -> Tuple[TProtocol, StreamHandlerFn]:
         """
         Negotiate performs protocol selection
         :param stream: stream to negotiate on

libp2p/security/security_multistream.py

@@ -1,17 +1,14 @@
 from abc import ABC
-from typing import TYPE_CHECKING, Dict, NewType
+from typing import Dict, NewType
 
+from libp2p.network.connection.raw_connection_interface import IRawConnection
+from libp2p.peer.id import ID
 from libp2p.protocol_muxer.multiselect import Multiselect
 from libp2p.protocol_muxer.multiselect_client import MultiselectClient
+from libp2p.typing import TProtocol
 
-if TYPE_CHECKING:
-    from libp2p.network.connection.raw_connection_interface import IRawConnection
-    from libp2p.peer.id import ID
-    from .secure_conn_interface import ISecureConn
-    from .secure_transport_interface import ISecureTransport
-
-
-TProtocol = NewType("TProtocol", str)
+from .secure_conn_interface import ISecureConn
+from .secure_transport_interface import ISecureTransport
 
 
 """

libp2p/transport/upgrader.py

@@ -5,8 +5,9 @@ from libp2p.network.typing import GenericProtocolHandlerFn
 from libp2p.peer.id import ID
 from libp2p.security.secure_conn_interface import ISecureConn
 from libp2p.security.secure_transport_interface import ISecureTransport
-from libp2p.security.security_multistream import SecurityMultistream, TProtocol
+from libp2p.security.security_multistream import SecurityMultistream
 from libp2p.stream_muxer.mplex.mplex import Mplex
+from libp2p.typing import TProtocol
 
 from .listener_interface import IListener
 from .transport_interface import ITransport

libp2p/typing.py

@@ -0,0 +1,6 @@
+from typing import Awaitable, Callable, NewType
+
+from libp2p.network.stream.net_stream_interface import INetStream
+
+TProtocol = NewType("TProtocol", str)
+StreamHandlerFn = Callable[[INetStream], Awaitable[None]]