ci: Add and configure mypy, and make it happy

ssb/muxrpc.py

@@ -1,6 +1,15 @@
 """MuxRPC"""
 
-from ssb.packet_stream import PSMessageType
+from typing import Any, AsyncIterator, Callable, Dict, Generator, List, Literal, Optional, Union
+
+from typing_extensions import Self
+
+from .packet_stream import PacketStream, PSMessage, PSMessageType, PSRequestHandler, PSStreamHandler
+
+MuxRPCJSON = Dict[str, Any]
+MuxRPCCallType = Literal["async", "duplex", "sink", "source", "sync"]
+MuxRPCRequestHandlerType = Callable[[PacketStream, "MuxRPCRequest"], None]
+MuxRPCRequestParam = Union[bytes, str, MuxRPCJSON]  # pylint: disable=invalid-name
 
 
 class MuxRPCAPIException(Exception):
@@ -10,7 +19,7 @@ class MuxRPCAPIException(Exception):
 class MuxRPCHandler:  # pylint: disable=too-few-public-methods
     """Base MuxRPC handler class"""
 
-    def check_message(self, msg):
+    def check_message(self, msg: PSMessage) -> None:
         """Check message validity"""
 
         body = msg.body
@@ -18,30 +27,48 @@ class MuxRPCHandler:  # pylint: disable=too-few-public-methods
         if isinstance(body, dict) and "name" in body and body["name"] == "Error":
             raise MuxRPCAPIException(body["message"])
 
+    def __await__(self):
+        raise NotImplementedError()
 
-class MuxRPCRequestHandler(MuxRPCHandler):
+    def __aiter__(self):
+        raise NotImplementedError()
+
+    async def __anext__(self):
+        raise NotImplementedError()
+
+    def send(self, msg: Any, msg_type: PSMessageType = PSMessageType.JSON, end: bool = False) -> None:
+        """Send a message through the stream"""
+
+        raise NotImplementedError()
+
+
+class MuxRPCRequestHandler(MuxRPCHandler):  # pylint: disable=abstract-method
     """Base class for MuxRPC request handlers"""
 
-    def __init__(self, ps_handler):
+    def __init__(self, ps_handler: PSRequestHandler):
         self.ps_handler = ps_handler
 
     def __await__(self):
         msg = yield from self.ps_handler.__await__()
         self.check_message(msg)
+
         return msg
 
 
-class MuxRPCSourceHandler(MuxRPCHandler):
+class MuxRPCSourceHandler(MuxRPCHandler):  # pylint: disable=abstract-method
     """MuxRPC handler for sources"""
 
-    def __init__(self, ps_handler):
+    def __init__(self, ps_handler: PSStreamHandler):
         self.ps_handler = ps_handler
 
-    def __aiter__(self):
+    def __aiter__(self) -> AsyncIterator[Optional[PSMessage]]:
         return self
 
-    async def __anext__(self):
+    async def __anext__(self) -> Optional[PSMessage]:
         msg = await self.ps_handler.__anext__()
+
+        assert msg
+
         self.check_message(msg)
 
         return msg
@@ -50,64 +77,72 @@ class MuxRPCSourceHandler(MuxRPCHandler):
 class MuxRPCSinkHandlerMixin:  # pylint: disable=too-few-public-methods
     """Mixin for sink-type MuxRPC handlers"""
 
-    def send(self, msg, msg_type=PSMessageType.JSON, end=False):
+    connection: PacketStream
+    req: int
+
+    def send(self, msg: Any, msg_type: PSMessageType = PSMessageType.JSON, end: bool = False) -> None:
         """Send a message through the stream"""
 
         self.connection.send(msg, stream=True, msg_type=msg_type, req=self.req, end_err=end)
 
 
-class MuxRPCDuplexHandler(MuxRPCSinkHandlerMixin, MuxRPCSourceHandler):
+class MuxRPCDuplexHandler(MuxRPCSinkHandlerMixin, MuxRPCSourceHandler):  # pylint: disable=abstract-method
     """MuxRPC handler for duplex streams"""
 
-    def __init__(self, ps_handler, connection, req):
+    def __init__(self, ps_handler: PSStreamHandler, connection: PacketStream, req: int):
         super().__init__(ps_handler)
 
         self.connection = connection
         self.req = req
 
 
-class MuxRPCSinkHandler(MuxRPCHandler, MuxRPCSinkHandlerMixin):
+class MuxRPCSinkHandler(MuxRPCHandler, MuxRPCSinkHandlerMixin):  # pylint: disable=abstract-method
     """MuxRPC handler for sinks"""
 
-    def __init__(self, connection, req):
+    def __init__(self, connection: PacketStream, req: int):
         self.connection = connection
         self.req = req
 
 
