add typing to protocol_muxer

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