Fix on type hints

libp2p/pubsub/pubsub_notifee.py

@@ -5,8 +5,11 @@ from typing import (
 
 from multiaddr import Multiaddr
 
-from libp2p.stream_muxer.mplex.mplex import (
-    Mplex,
+from libp2p.peer.id import (
+    ID,
+)
+from libp2p.stream_muxer.muxed_connection_interface import (
+    IMuxedConn,
 )
 from libp2p.network.notifee_interface import (
     INotifee,
@@ -22,10 +25,9 @@ from libp2p.network.stream.net_stream_interface import (
 class PubsubNotifee(INotifee):
     # pylint: disable=too-many-instance-attributes, cell-var-from-loop
 
-    # FIXME: Should be changed to type 'peer.ID'
-    initiator_peers_queue: asyncio.Queue[str]
+    initiator_peers_queue: asyncio.Queue[ID]
 
-    def __init__(self, initiator_peers_queue: asyncio.Queue[str]) -> None:
+    def __init__(self, initiator_peers_queue: asyncio.Queue[ID]) -> None:
         """
         :param initiator_peers_queue: queue to add new peers to so that pubsub
         can process new peers after we connect to them
@@ -38,7 +40,7 @@ class PubsubNotifee(INotifee):
     async def closed_stream(self, network: INetwork, stream: INetStream) -> None:
         pass
 
-    async def connected(self, network: INetwork, conn: Mplex) -> None:
+    async def connected(self, network: INetwork, conn: IMuxedConn) -> None:
         """
         Add peer_id to initiator_peers_queue, so that this peer_id can be used to
         create a stream and we only want to have one pubsub stream with each peer.
@@ -51,7 +53,7 @@ class PubsubNotifee(INotifee):
         if conn.initiator:
             await self.initiator_peers_queue.put(conn.peer_id)
 
-    async def disconnected(self, network: INetwork, conn: Mplex) -> None:
+    async def disconnected(self, network: INetwork, conn: IMuxedConn) -> None:
         pass
 
     async def listen(self, network: INetwork, multiaddr: Multiaddr) -> None: