ci: Add and configure mypy, and make it happy

ssb/packet_stream.py

@@ -6,11 +6,15 @@ import logging
 from math import ceil
 import struct
 from time import time
+from typing import Any, AsyncIterator, Dict, Optional, Union
 
+from secret_handshake.network import SHSDuplexStream
 import simplejson
 
 logger = logging.getLogger("packet_stream")
 
+PSMessageData = Union[bytes, bool, Dict[str, Any], str]
+
 
 class PSMessageType(Enum):
     """Available message types"""
@@ -41,7 +45,7 @@ class PSStreamHandler:
     def __aiter__(self):
         return self
 
-    async def __anext__(self):
+    async def __anext__(self) -> Optional["PSMessage"]:
         elem = await self.queue.get()
 
         if not elem:
@@ -93,7 +97,7 @@ class PSMessage:
         return cls(type_, body, bool(flags & 0x08), bool(flags & 0x04), req=req)
 
     @property
-    def data(self):
+    def data(self) -> bytes:
         """The raw message data"""
 
         if self.type == PSMessageType.TEXT:
@@ -127,7 +131,7 @@ class PSMessage:
 class PacketStream:
     """SSB Packet stream"""
 
-    def __init__(self, connection):
+    def __init__(self, connection: SHSDuplexStream):
         self.connection = connection
         self.req_counter = 1
         self._event_map = {}
@@ -144,10 +148,10 @@ class PacketStream:
 
         return self.connection.is_connected
 
-    def __aiter__(self):
+    def __aiter__(self) -> AsyncIterator[Optional[PSMessage]]:
         return self
 
-    async def __anext__(self):
+    async def __anext__(self) -> Optional[PSMessage]:
         msg = await self.read()
 
         if not msg:
@@ -202,7 +206,7 @@ class PacketStream:
                 logger.info("RESPONSE [%d]: EOS", -msg.req)
         return msg
 
-    def _write(self, msg):
+    def _write(self, msg: PSMessage) -> None:
         logger.info("SEND [%d]: %r", msg.req, msg)
         header = struct.pack(
             ">BIi", (int(msg.stream) << 3) | (int(msg.end_err) << 2) | msg.type.value, len(msg.data), msg.req
@@ -213,11 +217,17 @@ class PacketStream:
         logger.debug("WRITE DATA: %s", msg.data)
 
     def send(  # pylint: disable=too-many-arguments
-        self, data, msg_type=PSMessageType.JSON, stream=False, end_err=False, req=None
+        self,
+        data: Any,
+        msg_type: PSMessageType = PSMessageType.JSON,
+        stream: bool = False,
+        end_err: bool = False,
+        req: Optional[int] = None,
     ):
         """Send data through the packet stream"""
 
         update_counter = False
+
         if req is None:
             update_counter = True
             req = self.req_counter
@@ -231,10 +241,12 @@ class PacketStream:
             handler = PSStreamHandler(self.req_counter)
         else:
             handler = PSRequestHandler(self.req_counter)
+
         self.register_handler(handler)
 
         if update_counter:
             self.req_counter += 1
+
         return handler
 
     def disconnect(self):