Add type hints to network folder

libp2p/network/connection/raw_connection.py

@@ -1,8 +1,22 @@
+import asyncio
+
 from .raw_connection_interface import IRawConnection
 
 class RawConnection(IRawConnection):
 
-    def __init__(self, ip, port, reader, writer, initiator):
+    conn_ip: str
+    conn_port: str
+    reader: asyncio.StreamReader
+    writer: asyncio.StreamWriter
+    _next_id: int
+    initiator: bool
+
+    def __init__(self,
+                 ip: str,
+                 port: str,
+                 reader: asyncio.StreamReader,
+                 writer: asyncio.StreamWriter,
+                 initiator: bool) -> None:
         # pylint: disable=too-many-arguments
         self.conn_ip = ip
         self.conn_port = port
@@ -11,12 +25,12 @@ class RawConnection(IRawConnection):
         self._next_id = 0 if initiator else 1
         self.initiator = initiator
 
-    async def write(self, data):
+    async def write(self, data: bytes) -> None:
         self.writer.write(data)
         self.writer.write("\n".encode())
         await self.writer.drain()
 
-    async def read(self):
+    async def read(self) -> bytes:
         line = await self.reader.readline()
         adjusted_line = line.decode().rstrip('\n')
 
@@ -24,10 +38,10 @@ class RawConnection(IRawConnection):
         # encoding and decoding
         return adjusted_line.encode()
 
-    def close(self):
+    def close(self) -> None:
         self.writer.close()
 
-    def next_stream_id(self):
+    def next_stream_id(self) -> int:
         """
         Get next available stream id
         :return: next available stream id for the connection

libp2p/network/network_interface.py

@@ -1,16 +1,31 @@
 from abc import ABC, abstractmethod
+from typing import (
+    Any,
+    Callable,
+    Coroutine,
+    Sequence,
+)
+
+from multiaddr import Multiaddr
+
+from .notifee_interface import INotifee
+from .stream.net_stream import NetStream
+from libp2p.peer.id import (
+    ID,
+)
+from libp2p.stream_muxer.muxed_connection_interface import IMuxedConn
 
 
 class INetwork(ABC):
 
     @abstractmethod
-    def get_peer_id(self):
+    def get_peer_id(self) -> ID:
         """
         :return: the peer id
         """
 
     @abstractmethod
-    def dial_peer(self, peer_id):
+    def dial_peer(self, peer_id: ID) -> Coroutine[Any, Any, IMuxedConn]:
         """
         dial_peer try to create a connection to peer_id
 
@@ -20,7 +35,7 @@ class INetwork(ABC):
         """
 
     @abstractmethod
-    def set_stream_handler(self, protocol_id, stream_handler):
+    def set_stream_handler(self, protocol_id: str, stream_handler: Callable[[NetStream], None]) -> bool:
         """
         :param protocol_id: protocol id used on stream
         :param stream_handler: a stream handler instance
@@ -28,7 +43,7 @@ class INetwork(ABC):
         """
 
     @abstractmethod
-    def new_stream(self, peer_id, protocol_ids):
+    def new_stream(self, peer_id: ID, protocol_ids: Sequence[str]) -> Coroutine[Any, Any, NetStream]:
         """
         :param peer_id: peer_id of destination
         :param protocol_ids: available protocol ids to use for stream
@@ -36,14 +51,14 @@ class INetwork(ABC):
         """
 
     @abstractmethod
-    def listen(self, *args):
+    def listen(self, *args: Multiaddr) -> Coroutine[Any, Any, bool]:
         """
         :param *args: one or many multiaddrs to start listening on
         :return: True if at least one success
         """
 
     @abstractmethod
-    def notify(self, notifee):
+    def notify(self, notifee: INotifee) -> bool:
         """
         :param notifee: object implementing Notifee interface
         :return: true if notifee registered successfully, false otherwise

libp2p/network/notifee_interface.py

@@ -1,44 +1,51 @@
 from abc import ABC, abstractmethod
 
+from multiaddr import Multiaddr
+
+from libp2p.network.network_interface import INetwork
+from libp2p.network.stream.net_stream_interface import INetStream
+from libp2p.stream_muxer.muxed_connection_interface import IMuxedConn
+
+
 class INotifee(ABC):
 
     @abstractmethod
-    async def opened_stream(self, network, stream):
+    async def opened_stream(self, network: INetwork, stream: INetStream) -> None:
         """
         :param network: network the stream was opened on
         :param stream: stream that was opened
         """
 
     @abstractmethod
-    async def closed_stream(self, network, stream):
+    async def closed_stream(self, network: INetwork, stream: INetStream) -> None:
         """
         :param network: network the stream was closed on
         :param stream: stream that was closed
         """
 
     @abstractmethod
-    async def connected(self, network, conn):
+    async def connected(self, network: INetwork, conn: IMuxedConn) -> None:
         """
         :param network: network the connection was opened on
         :param conn: connection that was opened
         """
 
     @abstractmethod
-    async def disconnected(self, network, conn):
+    async def disconnected(self, network: INetwork, conn: IMuxedConn) -> None:
         """
         :param network: network the connection was closed on
         :param conn: connection that was closed
         """
 
     @abstractmethod
-    async def listen(self, network, multiaddr):
+    async def listen(self, network: INetwork, multiaddr: Multiaddr) -> None:
         """
         :param network: network the listener is listening on
         :param multiaddr: multiaddress listener is listening on
         """
 
     @abstractmethod
-    async def listen_close(self, network, multiaddr):
+    async def listen_close(self, network: INetwork, multiaddr: Multiaddr) -> None:
         """
         :param network: network the connection was opened on
         :param multiaddr: multiaddress listener is no longer listening on

libp2p/network/stream/net_stream.py

@@ -1,41 +1,47 @@
 from .net_stream_interface import INetStream
+from libp2p.stream_muxer.mplex.mplex_stream import MplexStream
+from libp2p.stream_muxer.mplex.mplex import Mplex
 
 
 class NetStream(INetStream):
 
-    def __init__(self, muxed_stream):
+    muxed_stream: MplexStream
+    mplex_conn: Mplex
+    protocol_id: str
+
+    def __init__(self, muxed_stream: MplexStream) -> None:
         self.muxed_stream = muxed_stream
         self.mplex_conn = muxed_stream.mplex_conn
         self.protocol_id = None
 
-    def get_protocol(self):
+    def get_protocol(self) -> str:
         """
         :return: protocol id that stream runs on
         """
         return self.protocol_id
 
-    def set_protocol(self, protocol_id):
+    def set_protocol(self, protocol_id: str) -> None:
         """
         :param protocol_id: protocol id that stream runs on
         :return: true if successful
         """
         self.protocol_id = protocol_id
 
-    async def read(self):
+    async def read(self) -> bytes:
         """
         read from stream
         :return: bytes of input until EOF
         """
         return await self.muxed_stream.read()
 
-    async def write(self, data):
+    async def write(self, data: bytes) -> int:
         """
         write to stream
         :return: number of bytes written
         """
         return await self.muxed_stream.write(data)
 
-    async def close(self):
+    async def close(self) -> bool:
         """
         close stream
         :return: true if successful

libp2p/network/stream/net_stream_interface.py

@@ -1,37 +1,41 @@
 from abc import ABC, abstractmethod
+from typing import (
+    Any,
+    Coroutine,
+)
 
 
 class INetStream(ABC):
 
     @abstractmethod
-    def get_protocol(self):
+    def get_protocol(self) -> str:
         """
         :return: protocol id that stream runs on
         """
 
     @abstractmethod
-    def set_protocol(self, protocol_id):
+    def set_protocol(self, protocol_id: str) -> bool:
         """
         :param protocol_id: protocol id that stream runs on
         :return: true if successful
         """
 
     @abstractmethod
-    def read(self):
+    def read(self) -> Coroutine[Any, Any, bytes]:
         """
         reads from the underlying muxed_stream
         :return: bytes of input
         """
 
     @abstractmethod
-    def write(self, _bytes):
+    def write(self, data: bytes) -> Coroutine[Any, Any, int]:
         """
         write to the underlying muxed_stream
         :return: number of bytes written
         """
 
     @abstractmethod
-    def close(self):
+    def close(self) -> Coroutine[Any, Any, bool]:
         """
         close the underlying muxed stream
         :return: true if successful

libp2p/network/swarm.py

@@ -1,18 +1,57 @@
 import asyncio
+from typing import (
+    Any,
+    Callable,
+    Coroutine,
+    Dict,
+    List,
+    Sequence,
+)
 
-from libp2p.protocol_muxer.multiselect_client import MultiselectClient
+from multiaddr import Multiaddr
-from libp2p.protocol_muxer.multiselect import Multiselect
-from libp2p.peer.id import id_b58_decode
 
 from .network_interface import INetwork
 from .notifee_interface import INotifee
-from .stream.net_stream import NetStream
 from .connection.raw_connection import RawConnection
+from .stream.net_stream import NetStream
+from libp2p.peer.id import (
+    ID,
+    id_b58_decode,
+)
+from libp2p.peer.peerstore import PeerStore
+from libp2p.protocol_muxer.multiselect import Multiselect
+from libp2p.protocol_muxer.multiselect_client import MultiselectClient
+from libp2p.routing.interfaces import IPeerRouting
+from libp2p.transport.upgrader import TransportUpgrader
+from libp2p.transport.transport_interface import ITransport
+from libp2p.transport.listener_interface import IListener
+from libp2p.stream_muxer.mplex.mplex_stream import MplexStream
+from libp2p.stream_muxer.muxed_connection_interface import IMuxedConn
+
 
 class Swarm(INetwork):
     # pylint: disable=too-many-instance-attributes,cell-var-from-loop,too-many-arguments
 
-    def __init__(self, peer_id, peerstore, upgrader, transport, router):
+    self_id: ID
+    peerstore: PeerStore
+    upgrader: TransportUpgrader
+    transport: ITransport
+    router: IPeerRouting
+    connections: Dict[ID, IMuxedConn]
+    listeners: Dict[str, IListener]
+    stream_handlers: Dict[NetStream, Callable[[NetStream], None]]
+
+    multiselect: Multiselect
+    multiselect_client: MultiselectClient
+
+    notifees: List[INotifee]
+
+    def __init__(self,
+                 peer_id: ID,
+                 peerstore: PeerStore,
+                 upgrader: TransportUpgrader,
+                 transport: ITransport,
+                 router: IPeerRouting):
         self.self_id = peer_id
         self.peerstore = peerstore
         self.upgrader = upgrader
@@ -32,10 +71,10 @@ class Swarm(INetwork):
         # Create generic protocol handler
         self.generic_protocol_handler = create_generic_protocol_handler(self)
 
-    def get_peer_id(self):
+    def get_peer_id(self) -> ID:
         return self.self_id
 
-    def set_stream_handler(self, protocol_id, stream_handler):
+    def set_stream_handler(self, protocol_id: str, stream_handler: Callable[[NetStream], None]) -> bool:
         """
         :param protocol_id: protocol id used on stream
         :param stream_handler: a stream handler instance
