diff --git a/libp2p/peer/addrbook_interface.py b/libp2p/peer/addrbook_interface.py index 914d937..5ac34f0 100644 --- a/libp2p/peer/addrbook_interface.py +++ b/libp2p/peer/addrbook_interface.py @@ -1,13 +1,22 @@ from abc import ABC, abstractmethod +from typing import ( + List, + Sequence, +) + + +from multiaddr import Multiaddr + +from .id import ID class IAddrBook(ABC): - def __init__(self): + def __init__(self) -> None: pass @abstractmethod - def add_addr(self, peer_id, addr, ttl): + def add_addr(self, peer_id: ID, addr: Multiaddr, ttl: int) -> None: """ Calls add_addrs(peer_id, [addr], ttl) :param peer_id: the peer to add address for @@ -16,7 +25,7 @@ class IAddrBook(ABC): """ @abstractmethod - def add_addrs(self, peer_id, addrs, ttl): + def add_addrs(self, peer_id: ID, addrs: Sequence[Multiaddr], ttl: int) -> None: """ Adds addresses for a given peer all with the same time-to-live. If one of the addresses already exists for the peer and has a longer TTL, no operation should take place. @@ -27,21 +36,21 @@ class IAddrBook(ABC): """ @abstractmethod - def addrs(self, peer_id): + def addrs(self, peer_id: ID) -> List[Multiaddr]: """ :param peer_id: peer to get addresses of :return: all known (and valid) addresses for the given peer """ @abstractmethod - def clear_addrs(self, peer_id): + def clear_addrs(self, peer_id: ID) -> None: """ Removes all previously stored addresses :param peer_id: peer to remove addresses of """ @abstractmethod - def peers_with_addrs(self): + def peers_with_addrs(self) -> List[ID]: """ :return: all of the peer IDs stored with addresses """ diff --git a/libp2p/peer/id.py b/libp2p/peer/id.py index 77e2f87..578f67f 100644 --- a/libp2p/peer/id.py +++ b/libp2p/peer/id.py @@ -1,3 +1,10 @@ +from typing import ( + Union, +) + +from Crypto.PublicKey.RSA import ( + RsaKey, +) import hashlib import base58 import multihash @@ -13,67 +20,71 @@ MAX_INLINE_KEY_LENGTH = 42 class ID: - def __init__(self, id_str): + _id_str: str + + def __init__(self, id_str: str) -> None: self._id_str = id_str def to_bytes(self) -> bytes: return self._id_str - def get_raw_id(self): + def get_raw_id(self) -> str: return self._id_str - def pretty(self): + def pretty(self) -> str: return base58.b58encode(self._id_str).decode() - def get_xor_id(self): + def get_xor_id(self) -> int: return int(digest(self.get_raw_id()).hex(), 16) - def __str__(self): + def __str__(self) -> str: pid = self.pretty() return pid __repr__ = __str__ - def __eq__(self, other): + def __eq__(self, other: object) -> bool: #pylint: disable=protected-access + if not isinstance(other, ID): + return NotImplemented return self._id_str == other._id_str - def __hash__(self): + def __hash__(self) -> int: return hash(self._id_str) -def id_b58_encode(peer_id): +def id_b58_encode(peer_id: ID) -> str: """ return a b58-encoded string """ #pylint: disable=protected-access - return base58.b58encode(peer_id._id_str).decode() + return base58.b58encode(peer_id.get_raw_id()).decode() -def id_b58_decode(peer_id_str): +def id_b58_decode(peer_id_str: str) -> ID: """ return a base58-decoded peer ID """ return ID(base58.b58decode(peer_id_str)) -def id_from_public_key(key): +def id_from_public_key(key: RsaKey) -> ID: # export into binary format key_bin = key.exportKey("DER") - algo = multihash.Func.sha2_256 + algo: int = multihash.Func.sha2_256 # TODO: seems identity is not yet supported in pymultihash # if len(b) <= MAX_INLINE_KEY_LENGTH: # algo multihash.func.identity - mh_digest = multihash.digest(key_bin, algo) + mh_digest: multihash.Multihash = multihash.digest(key_bin, algo) return ID(mh_digest.encode()) -def id_from_private_key(key): +def id_from_private_key(key: RsaKey) -> ID: return id_from_public_key(key.publickey()) -def digest(string): - if not isinstance(string, bytes): - string = str(string).encode('utf8') - return hashlib.sha1(string).digest() +def digest(data: Union[str, bytes]) -> bytes: + if not isinstance(data, bytes): + data_bytes = str(data).encode('utf8') + return hashlib.sha1(data_bytes).digest() diff --git a/libp2p/peer/peerdata.py b/libp2p/peer/peerdata.py index c505bad..696e736 100644 --- a/libp2p/peer/peerdata.py +++ b/libp2p/peer/peerdata.py @@ -1,35 +1,48 @@ +from typing import ( + Any, + Dict, + List, + Sequence, +) + +from multiaddr import Multiaddr + from .peerdata_interface import IPeerData class PeerData(IPeerData): - def __init__(self): + metadata: Dict[Any, Any] + protocols: List[str] + addrs: List[Multiaddr] + + def __init__(self) -> None: self.metadata = {} self.protocols = [] self.addrs = [] - def get_protocols(self): + def get_protocols(self) -> List[str]: return self.protocols - def add_protocols(self, protocols): - self.protocols.extend(protocols) + def add_protocols(self, protocols: Sequence[str]) -> None: + self.protocols.extend(list(protocols)) - def set_protocols(self, protocols): - self.protocols = protocols + def set_protocols(self, protocols: Sequence[str]) -> None: + self.protocols = list(protocols) - def add_addrs(self, addrs): - self.addrs.extend(addrs) + def add_addrs(self, addrs: Sequence[Multiaddr]) -> None: + self.addrs.extend(list(addrs)) - def get_addrs(self): + def get_addrs(self) -> List[Multiaddr]: return self.addrs - def clear_addrs(self): + def clear_addrs(self) -> None: self.addrs = [] - def put_metadata(self, key, val): + def put_metadata(self, key: Any, val: Any) -> None: self.metadata[key] = val - def get_metadata(self, key): + def get_metadata(self, key: Any) -> Any: if key in self.metadata: return self.metadata[key] raise PeerDataError("key not found") diff --git a/libp2p/peer/peerdata_interface.py b/libp2p/peer/peerdata_interface.py index cf4a292..e2f150e 100644 --- a/libp2p/peer/peerdata_interface.py +++ b/libp2p/peer/peerdata_interface.py @@ -1,46 +1,53 @@ from abc import ABC, abstractmethod +from typing import ( + Any, + List, + Sequence, +) + +from multiaddr import Multiaddr class IPeerData(ABC): @abstractmethod - def get_protocols(self): + def get_protocols(self) -> List[str]: """ :return: all protocols associated with given peer """ @abstractmethod - def add_protocols(self, protocols): + def add_protocols(self, protocols: Sequence[str]) -> None: """ :param protocols: protocols to add """ @abstractmethod - def set_protocols(self, protocols): + def set_protocols(self, protocols: Sequence[str]) -> None: """ :param protocols: protocols to add """ @abstractmethod - def add_addrs(self, addrs): + def add_addrs(self, addrs: Sequence[Multiaddr]) -> None: """ :param addrs: multiaddresses to add """ @abstractmethod - def get_addrs(self): + def get_addrs(self) -> List[Multiaddr]: """ :return: all multiaddresses """ @abstractmethod - def clear_addrs(self): + def clear_addrs(self) -> None: """ Clear all addresses """ @abstractmethod - def put_metadata(self, key, val): + def put_metadata(self, key: Any, val: Any) -> None: """ :param key: key in KV pair :param val: val to associate with key @@ -48,7 +55,7 @@ class IPeerData(ABC): """ @abstractmethod - def get_metadata(self, key): + def get_metadata(self, key: Any) -> Any: """ :param key: key in KV pair :return: val for key diff --git a/libp2p/peer/peerinfo.py b/libp2p/peer/peerinfo.py index e173066..6271f18 100644 --- a/libp2p/peer/peerinfo.py +++ b/libp2p/peer/peerinfo.py @@ -1,12 +1,23 @@ +from typing import ( + List, +) + import multiaddr -from .id import id_b58_decode +from .id import ( + ID, + id_b58_decode, +) from .peerdata import PeerData class PeerInfo: # pylint: disable=too-few-public-methods - def __init__(self, peer_id, peer_data=None): + + peer_id: ID + addrs: List[multiaddr.Multiaddr] + + def __init__(self, peer_id: ID, peer_data: PeerData=None) -> None: self.peer_id = peer_id self.addrs = peer_data.get_addrs() if peer_data else None @@ -30,16 +41,16 @@ def info_from_p2p_addr(addr: multiaddr.Multiaddr) -> PeerInfo: ) # make sure the /p2p value parses as a peer.ID - peer_id_str = p2p_part.value_for_protocol(multiaddr.protocols.P_P2P) - peer_id = id_b58_decode(peer_id_str) + peer_id_str: str = p2p_part.value_for_protocol(multiaddr.protocols.P_P2P) + peer_id: ID = id_b58_decode(peer_id_str) # we might have received just an / p2p part, which means there's no addr. if len(parts) > 1: addr = multiaddr.Multiaddr.join(*parts[:-1]) peer_data = PeerData() - peer_data.addrs = [addr] - peer_data.protocols = [p.code for p in addr.protocols()] + peer_data.add_addrs(addr) + peer_data.set_protocols([p.code for p in addr.protocols()]) return PeerInfo(peer_id, peer_data) diff --git a/libp2p/peer/peermetadata_interface.py b/libp2p/peer/peermetadata_interface.py index 1badabd..6fcf8de 100644 --- a/libp2p/peer/peermetadata_interface.py +++ b/libp2p/peer/peermetadata_interface.py @@ -1,13 +1,20 @@ from abc import ABC, abstractmethod +from typing import ( + Any, +) + +from .id import ( + ID, +) class IPeerMetadata(ABC): - def __init__(self): + def __init__(self) -> None: pass @abstractmethod - def get(self, peer_id, key): + def get(self, peer_id: ID, key: Any) -> Any: """ :param peer_id: peer ID to lookup key for :param key: key to look up @@ -16,7 +23,7 @@ class IPeerMetadata(ABC): """ @abstractmethod - def put(self, peer_id, key, val): + def put(self, peer_id: ID, key: Any, val: Any) -> None: """ :param peer_id: peer ID to lookup key for :param key: key to associate with peer diff --git a/libp2p/peer/peerstore.py b/libp2p/peer/peerstore.py index 5a6a77b..5d32db1 100644 --- a/libp2p/peer/peerstore.py +++ b/libp2p/peer/peerstore.py @@ -1,3 +1,14 @@ +from typing import ( + Any, + Dict, + List, + Optional, + Sequence, +) + +from multiaddr import Multiaddr + +from .id import ID from .peerstore_interface import IPeerStore from .peerdata import PeerData from .peerinfo import PeerInfo @@ -5,11 +16,13 @@ from .peerinfo import PeerInfo class PeerStore(IPeerStore): - def __init__(self): + peer_map: Dict[ID, PeerData] + + def __init__(self) -> None: IPeerStore.__init__(self) self.peer_map = {} - def __create_or_get_peer(self, peer_id): + def __create_or_get_peer(self, peer_id: ID) -> PeerData: """ Returns the peer data for peer_id or creates a new peer data (and stores it in peer_map) if peer @@ -23,65 +36,65 @@ class PeerStore(IPeerStore): self.peer_map[peer_id] = data return self.peer_map[peer_id] - def peer_info(self, peer_id): + def peer_info(self, peer_id: ID) -> Optional[PeerInfo]: if peer_id in self.peer_map: - peer = self.peer_map[peer_id] - return PeerInfo(peer_id, peer) + peer_data = self.peer_map[peer_id] + return PeerInfo(peer_id, peer_data) return None - def get_protocols(self, peer_id): + def get_protocols(self, peer_id: ID) -> List[str]: if peer_id in self.peer_map: return self.peer_map[peer_id].get_protocols() raise PeerStoreError("peer ID not found") - def add_protocols(self, peer_id, protocols): + def add_protocols(self, peer_id: ID, protocols: Sequence[str]) -> None: peer = self.__create_or_get_peer(peer_id) - peer.add_protocols(protocols) + peer.add_protocols(list(protocols)) - def set_protocols(self, peer_id, protocols): + def set_protocols(self, peer_id: ID, protocols: Sequence[str]) -> None: peer = self.__create_or_get_peer(peer_id) - peer.set_protocols(protocols) + peer.set_protocols(list(protocols)) - def peers(self): + def peer_ids(self) -> List[ID]: return list(self.peer_map.keys()) - def get(self, peer_id, key): + def get(self, peer_id: ID, key: Any) -> Any: if peer_id in self.peer_map: val = self.peer_map[peer_id].get_metadata(key) return val raise PeerStoreError("peer ID not found") - def put(self, peer_id, key, val): + def put(self, peer_id: ID, key: Any, val: Any) -> None: # <> # 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) - def add_addr(self, peer_id, addr, ttl): + def add_addr(self, peer_id: ID, addr: Multiaddr, ttl: int) -> None: self.add_addrs(peer_id, [addr], ttl) - def add_addrs(self, peer_id, addrs, ttl): + def add_addrs(self, peer_id: ID, addrs: Sequence[Multiaddr], ttl: int) -> None: # Ignore ttl for now peer = self.__create_or_get_peer(peer_id) - peer.add_addrs(addrs) + peer.add_addrs(list(addrs)) - def addrs(self, peer_id): + def addrs(self, peer_id: ID) -> List[Multiaddr]: if peer_id in self.peer_map: return self.peer_map[peer_id].get_addrs() raise PeerStoreError("peer ID not found") - def clear_addrs(self, peer_id): + def clear_addrs(self, peer_id: ID) -> None: # Only clear addresses if the peer is in peer map if peer_id in self.peer_map: self.peer_map[peer_id].clear_addrs() - def peers_with_addrs(self): + def peers_with_addrs(self) -> List[ID]: # Add all peers with addrs at least 1 to output - output = [] + output: List[ID] = [] - for key in self.peer_map: - if len(self.peer_map[key].get_addrs()) >= 1: - output.append(key) + for peer_id in self.peer_map: + if len(self.peer_map[peer_id].get_addrs()) >= 1: + output.append(peer_id) return output diff --git a/libp2p/peer/peerstore_interface.py b/libp2p/peer/peerstore_interface.py index b368d8f..01b5b83 100644 --- a/libp2p/peer/peerstore_interface.py +++ b/libp2p/peer/peerstore_interface.py @@ -1,24 +1,31 @@ from abc import abstractmethod +from typing import ( + List, + Sequence, +) + +from .id import ID +from .peerinfo import PeerInfo from .addrbook_interface import IAddrBook from .peermetadata_interface import IPeerMetadata class IPeerStore(IAddrBook, IPeerMetadata): - def __init__(self): + def __init__(self) -> None: IPeerMetadata.__init__(self) IAddrBook.__init__(self) @abstractmethod - def peer_info(self, peer_id): + def peer_info(self, peer_id: ID) -> PeerInfo: """ :param peer_id: peer ID to get info for :return: peer info object """ @abstractmethod - def get_protocols(self, peer_id): + def get_protocols(self, peer_id: ID) -> List[str]: """ :param peer_id: peer ID to get protocols for :return: protocols (as strings) @@ -26,7 +33,7 @@ class IPeerStore(IAddrBook, IPeerMetadata): """ @abstractmethod - def add_protocols(self, peer_id, protocols): + def add_protocols(self, peer_id: ID, protocols: Sequence[str]) -> None: """ :param peer_id: peer ID to add protocols for :param protocols: protocols to add @@ -34,7 +41,7 @@ class IPeerStore(IAddrBook, IPeerMetadata): """ @abstractmethod - def set_protocols(self, peer_id, protocols): + def set_protocols(self, peer_id: ID, protocols: Sequence[str]) -> None: """ :param peer_id: peer ID to set protocols for :param protocols: protocols to set @@ -42,7 +49,7 @@ class IPeerStore(IAddrBook, IPeerMetadata): """ @abstractmethod - def peers(self): + def peer_ids(self) -> List[ID]: """ :return: all of the peer IDs stored in peer store """ diff --git a/tests/libp2p/test_libp2p.py b/tests/libp2p/test_libp2p.py index e9bfb83..a746503 100644 --- a/tests/libp2p/test_libp2p.py +++ b/tests/libp2p/test_libp2p.py @@ -316,20 +316,20 @@ async def test_host_connect(): transport_opt_list = [["/ip4/127.0.0.1/tcp/0"], ["/ip4/127.0.0.1/tcp/0"]] (node_a, node_b) = await set_up_nodes_by_transport_opt(transport_opt_list) - assert not node_a.get_peerstore().peers() + assert not node_a.get_peerstore().peer_ids() addr = node_b.get_addrs()[0] info = info_from_p2p_addr(addr) await node_a.connect(info) - assert len(node_a.get_peerstore().peers()) == 1 + assert len(node_a.get_peerstore().peer_ids()) == 1 await node_a.connect(info) # make sure we don't do double connection - assert len(node_a.get_peerstore().peers()) == 1 + assert len(node_a.get_peerstore().peer_ids()) == 1 - assert node_b.get_id() in node_a.get_peerstore().peers() + assert node_b.get_id() in node_a.get_peerstore().peer_ids() ma_node_b = multiaddr.Multiaddr('/p2p/%s' % node_b.get_id().pretty()) for addr in node_a.get_peerstore().addrs(node_b.get_id()): assert addr.encapsulate(ma_node_b) in node_b.get_addrs() diff --git a/tests/peer/test_peerstore.py b/tests/peer/test_peerstore.py index 52029ba..ffb7b2d 100644 --- a/tests/peer/test_peerstore.py +++ b/tests/peer/test_peerstore.py @@ -55,4 +55,4 @@ def test_peers(): store.put("peer2", "key", "val") store.add_addr("peer3", "/foo", 10) - assert set(store.peers()) == set(["peer1", "peer2", "peer3"]) + assert set(store.peer_ids()) == set(["peer1", "peer2", "peer3"])