Feat/fix return type for mux

libp2p/abc.py

@@ -1485,6 +1485,107 @@ async def listen_close(self, network: "INetwork", multiaddr: Multiaddr) -> None:
         """
 
 
+class IMultiselectCommunicator(ABC):
+    """
+    Communicator helper for multiselect.
+
+    Ensures that both the client and multistream module follow the same
+    multistream protocol.
+    """
+
+    @abstractmethod
+    async def write(self, msg_str: str) -> None:
+        """
+        Write a message to the stream.
+
+        Parameters
+        ----------
+        msg_str : str
+            The message string to write.
+
+        """
+
+    @abstractmethod
+    async def read(self) -> str:
+        """
+        Read a message from the stream until EOF.
+
+        Returns
+        -------
+        str
+            The message read from the stream.
+
+        """
+
+
+# -------------------------- multiselect_muxer interface.py --------------------------
+
+
+class IMultiselectMuxer(ABC):
+    """
+    Multiselect module for protocol negotiation.
+
+    Responsible for responding to a multiselect client by selecting a protocol
+    and its corresponding handler for communication.
+    """
+
+    handlers: dict[TProtocol | None, StreamHandlerFn | None]
+
+    @abstractmethod
+    def add_handler(self, protocol: TProtocol, handler: StreamHandlerFn) -> None:
+        """
+        Store a handler for the specified protocol.
+
+        Parameters
+        ----------
+        protocol : TProtocol
+            The protocol name.
+        handler : StreamHandlerFn
+            The handler function associated with the protocol.
+
+        """
+
+    @abstractmethod  # Ensure this is present if it was implicitly removed
+    def get_protocols(self) -> tuple[TProtocol, ...]:
+        """
+        Retrieve the protocols for which handlers have been registered.
+
+        Returns
+        -------
+        tuple[TProtocol, ...]
+            A tuple of registered protocol names.
+
+        """
+        # For an abstract method, the body might be empty or a simple `pass`.
+        # If it was `return tuple(self.handlers.keys())`, it should remain unchanged.
+        # The key is the type annotation above.
+        pass  # Or whatever the abstract method's body was
+
+    @abstractmethod
+    async def negotiate(
+        self, communicator: IMultiselectCommunicator
+    ) -> tuple[TProtocol | None, StreamHandlerFn | None]:
+        """
+        Negotiate a protocol selection with a multiselect client.
+
+        Parameters
+        ----------
+        communicator : IMultiselectCommunicator
+            The communicator used to negotiate the protocol.
+
+        Returns
+        -------
+        tuple[TProtocol, StreamHandlerFn]
+            A tuple containing the selected protocol and its handler.
+
+        Raises
+        ------
+        Exception
+            If negotiation fails.
+
+        """
+
+
 # -------------------------- host interface.py --------------------------
 
 
@@ -1545,15 +1646,14 @@ def get_network(self) -> INetworkService:
 
         """
 
-    # FIXME: Replace with correct return type
     @abstractmethod
-    def get_mux(self) -> Any:
+    def get_mux(self) -> IMultiselectMuxer:
         """
         Retrieve the muxer instance for the host.
 
         Returns
         -------
-        Any
+        IMultiselectMuxer
             The muxer instance of the host.
 
         """
@@ -2016,39 +2116,6 @@ def is_expired(self) -> bool:
 # ------------------ multiselect_communicator interface.py ------------------
 
 
-class IMultiselectCommunicator(ABC):
-    """
-    Communicator helper for multiselect.
-
-    Ensures that both the client and multistream module follow the same
-    multistream protocol.
-    """
-
-    @abstractmethod
-    async def write(self, msg_str: str) -> None:
-        """
-        Write a message to the stream.
-
-        Parameters
-        ----------
-        msg_str : str
-            The message string to write.
-
-        """
-
-    @abstractmethod
-    async def read(self) -> str:
-        """
-        Read a message from the stream until EOF.
-
-        Returns
-        -------
-        str
-            The message read from the stream.
-
-        """
-
-
 # -------------------------- multiselect_client interface.py --------------------------
 
 
@@ -2131,70 +2198,6 @@ async def try_select(
         """
 
 
-# -------------------------- multiselect_muxer interface.py --------------------------
-
-
-class IMultiselectMuxer(ABC):
-    """
-    Multiselect module for protocol negotiation.
-
-    Responsible for responding to a multiselect client by selecting a protocol
-    and its corresponding handler for communication.
-    """
-
-    handlers: dict[TProtocol | None, StreamHandlerFn | None]
-
-    @abstractmethod
-    def add_handler(self, protocol: TProtocol, handler: StreamHandlerFn) -> None:
-        """
-        Store a handler for the specified protocol.
-
-        Parameters
-        ----------
-        protocol : TProtocol
-            The protocol name.
-        handler : StreamHandlerFn
-            The handler function associated with the protocol.
-
-        """
-
-    def get_protocols(self) -> tuple[TProtocol | None, ...]:
-        """
-        Retrieve the protocols for which handlers have been registered.
-
-        Returns
-        -------
-        tuple[TProtocol, ...]
-            A tuple of registered protocol names.
-
-        """
-        return tuple(self.handlers.keys())
-
-    @abstractmethod
-    async def negotiate(
-        self, communicator: IMultiselectCommunicator
-    ) -> tuple[TProtocol | None, StreamHandlerFn | None]:
-        """
-        Negotiate a protocol selection with a multiselect client.
-
-        Parameters
-        ----------
-        communicator : IMultiselectCommunicator
-            The communicator used to negotiate the protocol.
-
-        Returns
-        -------
-        tuple[TProtocol, StreamHandlerFn]
-            A tuple containing the selected protocol and its handler.
-
-        Raises
-        ------
-        Exception
-            If negotiation fails.
-
-        """
-
-
 # -------------------------- routing interface.py --------------------------
 
 

libp2p/host/basic_host.py

@@ -16,6 +16,7 @@
 
 from libp2p.abc import (
     IHost,
+    IMultiselectMuxer,
     INetConn,
     INetStream,
     INetworkService,
@@ -130,7 +131,7 @@ def get_peerstore(self) -> IPeerStore:
         """
         return self.peerstore
 
-    def get_mux(self) -> Multiselect:
+    def get_mux(self) -> IMultiselectMuxer:
         """
         :return: mux instance of host
         """

libp2p/protocol_muxer/multiselect.py

@@ -48,6 +48,19 @@ def add_handler(
         """
         self.handlers[protocol] = handler
 
+    def get_protocols(self) -> tuple[TProtocol, ...]:
+        """
+        Retrieve the protocols for which handlers have been registered.
+
+        Returns
+        -------
+        tuple[TProtocol, ...]
+            A tuple of registered protocol names.
+
+        """
+        # Filter out None values, as they are not considered valid "protocols"
+        return tuple(p for p in self.handlers.keys() if p is not None)
+
     # FIXME: Make TProtocol Optional[TProtocol] to keep types consistent
     async def negotiate(
         self,

pyproject.toml

@@ -279,4 +279,4 @@ project_excludes = [
     "**/*pb2.py",
     "**/*.pyi",
     ".venv/**",
-]
+]

tests/conftest.py

@@ -1,6 +1,11 @@
 import pytest
 
+pytest_plugins = ["trio"]
+
 
 @pytest.fixture
 def security_protocol():
     return None
+
+
+pytest_plugins = ["trio"]