diff --git a/libp2p/pubsub/pubsub.py b/libp2p/pubsub/pubsub.py index 6e9d948..2d5625f 100644 --- a/libp2p/pubsub/pubsub.py +++ b/libp2p/pubsub/pubsub.py @@ -8,6 +8,7 @@ from typing import ( Dict, List, NamedTuple, + Set, Tuple, Union, cast, @@ -73,7 +74,7 @@ class Pubsub: my_topics: Dict[str, "asyncio.Queue[rpc_pb2.Message]"] - peer_topics: Dict[str, List[ID]] + peer_topics: Dict[str, Set[ID]] peers: Dict[ID, INetStream] topic_validators: Dict[str, TopicValidator] @@ -314,7 +315,7 @@ class Pubsub: for topic in self.peer_topics: if peer_id in self.peer_topics[topic]: - self.peer_topics[topic].remove(peer_id) + self.peer_topics[topic].discard(peer_id) self.router.remove_peer(peer_id) @@ -353,11 +354,11 @@ class Pubsub: self.peer_topics[sub_message.topicid] = [origin_id] elif origin_id not in self.peer_topics[sub_message.topicid]: # Add peer to topic - self.peer_topics[sub_message.topicid].append(origin_id) + self.peer_topics[sub_message.topicid].add(origin_id) else: if sub_message.topicid in self.peer_topics: if origin_id in self.peer_topics[sub_message.topicid]: - self.peer_topics[sub_message.topicid].remove(origin_id) + self.peer_topics[sub_message.topicid].discard(origin_id) # FIXME(mhchia): Change the function name? async def handle_talk(self, publish_message: rpc_pb2.Message) -> None: