Add type hints to network folder

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 libp2p.protocol_muxer.multiselect import Multiselect
-from libp2p.peer.id import id_b58_decode
+from multiaddr import Multiaddr
 
 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