-def _get_appropriate_api_handler(type_, connection, ps_handler, req):
+def _get_appropriate_api_handler(
+    type_: MuxRPCCallType, connection: PacketStream, ps_handler: Union[PSRequestHandler, PSStreamHandler], req: int
+) -> MuxRPCHandler:
     """Find the appropriate MuxRPC handler"""
 
     if type_ in {"sync", "async"}:
+        assert isinstance(ps_handler, PSRequestHandler)
         return MuxRPCRequestHandler(ps_handler)
 
     if type_ == "source":
+        assert isinstance(ps_handler, PSStreamHandler)
         return MuxRPCSourceHandler(ps_handler)
 
     if type_ == "sink":
         return MuxRPCSinkHandler(connection, req)
 
     if type_ == "duplex":
+        assert isinstance(ps_handler, PSStreamHandler)
         return MuxRPCDuplexHandler(ps_handler, connection, req)
 
-    return None
+    raise TypeError(f"Unknown request type {type_}")
 
 
 class MuxRPCRequest:
     """MuxRPC request"""
 
     @classmethod
-    def from_message(cls, message):
+    def from_message(cls, message: PSMessage) -> Self:
         """Initialise a request from a raw packet stream message"""
 
         body = message.body
 
         return cls(".".join(body["name"]), body["args"])
 
-    def __init__(self, name, args):
+    def __init__(self, name: str, args: List[MuxRPCRequestParam]):
         self.name = name
         self.args = args
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"<MuxRPCRequest {self.name} {self.args}>"
 
 
@@ -115,7 +150,7 @@ class MuxRPCMessage:
     """MuxRPC message"""
 
     @classmethod
-    def from_message(cls, message):
+    def from_message(cls, message: PSMessage) -> Self:
         """Initialise a MuxRPC message from a raw packet stream message"""
 
         return cls(message.body)
@@ -123,21 +158,23 @@ class MuxRPCMessage:
     def __init__(self, body):
         self.body = body
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"<MuxRPCMessage {self.body}>"
 
 
 class MuxRPCAPI:
     """Generit MuxRPC API"""
 
-    def __init__(self):
-        self.handlers = {}
-        self.connection = None
+    def __init__(self) -> None:
+        self.handlers: Dict[str, MuxRPCRequestHandlerType] = {}
+        self.connection: Optional[PacketStream] = None
 
-    def __aiter__(self):
+    def __aiter__(self) -> AsyncIterator[None]:
         return self
 
-    async def __anext__(self):
+    async def __anext__(self) -> None:
+        assert self.connection
+
         req_message = await self.connection.__anext__()
 
         if req_message is None:
@@ -148,25 +185,25 @@ class MuxRPCAPI:
         if isinstance(body, dict) and body.get("name"):
             self.process(self.connection, MuxRPCRequest.from_message(req_message))
 
-    def __await__(self):
+    def __await__(self) -> Generator[None, None, None]:
         yield from self.__anext__().__await__()
 
-    def add_connection(self, connection):
+    def add_connection(self, connection: PacketStream) -> None:
         """Set the packet stream connection of this RPC API"""
 
         self.connection = connection
 
-    def define(self, name):
+    def define(self, name: str) -> Callable[[MuxRPCRequestHandlerType], MuxRPCRequestHandlerType]:
         """Decorator to define an RPC method handler"""
 
-        def _handle(f):
+        def _handle(f: MuxRPCRequestHandlerType) -> MuxRPCRequestHandlerType:
             self.handlers[name] = f
 
             return f
 
         return _handle
 
-    def process(self, connection, request):
+    def process(self, connection: PacketStream, request: MuxRPCRequest) -> None:
         """Process an incoming request"""
 
         handler = self.handlers.get(request.name)
@@ -176,9 +213,11 @@ class MuxRPCAPI:
 
         handler(connection, request)
 
-    def call(self, name, args, type_="sync"):
+    def call(self, name: str, args: List[MuxRPCRequestParam], type_: MuxRPCCallType = "sync") -> MuxRPCHandler:
         """Call an RPC method"""
 
+        assert self.connection
+
         if not self.connection.is_connected:
             raise Exception("not connected")  # pylint: disable=broad-exception-raised
 

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):

ssb/util.py

@@ -2,11 +2,19 @@
 
 from base64 import b64decode, b64encode
 import os
+from typing import TypedDict
 
 from nacl.signing import SigningKey
 import yaml
 
 
+class SSBSecret(TypedDict):
+    """Dictionary to hold an SSB identity"""
+
+    keypair: SigningKey
+    id: str
+
+
 class ConfigException(Exception):
     """Exception to raise if there is a problem with the configuration data"""
 
@@ -17,7 +25,7 @@ def tag(key):
     return b"@" + b64encode(bytes(key)) + b".ed25519"
 
 
-def load_ssb_secret():
+def load_ssb_secret() -> SSBSecret:
     """Load SSB keys from ~/.ssb"""
 
     with open(os.path.expanduser("~/.ssb/secret"), encoding="utf-8") as f: