py-libp2p/libp2p/peer/peerstore.py

100 lines
3.3 KiB
Python
Raw Normal View History

2019-08-01 06:00:12 +08:00
from typing import Any, Dict, List, Optional, Sequence
2019-07-27 16:27:01 +08:00
from multiaddr import Multiaddr
from .id import ID
from .peerdata import PeerData, PeerDataError
from .peerinfo import PeerInfo
2019-07-28 22:30:51 +08:00
from .peerstore_interface import IPeerStore
2019-01-10 02:38:56 +08:00
class PeerStore(IPeerStore):
2019-07-27 16:27:01 +08:00
peer_map: Dict[ID, PeerData]
def __init__(self) -> None:
IPeerStore.__init__(self)
2018-10-22 03:15:53 +08:00
self.peer_map = {}
2019-07-27 16:27:01 +08:00
def __create_or_get_peer(self, peer_id: ID) -> PeerData:
2018-10-22 03:15:53 +08:00
"""
Returns the peer data for peer_id or creates a new peer data (and
stores it in peer_map) if peer data for peer_id does not yet exist.
2018-10-22 03:15:53 +08:00
:param peer_id: peer ID
:return: peer data
"""
if peer_id in self.peer_map:
return self.peer_map[peer_id]
data = PeerData()
self.peer_map[peer_id] = data
return self.peer_map[peer_id]
2019-07-27 16:27:01 +08:00
def peer_info(self, peer_id: ID) -> Optional[PeerInfo]:
2018-10-22 03:15:53 +08:00
if peer_id in self.peer_map:
2019-07-27 16:27:01 +08:00
peer_data = self.peer_map[peer_id]
return PeerInfo(peer_id, peer_data.addrs)
2018-10-22 03:15:53 +08:00
return None
2019-07-27 16:27:01 +08:00
def get_protocols(self, peer_id: ID) -> List[str]:
2018-10-22 03:15:53 +08:00
if peer_id in self.peer_map:
2018-10-29 20:13:39 +08:00
return self.peer_map[peer_id].get_protocols()
raise PeerStoreError("peer ID not found")
2019-07-27 16:27:01 +08:00
def add_protocols(self, peer_id: ID, protocols: Sequence[str]) -> None:
2018-10-22 03:15:53 +08:00
peer = self.__create_or_get_peer(peer_id)
2019-07-27 16:27:01 +08:00
peer.add_protocols(list(protocols))
2019-07-27 16:27:01 +08:00
def set_protocols(self, peer_id: ID, protocols: Sequence[str]) -> None:
peer = self.__create_or_get_peer(peer_id)
2019-07-27 16:27:01 +08:00
peer.set_protocols(list(protocols))
2019-07-27 16:27:01 +08:00
def peer_ids(self) -> List[ID]:
return list(self.peer_map.keys())
2019-07-30 15:31:02 +08:00
def get(self, peer_id: ID, key: str) -> Any:
2018-10-22 03:15:53 +08:00
if peer_id in self.peer_map:
try:
val = self.peer_map[peer_id].get_metadata(key)
except PeerDataError as error:
raise PeerStoreError(error)
2018-10-29 20:13:39 +08:00
return val
raise PeerStoreError("peer ID not found")
2019-07-30 15:31:02 +08:00
def put(self, peer_id: ID, key: str, val: Any) -> None:
2018-10-22 03:15:53 +08:00
# <<?>>
# This can output an error, not sure what the possible errors are
peer = self.__create_or_get_peer(peer_id)
peer.put_metadata(key, val)
2019-07-27 16:27:01 +08:00
def add_addr(self, peer_id: ID, addr: Multiaddr, ttl: int) -> None:
self.add_addrs(peer_id, [addr], ttl)
2019-07-27 16:27:01 +08:00
def add_addrs(self, peer_id: ID, addrs: Sequence[Multiaddr], ttl: int) -> None:
2018-10-22 03:15:53 +08:00
# Ignore ttl for now
peer = self.__create_or_get_peer(peer_id)
2019-07-27 16:27:01 +08:00
peer.add_addrs(list(addrs))
2019-07-27 16:27:01 +08:00
def addrs(self, peer_id: ID) -> List[Multiaddr]:
2018-10-22 03:15:53 +08:00
if peer_id in self.peer_map:
2018-10-29 20:13:39 +08:00
return self.peer_map[peer_id].get_addrs()
raise PeerStoreError("peer ID not found")
2019-07-27 16:27:01 +08:00
def clear_addrs(self, peer_id: ID) -> None:
2018-10-22 03:15:53 +08:00
# Only clear addresses if the peer is in peer map
if peer_id in self.peer_map:
self.peer_map[peer_id].clear_addrs()
2019-07-27 16:27:01 +08:00
def peers_with_addrs(self) -> List[ID]:
2018-10-22 03:15:53 +08:00
# Add all peers with addrs at least 1 to output
2019-07-27 16:27:01 +08:00
output: List[ID] = []
2019-07-27 16:27:01 +08:00
for peer_id in self.peer_map:
if len(self.peer_map[peer_id].get_addrs()) >= 1:
output.append(peer_id)
2018-10-22 03:15:53 +08:00
return output
2018-10-30 07:03:19 +08:00
2019-01-10 02:38:56 +08:00
2018-10-30 07:03:19 +08:00
class PeerStoreError(KeyError):
"""Raised when peer ID is not found in peer store."""