Type-annotate all of SocketShark (#328)

* Type-annotate all of SocketShark.

* More accurate annotation for `ping_timeout_handler`.

Previous annotation was too loose. The `ping_timeout_handler` is only ever
called with the result of `await self.websocket.ping()`, which has the
following signature:
```
async def ping(self, data: Optional[Data] = None) -> Awaitable[None]:
```

* More precise `client` annotation.

* Type-annotate the re-added `LogMetrics`, too.

I briefly removed it from master before realizing it's still used (in a very
counter-intuitive fashion).

Author: wojcikstefan

setup.cfg

@@ -109,40 +109,3 @@ check_untyped_defs = true
 disallow_any_generics = true
 disallow_untyped_defs = true
 disallow_incomplete_defs = true
-
-# The following modules have not yet been fully type-annotated. Allow untyped
-# defs in them until they get typed, then remove the corresponding override.
-[mypy-socketshark]
-allow_untyped_defs = True
-
-[mypy-socketshark.backend.websockets]
-allow_untyped_defs = True
-
-[mypy-socketshark.events]
-allow_untyped_defs = True
-
-[mypy-socketshark.exceptions]
-allow_untyped_defs = True
-
-[mypy-socketshark.metrics]
-allow_untyped_defs = True
-
-[mypy-socketshark.metrics.prometheus]
-allow_untyped_defs = True
-
-[mypy-socketshark.receiver]
-allow_untyped_defs = True
-
-[mypy-socketshark.redis_connection]
-allow_untyped_defs = True
-allow_incomplete_defs = True
-
-[mypy-socketshark.session]
-allow_untyped_defs = True
-allow_incomplete_defs = True
-
-[mypy-socketshark.subscription]
-allow_untyped_defs = True
-
-[mypy-socketshark.utils]
-allow_untyped_defs = True

socketshark/__init__.py

@@ -5,6 +5,7 @@
 import signal
 import ssl
 import sys
+from typing import Any
 
 import aioredis
 import click
@@ -14,14 +15,16 @@
 from .metrics import Metrics
 from .receiver import ServiceReceiver
 from .redis_connection import RedisConnection
+from .session import Session
+from .types import Config, LogConfig
 
 
-def setup_logging(log_config):
+def setup_logging(log_config: LogConfig) -> None:
     # Configure root logger if logging level is specified in config
     if log_config['level']:
         level = getattr(logging, log_config['level'])
         formatter = logging.Formatter(log_config['format'])
-        sh = logging.StreamHandler()
+        sh: logging.StreamHandler[Any] = logging.StreamHandler()
         sh.setFormatter(formatter)
 
         logger = logging.getLogger()
@@ -36,7 +39,7 @@ def setup_logging(log_config):
         setup_structlog(sys.stdout.isatty())
 
 
-def setup_structlog(tty=False):
+def setup_structlog(tty: bool = False) -> None:
     processors = [
         structlog.stdlib.filter_by_level,
         structlog.stdlib.add_log_level,
@@ -59,7 +62,7 @@ def setup_structlog(tty=False):
     )
 
 
-def load_backend(config):
+def load_backend(config: Config) -> Any:
     """
     Return the backend module from the given SocketShark configuration.
     """
@@ -69,30 +72,34 @@ def load_backend(config):
 
 
 class SocketShark:
-    def __init__(self, config):
+    def __init__(self, config: Config) -> None:
         self.config = config
         backend_module = load_backend(config)
         backend_cls = backend_module.Backend
         self.backend = backend_cls(self)
         self._init_logging()
-        self._task = None
+        self._task: asyncio.Task[None] | None = None
         self._shutdown = False
-        self.sessions = set()
+        self.sessions: set[Session] = set()
         self.metrics = Metrics(self)
         self.metrics.initialize()
         self.metrics.set_ready(False)
-        self.redis_connections = []
+        self.redis_connections: list[RedisConnection] = []
 
-    def _init_logging(self):
+    def _init_logging(self) -> None:
         logger_name = self.config['LOG']['logger_name']
         trace_logger_prefix = self.config['LOG']['trace_logger_prefix']
         trace_logger_name = '{}.{}'.format(trace_logger_prefix, logger_name)
         pid = os.getpid()
-        self.log = structlog.get_logger(logger_name).bind(pid=pid)
-        self.trace_log = structlog.get_logger(trace_logger_name).bind(pid=pid)
+        self.log: structlog.stdlib.BoundLogger = structlog.get_logger(
+            logger_name
+        ).bind(pid=pid)
+        self.trace_log: structlog.stdlib.BoundLogger = structlog.get_logger(
+            trace_logger_name
+        ).bind(pid=pid)
         self.trace_log.debug('trace')
 
-    def signal_ready(self):
+    def signal_ready(self) -> None:
         """
         Notify that the backend is ready.
         """
@@ -104,14 +111,14 @@ def signal_ready(self):
         )
         self.metrics.set_ready(True)
 
-    def signal_shutdown(self):
+    def signal_shutdown(self) -> None:
         """
         Notify that the backend shut down.
         """
         self.log.info('done')
         self.metrics.set_ready(False)
 
-    async def _redis_connection_handler(self):
+    async def _redis_connection_handler(self) -> None:
         """
         Handle Redis connection errors.
 
@@ -132,7 +139,7 @@ async def _redis_connection_handler(self):
         # Redis goes down so they can reconnect and restore subscriptions.
         asyncio.ensure_future(self.shutdown())
 
-    async def prepare(self):
+    async def prepare(self) -> None:
         """
         Callback called by the backend to prepare SocketShark.
 
@@ -155,12 +162,12 @@ async def prepare(self):
 
         self.service_receiver = ServiceReceiver(self)
 
-    def _cleanup(self):
+    def _cleanup(self) -> None:
         self._redis_connection_handler_task.cancel()
         for c in self.redis_connections:
             c.redis.close()
 
-    async def shutdown(self):
+    async def shutdown(self) -> None:
         """
         Shut down SocketShark cleanly.
         """
@@ -201,20 +208,22 @@ async def shutdown(self):
         self._uninstall_signal_handlers()
         self._shutdown = False
 
-    async def run_service_receiver(self, once=False):
+    async def run_service_receiver(
+        self, once: bool = False
+    ) -> list[bool | None] | None:
         return await self.service_receiver.reader(once=once)
 
-    def start(self):
+    def start(self) -> None:
         """
         Start the backend (main entrypoint into SocketShark).
         """
         self.backend.start()
 
-    async def _run(self, once=False):
+    async def _run(self, once: bool = False) -> None:
         await self.run_service_receiver()
         asyncio.ensure_future(self.shutdown())
 
