StarMirror/py-libp2p
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add type hints to peer folder

Browse Source
This commit is contained in:
NIC619 2019-07-27 16:27:01 +08:00
parent b695b0e1ec
commit f2de986c74
No known key found for this signature in database
GPG Key ID: 570C35F5C2D51B17
10 changed files with 165 additions and 87 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
libp2p/peer/addrbook_interface.py
View File

@ -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
"""

47
libp2p/peer/id.py
View File

@ -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()

37
libp2p/peer/peerdata.py
View File

@ -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")

23
libp2p/peer/peerdata_interface.py
View File

@ -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

23
libp2p/peer/peerinfo.py
View File

@ -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)

13
libp2p/peer/peermetadata_interface.py
View File

@ -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

59
libp2p/peer/peerstore.py
View File

@ -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

19
libp2p/peer/peerstore_interface.py
View File

@ -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
"""

8
tests/libp2p/test_libp2p.py
View File

@ -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()

2
tests/peer/test_peerstore.py
View File

@ -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"])