add typing to transport

libp2p/transport/tcp/tcp.py

@@ -1,68 +1,79 @@
 import asyncio
 
-import multiaddr
+from multiaddr import Multiaddr
 
 from libp2p.network.connection.raw_connection import RawConnection
 
-from ..listener_interface import IListener
-from ..transport_interface import ITransport
+from libp2p.transport.listener_interface import IListener
+from libp2p.transport.transport_interface import ITransport
+
+from typing import TYPE_CHECKING, List, Callable
+
+if TYPE_CHECKING:
+    from socket import socket
+    from libp2p.transport.typing import THandler
+    from libp2p.network.connection.raw_connection_interface import IRawConnection
+
+
+class TCPListener(IListener):
+    multiaddrs: List[Multiaddr]
+    server = None
+    handler = None
+
+    def __init__(self, handler_function: "THandler" = None) -> None:
+        self.multiaddrs = []
+        self.server = None
+        self.handler = handler_function
+
+    async def listen(self, maddr: Multiaddr) -> bool:
+        """
+        put listener in listening mode and wait for incoming connections
+        :param maddr: maddr of peer
+        :return: return True if successful
+        """
+        self.server = await asyncio.start_server(
+            self.handler,
+            maddr.value_for_protocol("ip4"),
+            maddr.value_for_protocol("tcp"),
+        )
+        socket = self.server.sockets[0]
+        self.multiaddrs.append(_multiaddr_from_socket(socket))
+
+        return True
+
+    def get_addrs(self) -> List[Multiaddr]:
+        """
+        retrieve list of addresses the listener is listening on
+        :return: return list of addrs
+        """
+        # TODO check if server is listening
+        return self.multiaddrs
+
+    def close(self) -> bool:
+        """
+        close the listener such that no more connections
+        can be open on this transport instance
+        :return: return True if successful
+        """
+        if self.server is None:
+            return False
+        self.server.close()
+        _loop = asyncio.get_event_loop()
+        _loop.run_until_complete(self.server.wait_closed())
+        _loop.close()
+        self.server = None
+        return True
 
 
 class TCP(ITransport):
-    def __init__(self):
-        self.listener = self.Listener()
+    def __init__(self) -> None:
+        self.listener = TCPListener()
 
-    class Listener(IListener):
-        def __init__(self, handler_function=None):
-            self.multiaddrs = []
-            self.server = None
-            self.handler = handler_function
-
-        async def listen(self, maddr):
-            """
-            put listener in listening mode and wait for incoming connections
-            :param maddr: maddr of peer
-            :return: return True if successful
-            """
-            self.server = await asyncio.start_server(
-                self.handler, maddr.value_for_protocol("ip4"), maddr.value_for_protocol("tcp")
-            )
-            socket = self.server.sockets[0]
-            self.multiaddrs.append(_multiaddr_from_socket(socket))
-
-            return True
-
-        def get_addrs(self):
-            """
-            retrieve list of addresses the listener is listening on
-            :return: return list of addrs
-            """
-            # TODO check if server is listening
-            return self.multiaddrs
-
-        def close(self, options=None):
-            """
-            close the listener such that no more connections
-            can be open on this transport instance
-            :param options: optional object potential with timeout
-            a timeout value in ms that fires and destroy all connections
-            :return: return True if successful
-            """
-            if self.server is None:
-                return False
-            self.server.close()
-            _loop = asyncio.get_event_loop()
-            _loop.run_until_complete(self.server.wait_closed())
-            _loop.close()
-            self.server = None
-            return True
-
-    async def dial(self, maddr, self_id, options=None):
+    async def dial(self, maddr: Multiaddr, self_id: ID) -> "IRawConnection":
         """
         dial a transport to peer listening on multiaddr
         :param maddr: multiaddr of peer
         :param self_id: peer_id of the dialer (to send to receiver)
-        :param options: optional object
         :return: True if successful
         """
         host = maddr.value_for_protocol("ip4")
@@ -81,18 +92,17 @@ class TCP(ITransport):
         if ack != expected_ack_str:
             raise Exception("Receiver did not receive peer id")
 
-        return RawConnection(host, port, reader, writer, True)
+        return RawConnection(host, str(port), reader, writer, True)
 
-    def create_listener(self, handler_function, options=None):
+    def create_listener(self, handler_function: "THandler") -> TCPListener:
         """
         create listener on transport
-        :param options: optional object with properties the listener must have
         :param handler_function: a function called when a new connection is received
         that takes a connection as argument which implements interface-connection
         :return: a listener object that implements listener_interface.py
         """
-        return self.Listener(handler_function)
+        return TCPListener(handler_function)
 
 
-def _multiaddr_from_socket(socket):
-    return multiaddr.Multiaddr("/ip4/%s/tcp/%s" % socket.getsockname())
+def _multiaddr_from_socket(socket: "socket") -> Multiaddr:
+    return Multiaddr("/ip4/%s/tcp/%s" % socket.getsockname())