-    async def run(self, once=False):
+    async def run(self, once: bool = False) -> None:
         """
         Set up SocketShark signal handlers and run the service receiver.
 
@@ -223,39 +232,40 @@ async def run(self, once=False):
         self._install_signal_handlers()
         self._task = asyncio.ensure_future(self._run())
 
-    def _install_signal_handlers(self):
+    def _install_signal_handlers(self) -> None:
         """
         Set up signal handlers for safely stopping the worker.
         """
 
-        def request_stop():
+        def request_stop() -> None:
             self.log.info('stop requested')
             asyncio.ensure_future(self.shutdown())
 
         loop = asyncio.get_event_loop()
         loop.add_signal_handler(signal.SIGINT, request_stop)
         loop.add_signal_handler(signal.SIGTERM, request_stop)
 
-    def _uninstall_signal_handlers(self):
+    def _uninstall_signal_handlers(self) -> None:
         """
         Restore default signal handlers.
         """
         loop = asyncio.get_event_loop()
         loop.remove_signal_handler(signal.SIGINT)
         loop.remove_signal_handler(signal.SIGTERM)
 
-    def get_ssl_context(self):
+    def get_ssl_context(self) -> ssl.SSLContext | None:
         ssl_settings = self.config.get('WS_SSL')
         if ssl_settings:
             ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
             ssl_context.load_cert_chain(
                 certfile=ssl_settings['cert'], keyfile=ssl_settings['key']
             )
             return ssl_context
+        return None
 
 
-def load_config(config_name):
-    config = {}
+def load_config(config_name: str) -> Config:
+    config: dict[str, Any] = {}
 
     # Get config defaults
     for key in dir(config_defaults):
@@ -272,13 +282,13 @@ def load_config(config_name):
             else:
                 config[key] = value
 
-    return config
+    return Config(config)
 
 
 @click.command()
 @click.option('-c', '--config', required=True, help='dotted path to config')
 @click.pass_context
-def run(context, config):
+def run(context: click.Context, config: str) -> None:
     config_obj = load_config(config)
 
     setup_logging(config_obj['LOG'])

socketshark/backend/websockets.py

@@ -1,15 +1,24 @@
 import asyncio
 import json
 import time
+from typing import TYPE_CHECKING
 
 import websockets
 
 from .. import constants as c
 from ..session import Session
+from ..types import ClientMessage
+
+if TYPE_CHECKING:
+    from .. import SocketShark
 
 
 class Client:
-    def __init__(self, shark, websocket):
+    def __init__(
+        self,
+        shark: 'SocketShark',
+        websocket: websockets.WebSocketServerProtocol,
+    ) -> None:
         self.websocket = websocket
         self.session = Session(
             shark,
@@ -20,7 +29,7 @@ def __init__(self, shark, websocket):
         )
         self.shark = shark
 
-    async def ping_timeout_handler(self, ping):
+    async def ping_timeout_handler(self, ping: asyncio.Future[None]) -> bool:
         ping_timeout = self.shark.config['WS_PING']['timeout']
         await asyncio.sleep(ping_timeout)
 
@@ -33,7 +42,7 @@ async def ping_timeout_handler(self, ping):
 
         return False
 
-    async def ping_handler(self):
+    async def ping_handler(self) -> None:
         ping_interval = self.shark.config['WS_PING']['interval']
         if not ping_interval:
             return
@@ -62,7 +71,7 @@ async def ping_handler(self):
             if not timeout_handler.cancel() and timeout_handler.result():
                 return
 
-    async def consumer_handler(self):
+    async def consumer_handler(self) -> None:
         try:
             ping_handler = asyncio.ensure_future(self.ping_handler())
             try:
@@ -73,10 +82,12 @@ async def consumer_handler(self):
                     except json.decoder.JSONDecodeError:
                         self.session.log.warn('received invalid json')
                         await self.send(
-                            {
-                                "status": "error",
-                                "error": c.ERR_INVALID_EVENT,
-                            }
+                            ClientMessage(
+                                {
+                                    "status": "error",
+                                    "error": c.ERR_INVALID_EVENT,
+                                }
+                            )
                         )
                     else:
                         await self.session.on_client_event(data)
@@ -87,23 +98,23 @@ async def consumer_handler(self):
         except Exception:
             self.session.log.exception('unhandled error in consumer handler')
 
-    async def send(self, event):
+    async def send(self, event: ClientMessage) -> None:
         try:
             await self.websocket.send(json.dumps(event))
         except websockets.ConnectionClosed:
             self.session.log.warn('attempted to send to closed socket')
 
-    async def close(self):
+    async def close(self) -> None:
         await self.websocket.close()
 
 
 class Backend:
-    def __init__(self, shark):
+    def __init__(self, shark: 'SocketShark') -> None:
         self.shark = shark
-        self.server = None
-        self._closed = False
+        self.server: websockets.WebSocketServer | None = None
+        self._closed: bool = False
 
-    def close(self):
+    def close(self) -> None:
         """
         Called by SocketShark to make the backend stop accepting connections.
         """
@@ -112,20 +123,22 @@ def close(self):
             self._closed = True
             self.server.server.close()
 
-    async def shutdown(self):
+    async def shutdown(self) -> None:
         """
         Called by SocketShark to close any open connections.
         """
         if self.server:
             self.server.close()
             await self.server.wait_closed()
 
-    def start(self):
+    def start(self) -> None:
         """
         Called by SocketShark to initialize the server, prepare & run.
         """
 
-        async def serve(websocket, path):
+        async def serve(
+            websocket: websockets.WebSocketServerProtocol, path: str
+        ) -> None:
             # If there are any pending connections that were established after
             # calling close() but before this callback was executed, close
             # them immediately.