diff --git a/libp2p/protocol_muxer/multiselect.py b/libp2p/protocol_muxer/multiselect.py index 6312e272..9a55ec39 100644 --- a/libp2p/protocol_muxer/multiselect.py +++ b/libp2p/protocol_muxer/multiselect.py @@ -1,7 +1,6 @@ from typing import Dict, Tuple -from libp2p.stream_muxer.abc import IMuxedStream -from libp2p.typing import StreamHandlerFn, TProtocol +from libp2p.typing import NegotiableTransport, StreamHandlerFn, TProtocol from .multiselect_communicator import MultiselectCommunicator from .multiselect_communicator_interface import IMultiselectCommunicator @@ -31,7 +30,7 @@ class Multiselect(IMultiselectMuxer): """ self.handlers[protocol] = handler - async def negotiate(self, stream: IMuxedStream) -> Tuple[TProtocol, StreamHandlerFn]: + async def negotiate(self, stream: NegotiableTransport) -> Tuple[TProtocol, StreamHandlerFn]: """ Negotiate performs protocol selection :param stream: stream to negotiate on diff --git a/libp2p/protocol_muxer/multiselect_client.py b/libp2p/protocol_muxer/multiselect_client.py index 8bfaabfe..b9b3240b 100644 --- a/libp2p/protocol_muxer/multiselect_client.py +++ b/libp2p/protocol_muxer/multiselect_client.py @@ -1,7 +1,7 @@ from typing import Sequence from libp2p.stream_muxer.abc import IMuxedStream -from libp2p.typing import TProtocol +from libp2p.typing import NegotiableTransport, TProtocol from .multiselect_client_interface import IMultiselectClient from .multiselect_communicator import MultiselectCommunicator @@ -59,7 +59,7 @@ class MultiselectClient(IMultiselectClient): return selected_protocol async def select_one_of( - self, protocols: Sequence[TProtocol], stream: IMuxedStream + self, protocols: Sequence[TProtocol], stream: NegotiableTransport ) -> TProtocol: """ For each protocol, send message to multiselect selecting protocol diff --git a/libp2p/protocol_muxer/multiselect_communicator.py b/libp2p/protocol_muxer/multiselect_communicator.py index 5aec962c..36513fcd 100644 --- a/libp2p/protocol_muxer/multiselect_communicator.py +++ b/libp2p/protocol_muxer/multiselect_communicator.py @@ -1,4 +1,4 @@ -from libp2p.stream_muxer.abc import IMuxedStream +from libp2p.typing import NegotiableTransport from .multiselect_communicator_interface import IMultiselectCommunicator @@ -10,9 +10,9 @@ class MultiselectCommunicator(IMultiselectCommunicator): which is necessary for them to work """ - reader_writer: IMuxedStream + reader_writer: NegotiableTransport - def __init__(self, reader_writer: IMuxedStream) -> None: + def __init__(self, reader_writer: NegotiableTransport) -> None: """ MultistreamCommunicator expects a reader_writer object that has an async read and an async write function (this could be a stream, diff --git a/libp2p/protocol_muxer/multiselect_muxer_interface.py b/libp2p/protocol_muxer/multiselect_muxer_interface.py index edacc6ba..54446460 100644 --- a/libp2p/protocol_muxer/multiselect_muxer_interface.py +++ b/libp2p/protocol_muxer/multiselect_muxer_interface.py @@ -1,8 +1,7 @@ from abc import ABC, abstractmethod from typing import Dict, Tuple -from libp2p.stream_muxer.abc import IMuxedStream -from libp2p.typing import StreamHandlerFn, TProtocol +from libp2p.typing import NegotiableTransport, StreamHandlerFn, TProtocol class IMultiselectMuxer(ABC): @@ -23,7 +22,7 @@ class IMultiselectMuxer(ABC): """ @abstractmethod - async def negotiate(self, stream: IMuxedStream) -> Tuple[TProtocol, StreamHandlerFn]: + async def negotiate(self, stream: NegotiableTransport) -> Tuple[TProtocol, StreamHandlerFn]: """ Negotiate performs protocol selection :param stream: stream to negotiate on diff --git a/libp2p/typing.py b/libp2p/typing.py index 72d0d2dc..4b5f5f23 100644 --- a/libp2p/typing.py +++ b/libp2p/typing.py @@ -1,6 +1,11 @@ -from typing import Awaitable, Callable, NewType +from typing import Awaitable, Callable, NewType, Union +from libp2p.network.connection.raw_connection_interface import IRawConnection from libp2p.network.stream.net_stream_interface import INetStream +from libp2p.stream_muxer.abc import IMuxedStream TProtocol = NewType("TProtocol", str) StreamHandlerFn = Callable[[INetStream], Awaitable[None]] + + +NegotiableTransport = Union[IMuxedStream, IRawConnection]