@@ -44,7 +83,7 @@ class Swarm(INetwork):
         self.multiselect.add_handler(protocol_id, stream_handler)
         return True
 
-    async def dial_peer(self, peer_id):
+    async def dial_peer(self, peer_id: ID) -> IMuxedConn:
         """
         dial_peer try to create a connection to peer_id
         :param peer_id: peer if we want to dial
@@ -87,7 +126,7 @@ class Swarm(INetwork):
 
         return muxed_conn
 
-    async def new_stream(self, peer_id, protocol_ids):
+    async def new_stream(self, peer_id: ID, protocol_ids: Sequence[str]) -> NetStream:
         """
         :param peer_id: peer_id of destination
         :param protocol_id: protocol id
@@ -109,7 +148,7 @@ class Swarm(INetwork):
         muxed_stream = await muxed_conn.open_stream(protocol_ids[0], multiaddr)
 
         # Perform protocol muxing to determine protocol to use
-        selected_protocol = await self.multiselect_client.select_one_of(protocol_ids, muxed_stream)
+        selected_protocol = await self.multiselect_client.select_one_of(list(protocol_ids), muxed_stream)
 
         # Create a net stream with the selected protocol
         net_stream = NetStream(muxed_stream)
@@ -121,7 +160,7 @@ class Swarm(INetwork):
 
         return net_stream
 
-    async def listen(self, *args):
+    async def listen(self, *args: Multiaddr) -> bool:
         """
         :param *args: one or many multiaddrs to start listening on
         :return: true if at least one success
@@ -139,7 +178,7 @@ class Swarm(INetwork):
             if str(multiaddr) in self.listeners:
                 return True
 
-            async def conn_handler(reader, writer):
+            async def conn_handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
                 # Read in first message (should be peer_id of initiator) and ack
                 peer_id = id_b58_decode((await reader.read(1024)).decode())
 
@@ -182,7 +221,7 @@ class Swarm(INetwork):
         # No multiaddr succeeded
         return False
 
-    def notify(self, notifee):
+    def notify(self, notifee: INotifee) -> bool:
         """
         :param notifee: object implementing Notifee interface
         :return: true if notifee registered successfully, false otherwise
@@ -192,7 +231,7 @@ class Swarm(INetwork):
             return True
         return False
 
-    def add_router(self, router):
+    def add_router(self, router: IPeerRouting) -> None:
         self.router = router
 
     # TODO: `tear_down`
@@ -204,7 +243,10 @@ class Swarm(INetwork):
     # TODO: `disconnect`?
 
 
-def create_generic_protocol_handler(swarm):
+GenericProtocolHandlerFn = Callable[[MplexStream], Coroutine[Any, Any, None]]
+
+
+def create_generic_protocol_handler(swarm: Swarm) -> GenericProtocolHandlerFn:
     """
     Create a generic protocol handler from the given swarm. We use swarm
     to extract the multiselect module so that generic_protocol_handler
@@ -213,7 +255,7 @@ def create_generic_protocol_handler(swarm):
     """
     multiselect = swarm.multiselect
 
-    async def generic_protocol_handler(muxed_stream):
+    async def generic_protocol_handler(muxed_stream: MplexStream) -> None:
         # Perform protocol muxing to determine protocol to use
         protocol, handler = await multiselect.negotiate(muxed_stream)
 
@@ -229,5 +271,6 @@ def create_generic_protocol_handler(swarm):
 
     return generic_protocol_handler
 
+
 class SwarmException(Exception):
     pass