Add type hints to peer folder

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

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):
+def digest(data: Union[str, bytes]) -> bytes:
-    if not isinstance(string, bytes):
+    if not isinstance(data, bytes):
-        string = str(string).encode('utf8')
+        data_bytes = str(data).encode('utf8')
-    return hashlib.sha1(string).digest()
+    return hashlib.sha1(data_bytes).digest()

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):
+    def add_protocols(self, protocols: Sequence[str]) -> None:
-        self.protocols.extend(protocols)
+        self.protocols.extend(list(protocols))
 
-    def set_protocols(self, protocols):
+    def set_protocols(self, protocols: Sequence[str]) -> None:
-        self.protocols = protocols
+        self.protocols = list(protocols)
 
-    def add_addrs(self, addrs):
+    def add_addrs(self, addrs: Sequence[Multiaddr]) -> None:
-        self.addrs.extend(addrs)
+        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")

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

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_str: str = p2p_part.value_for_protocol(multiaddr.protocols.P_P2P)
-    peer_id = id_b58_decode(peer_id_str)
+    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.add_addrs(addr)
-    peer_data.protocols = [p.code for p in addr.protocols()]
+    peer_data.set_protocols([p.code for p in addr.protocols()])
 
     return PeerInfo(peer_id, peer_data)
 

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

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]
+            peer_data = self.peer_map[peer_id]
-            return PeerInfo(peer_id, peer)
+            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:
+        for peer_id in self.peer_map:
-            if len(self.peer_map[key].get_addrs()) >= 1:
+            if len(self.peer_map[peer_id].get_addrs()) >= 1:
-                output.append(key)
+                output.append(peer_id)
         return output
 
 

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

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

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