feat: improve typing definitions, introduce mypy in CI (#338)

* feat: do not supress callback exceptions raised from calling callbacks

* fix: add explicit clauses for ConnectionClosedOK in send and heartbeat

* fix: change import order

* fix: apply ruff format

* feat: introduce flake.nix leveraging pyproject.nix

introduce flake.nix configuration leveraging the already existing
information in pyproject.toml, by using pyproject.nix to parse it and
generate the correct derivation. this way, any updates to
pyproject.toml should automatically be reflected in the flake
configuration, with low additional maintaining overhead required.

the only time when this may come as a problem is due to dependabot
constantly upgrading the packages to versions that do not exist yet in
nixpkgs. it may be required to force dependabot to also update the
nixpkgs version declared in the flake for it to not also break (or at
least to upgrade less often)

* chore: make `flake.nix` aware of dev dependencies, including ruff and pre-commit

by adding `dependency-groups`, `pyproject.nix` can now parse the
dependency groups and add these to the project, by passing `groups =
["dev"]` when calling the renderer.

special care was needed in order to expose `ruff` and `pre-commit` to
the toplevel of the environment. by default, it will treat them as
python dependencies and hide them in the python wrapper. by using
`groupBy`, we split the dependency list in "python" dependencies and
top level ones, such that they can be exposed.

sadly, these toplevel packages are not exposed in `python.pkgs` so we
must manually override them in a case by case basis.

* chore: refactor `dependencies-for`

* feat: add nix develop CI job

ensure that `nix develop` works on CI

* fix: change job name to be more compatible with the other actions

* feat: setup infra and run `python -m pytest` to run tests in nix setup action

do not rely on poetry to run tests, we'd like to not rely on it from
within the nix environment

* fix: do not rely on `npx` as its not on nix develop

use manual `supabase start` command instead

* fix: add `--command` typo

* chore: change name of nix develop command part

* feat: add coverage information and upload it to coveralls

* chore: run nix setup tests in both ubuntu and macos

* chore: undo macos latest as it apparently does not work

for some reason, `supabase start` does not seem to work on
`macos-latest` host in CI

* chore: switch from pylsp + mypy to pyright (basepyright)

it seems that pyright is better and faster than mypy in most if not
all cases, and the only reason I wasn't using it before is because it
was not working with eglot's default config.

* feat: add `basedpyright` to `pyproject.toml` instead of flake.nix only

this makes it available to all users, preparing for including it into
the standard CI later on

* fix: change back to python-lsp-server with pylsp-mypy

* chore: improve type definitions files to make mypy happy

there's still a lot of issues that I intend to improve upon. this is
just an initial set of changes to ensure that `mypy` doesn't
complain. we should further strengthen the mypy checking
rules (currently there are a lot of `Any`s in the code) as the code improves.

* chore: add mypy to CI as a step

* fix: run mypy through poetry

* fix: check for phx_ref_prev before calling del

* fix: use `python -m mypy` instead of `mypy` directly

* fix: run type check after `make run_tests` so that `poetry install` is ran

* fix: `StrEnum` does not exist in python 3.9

* fix: remove `phx_ref` from presence dict

* fix: add deprecation warning for calling send with dicts

* fix: fix config payload, improve more types

* fix: make typing definitions compatible with python 3.9

also enforce that mypy uses python 3.9 only rules.

* fix: import annotations for 3.9 to not complain about type error

* fix: move mypy to Makefile `run_tests`, run it before actual tests

* format: apply ruff reformating

* fix: try explicitly annotating `Callback` as a `TypeAlias`

* fix: change `dict` to `Dict` type due to 3.9

* fix: finally, import annotations to stop runtime from breaking

* format: apply ruff format one last time

* format: reorder import order in test_connection

* format: trim whitespace

Author: o-santi

.github/workflows/ci.yml

@@ -53,6 +53,8 @@ jobs:
           github_access_token: ${{ secrets.GITHUB_TOKEN }}
       - name: Clone Repository
         uses: actions/checkout@v4
+      - name: Type check
+        run: nix develop --command mypy ./realtime
       - name: Start Supabase local development setup
         run: nix develop --command supabase start --workdir infra -x studio,mailpit,edge-runtime,logflare,vector,supavisor,imgproxy,storage-api
       - name: Run python tests through nix

Makefile

@@ -5,11 +5,14 @@ install_poetry:
 	curl -sSL https://install.python-poetry.org | python -
 	poetry install
 
-tests: install tests_only tests_pre_commit
+tests: install run_mypy tests_only tests_pre_commit
 
 tests_pre_commit:
 	poetry run pre-commit run --all-files
 
+run_mypy:
+	poetry run mypy ./realtime
+
 run_infra:
 	npx supabase start --workdir infra -x studio,mailpit,edge-runtime,logflare,vector,supavisor,imgproxy,storage-api
 

poetry.lock

@@ -13,6 +13,7 @@ requires-python = ">=3.9"
 dependencies = [
   "websockets >=11,<16",
   "typing-extensions >=4.14.0",
+  "pydantic (>=2.11.7,<3.0.0)",
 ]
 
 [tool.poetry.group.dev.dependencies]
@@ -77,3 +78,13 @@ keep-runtime-typing = true
 [tool.pytest.ini_options]
 asyncio_mode = "strict"
 asyncio_default_fixture_loop_scope = "function"
+
+[tool.mypy]
+python_version = "3.9"
+check_untyped_defs = true
+allow_redefinition = true
+
+warn_return_any = true
+warn_unused_configs = true
+warn_redundant_casts = true
+warn_unused_ignores = true

pyproject.toml

@@ -3,7 +3,7 @@
 import asyncio
 import json
 import logging
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Mapping, Optional
 
 from realtime.types import (
     Binding,
@@ -16,6 +16,7 @@
     RealtimeSubscribeStates,
 )
 
+from ..message import Message
 from ..transformers import http_endpoint_url
 from .presence import (
     AsyncRealtimePresence,
@@ -52,23 +53,27 @@ def __init__(
         :param params: Optional parameters for connection.
         """
         self.socket = socket
-        self.params = params or {}
-        if self.params.get("config") is None:
-            self.params["config"] = {
-                "broadcast": {"ack": False, "self": False},
-                "presence": {"key": ""},
-                "private": False,
+        self.params: RealtimeChannelOptions = (
+            params
+            if params
+            else {
+                "config": {
+                    "broadcast": {"ack": False, "self": False},
+                    "presence": {"key": ""},
+                    "private": False,
+                }
             }
+        )
 
         self.topic = topic
         self._joined_once = False
-        self.bindings: Dict[str, List[Binding]] = {}
+        self.bindings: dict[str, list[Binding]] = {}
         self.presence = AsyncRealtimePresence(self)
         self.state = ChannelStates.CLOSED
-        self._push_buffer: List[AsyncPush] = []
+        self._push_buffer: list[AsyncPush] = []
         self.timeout = self.socket.timeout
 
-        self.join_push = AsyncPush(self, ChannelEvents.join, self.params)
+        self.join_push: AsyncPush = AsyncPush(self, ChannelEvents.join, self.params)
         self.rejoin_timer = AsyncTimer(
             self._rejoin_until_connected, lambda tries: 2**tries
         )
@@ -111,8 +116,9 @@ def on_error(payload, *args):
         self._on("close", on_close)
         self._on("error", on_error)
 
-        def on_reply(payload, ref):
-            self._trigger(self._reply_event_name(ref), payload)
+        def on_reply(payload: Dict[str, Any], ref: Optional[str]):
+            if ref:
+                self._trigger(self._reply_event_name(ref), payload)
 
         self._on(ChannelEvents.reply, on_reply)
 
@@ -169,22 +175,24 @@ async def subscribe(
             presence = config.get("presence", {})
             private = config.get("private", False)
 
-            access_token_payload = {}
-            config = {
-                "broadcast": broadcast,
-                "presence": presence,
-                "private": private,
-                "postgres_changes": list(
-                    map(lambda x: x.filter, self.bindings.get("postgres_changes", []))
-                ),
+            config_payload: Dict[str, Any] = {
+                "config": {
+                    "broadcast": broadcast,
+                    "presence": presence,
+                    "private": private,
+                    "postgres_changes": list(
+                        map(
+                            lambda x: x.filter,
+                            self.bindings.get("postgres_changes", []),
+                        )
+                    ),
+                }
             }
 
             if self.socket.access_token:
-                access_token_payload["access_token"] = self.socket.access_token
+                config_payload["access_token"] = self.socket.access_token
 
-            self.join_push.update_payload(
-                {**{"config": config}, **access_token_payload}
-            )
+            self.join_push.update_payload(config_payload)
             self._joined_once = True
 
             def on_join_push_ok(payload: Dict[str, Any]):
@@ -253,7 +261,7 @@ def on_join_push_timeout(*args):
 
         return self
 
-    async def unsubscribe(self):
+    async def unsubscribe(self) -> None:
         """
         Unsubscribe from the channel and leave the topic.
         Sets channel state to LEAVING and cleans up timers and pushes.
@@ -263,9 +271,9 @@ async def unsubscribe(self):
         self.rejoin_timer.reset()
         self.join_push.destroy()
 
-        def _close(*args):
+        def _close(*args) -> None:
             logger.info(f"channel {self.topic} leave")
-            self._trigger(ChannelEvents.close, "leave")
+            self._trigger(ChannelEvents.close, {})
 
         leave_push = AsyncPush(self, ChannelEvents.leave, {})
         leave_push.receive("ok", _close).receive("timeout", _close)
@@ -310,21 +318,24 @@ async def join(self) -> AsyncRealtimeChannel:
         :return: Channel
         """
         try:
-            await self.socket.send(
-                {
-                    "topic": self.topic,
-                    "event": "phx_join",
-                    "payload": {"config": self.params},
-                    "ref": None,
-                }
+            message = Message(
+                topic=self.topic,
+                event=ChannelEvents.join,
+                payload={"config": self.params},
+                ref=None,
             )
+            await self.socket.send(message)
+            return self
         except Exception as e:
             print(e)
             return self
 
     # Event handling methods
     def _on(
-        self, type: str, callback: Callback, filter: Optional[Dict[str, Any]] = None
+        self,
+        type: str,
+        callback: Callback[[Dict[str, Any], Optional[str]], None],
+        filter: Optional[Dict[str, Any]] = None,
     ) -> AsyncRealtimeChannel:
         """
         Set up a listener for a specific event.
@@ -411,7 +422,7 @@ def on_postgres_changes(
         )
 
     def on_system(
-        self, callback: Callable[[Dict[str, Any], None]]
+        self, callback: Callable[[Dict[str, Any]], None]
     ) -> AsyncRealtimeChannel:
         """
         Set up a listener for system events.
@@ -508,7 +519,7 @@ def _can_push(self):
     async def send_presence(self, event: str, data: Any) -> None:
         await self.push(ChannelEvents.presence, {"event": event, "payload": data})
 
-    def _trigger(self, type: str, payload: Optional[Any], ref: Optional[str] = None):
+    def _trigger(self, type: str, payload: Dict[str, Any], ref: Optional[str] = None):
         type_lowercase = type.lower()
         events = [
             ChannelEvents.close,
@@ -562,7 +573,7 @@ def _trigger(self, type: str, payload: Optional[Any], ref: Optional[str] = None)
                 elif binding.type == type_lowercase:
                     binding.callback(payload, ref)
 
-    def _reply_event_name(self, ref: str):
+    def _reply_event_name(self, ref: str) -> str:
         return f"chan_reply_{ref}"
 
     async def _rejoin_until_connected(self):

realtime/_async/channel.py

@@ -3,12 +3,12 @@
 import logging
 import re
 from functools import wraps
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional, Union
 from urllib.parse import urlencode, urlparse, urlunparse
 
 import websockets
 from websockets import connect
-from websockets.client import ClientProtocol
+from websockets.asyncio.client import ClientConnection
 
 from ..exceptions import NotConnectedError
 from ..message import Message
@@ -62,7 +62,7 @@ def __init__(
         :param timeout: Connection timeout in seconds. Defaults to DEFAULT_TIMEOUT.
         """
         if not is_ws_url(url):
-            ValueError("url must be a valid WebSocket URL or HTTP URL string")
+            raise ValueError("url must be a valid WebSocket URL or HTTP URL string")
         self.url = f"{re.sub(r'https://', 'wss://', re.sub(r'http://', 'ws://', url, flags=re.IGNORECASE), flags=re.IGNORECASE)}/websocket"
         if token:
             self.url += f"?apikey={token}"
@@ -72,7 +72,7 @@ def __init__(
         self.access_token = token
         self.send_buffer: List[Callable] = []
         self.hb_interval = hb_interval
-        self._ws_connection: Optional[ClientProtocol] = None
+        self._ws_connection: Optional[ClientConnection] = None
         self.ref = 0
         self.auto_reconnect = auto_reconnect
         self.channels: Dict[str, AsyncRealtimeChannel] = {}
@@ -97,13 +97,15 @@ async def _listen(self) -> None:
 
         try:
             async for msg in self._ws_connection:
-                logger.info(f"receive: {msg}")
+                logger.info(f"receive: {msg!r}")
 
-                msg = Message(**json.loads(msg))
-                channel = self.channels.get(msg.topic)
+                message = Message.model_validate_json(msg)
+                channel = self.channels.get(message.topic)
 
                 if channel:
-                    channel._trigger(msg.event, msg.payload, msg.ref)
+                    channel._trigger(
+                        message.event, dict(**message.payload), message.ref
+                    )
         except websockets.exceptions.ConnectionClosedError as e:
             await self._on_connect_error(e)
 
@@ -236,7 +238,7 @@ async def _heartbeat(self) -> None:
 
         while self.is_connected:
             try:
-                data = dict(
+                data = Message(
                     topic=PHOENIX_CHANNEL,
                     event=ChannelEvents.heartbeat,
                     payload={},
@@ -294,14 +296,6 @@ async def remove_all_channels(self) -> None:
 
         await self.close()
 
-    def summary(self) -> None:
-        """
-        Prints a list of topics and event the socket is listening to
-        :return: None
-        """
-        for topic, channel in self.channels.items():
-            print(f"Topic: {topic} | Events: {[e for e, _ in channel.listeners]}]")
-
     async def set_auth(self, token: Optional[str]) -> None:
         """
         Set the authentication token for the connection and update all joined channels.
@@ -325,7 +319,7 @@ def _make_ref(self) -> str:
         self.ref += 1
         return f"{self.ref}"
 
-    async def send(self, message: Dict[str, Any]) -> None:
+    async def send(self, message: Union[Message, Dict[str, Any]]) -> None:
         """
         Send a message through the WebSocket connection.
 
@@ -340,16 +334,22 @@ async def send(self, message: Dict[str, Any]) -> None:
         Returns:
             None
         """
-
-        message = json.dumps(message)
-        logger.info(f"send: {message}")
+        if isinstance(message, Message):
+            msg = message
+        else:
+            logger.warning(
+                "Warning: calling AsyncRealtimeClient.send with a dictionary is deprecated. Please call it with a Message object instead. This will be a hard error in the future."
+            )
+            msg = Message(**message)
+        message_str = msg.model_dump_json()
+        logger.info(f"send: {message_str}")
 
         async def send_message():
             if not self._ws_connection:
                 raise NotConnectedError("_send")
 
             try:
-                await self._ws_connection.send(message)
+                await self._ws_connection.send(message_str)
             except websockets.exceptions.ConnectionClosedError as e:
                 await self._on_connect_error(e)
             except websockets.exceptions.ConnectionClosedOK: