Add type hint to pubsub notifee/interface

libp2p/pubsub/pubsub_notifee.py

@@ -1,23 +1,45 @@
-from libp2p.network.notifee_interface import INotifee
+from typing import (
+    List,
+    Sequence,
+)
 
+from multiaddr import Multiaddr
+
+from libp2p.network.connection.raw_connection import (
+    RawConnection,
+)
+from libp2p.network.notifee_interface import (
+    INotifee,
+)
+from libp2p.network.network_interface import (
+    INetwork,
+)
+from libp2p.network.stream.net_stream_interface import (
+    INetStream,
+)
+from libp2p.peer.id import (
+    ID,
+)
 
 class PubsubNotifee(INotifee):
     # pylint: disable=too-many-instance-attributes, cell-var-from-loop
 
-    def __init__(self, initiator_peers_queue):
+    initiator_peers_queue: List[ID]
+
+    def __init__(self, initiator_peers_queue: Sequence[ID]) -> None:
         """
         :param initiator_peers_queue: queue to add new peers to so that pubsub
         can process new peers after we connect to them
         """
-        self.initiator_peers_queue = initiator_peers_queue
+        self.initiator_peers_queue = List(initiator_peers_queue)
 
-    async def opened_stream(self, network, stream):
+    async def opened_stream(self, network: INetwork, stream: INetStream) -> None:
         pass
 
-    async def closed_stream(self, network, stream):
+    async def closed_stream(self, network: INetwork, stream: INetStream) -> None:
         pass
 
-    async def connected(self, network, conn):
+    async def connected(self, network: INetwork, conn: RawConnection) -> 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.
@@ -30,11 +52,11 @@ class PubsubNotifee(INotifee):
         if conn.initiator:
             await self.initiator_peers_queue.put(conn.peer_id)
 
-    async def disconnected(self, network, conn):
+    async def disconnected(self, network: INetwork, conn: RawConnection) -> None:
         pass
 
-    async def listen(self, network, multiaddr):
+    async def listen(self, network: INetwork, multiaddr: Multiaddr) -> None:
         pass
 
-    async def listen_close(self, network, multiaddr):
+    async def listen_close(self, network: INetwork, multiaddr: Multiaddr) -> None:
         pass

libp2p/pubsub/pubsub_router_interface.py

@@ -1,15 +1,24 @@
 from abc import ABC, abstractmethod
+from typing import (
+    List,
+)
+
+from .pb import rpc_pb2
+from .pubsub import Pubsub
+from libp2p.peer.id import (
+    ID,
+)
 
 class IPubsubRouter(ABC):
 
     @abstractmethod
-    def get_protocols(self):
+    def get_protocols(self) -> List[str]:
         """
         :return: the list of protocols supported by the router
         """
 
     @abstractmethod
-    def attach(self, pubsub):
+    def attach(self, pubsub: Pubsub) -> None:
         """
         Attach is invoked by the PubSub constructor to attach the router to a
         freshly initialized PubSub instance.
@@ -17,21 +26,21 @@ class IPubsubRouter(ABC):
         """
 
     @abstractmethod
-    def add_peer(self, peer_id, protocol_id):
+    def add_peer(self, peer_id: ID, protocol_id: str) -> None:
         """
         Notifies the router that a new peer has been connected
         :param peer_id: id of peer to add
         """
 
     @abstractmethod
-    def remove_peer(self, peer_id):
+    def remove_peer(self, peer_id: ID) -> None:
         """
         Notifies the router that a peer has been disconnected
         :param peer_id: id of peer to remove
         """
 
     @abstractmethod
-    def handle_rpc(self, rpc, sender_peer_id):
+    def handle_rpc(self, rpc: rpc_pb2.ControlMessage, sender_peer_id: ID) -> None:
         """
         Invoked to process control messages in the RPC envelope.
         It is invoked after subscriptions and payload messages have been processed
@@ -42,7 +51,7 @@ class IPubsubRouter(ABC):
         """
 
     @abstractmethod
-    async def publish(self, msg_forwarder, pubsub_msg):
+    async def publish(self, msg_forwarder: ID, pubsub_msg: rpc_pb2.Message):
         """
         Invoked to forward a new message that has been validated
         :param msg_forwarder: peer_id of message sender
@@ -50,7 +59,7 @@ class IPubsubRouter(ABC):
         """
 
     @abstractmethod
-    def join(self, topic):
+    def join(self, topic: str) -> None:
         """
         Join notifies the router that we want to receive and
         forward messages in a topic. It is invoked after the
@@ -59,7 +68,7 @@ class IPubsubRouter(ABC):
         """
 
     @abstractmethod
-    def leave(self, topic):
+    def leave(self, topic: str) -> None:
         """
         Leave notifies the router that we are no longer interested in a topic.
         It is invoked after the unsubscription announcement.