Add type hints to network folder

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