diff --git a/libp2p/pubsub/pubsub_notifee.py b/libp2p/pubsub/pubsub_notifee.py index 4173bd8f..1fa3b416 100644 --- a/libp2p/pubsub/pubsub_notifee.py +++ b/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 diff --git a/libp2p/pubsub/pubsub_router_interface.py b/libp2p/pubsub/pubsub_router_interface.py index 8819e5f0..e0787e27 100644 --- a/libp2p/pubsub/pubsub_router_interface.py +++ b/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.