Add mypy type checking for rooms/ module (#133)

* fix all mypy errors under rooms/

* add typing-checks workflow

* rename CI job

* add type annotation to _rooms_by_id

Author: dlqqq

.github/workflows/typing-checks.yml

@@ -0,0 +1,44 @@
+name: Typing Checks
+
+on:
+  push:
+    # IMPORTANT: update these paths when we expand type-checking coverage
+    branches: [main]
+    paths:
+      - 'jupyter_server_documents/rooms/**'
+      - '.github/workflows/typing-checks.yml'
+  pull_request:
+    branches: [main]
+    paths:
+      - 'jupyter_server_documents/rooms/**'
+      - '.github/workflows/typing-checks.yml'
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  python:
+    name: 'Python'
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ['3.10', '3.11', '3.12', '3.13']
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Base Setup
+        uses: jupyterlab/maintainer-tools/.github/actions/base-setup@v1
+        with:
+          python_version: ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: |
+          python -m pip install mypy
+          python -m pip install -e .[test]
+
+      - name: Type-check `rooms` module
+        run: |
+          mypy jupyter_server_documents/rooms

jupyter_server_documents/rooms/yroom.py

@@ -17,11 +17,11 @@
 
 if TYPE_CHECKING:
     import logging
-    from typing import Coroutine, Literal, Tuple, Any
+    from typing import Callable, Coroutine, Literal, Tuple, Any
     from .yroom_manager import YRoomManager
-    from jupyter_server_fileid.manager import BaseFileIdManager
+    from jupyter_server_fileid.manager import BaseFileIdManager  # type: ignore
     from jupyter_server.services.contents.manager import ContentsManager
-    from pycrdt import TransactionEvent
+    from pycrdt import TransactionEvent, Subscription
     from ..outputs.manager import OutputsManager
 
 class YRoom(LoggingConfigurable):
@@ -136,7 +136,7 @@ class YRoom(LoggingConfigurable):
     documentation for more info.
     """
 
-    _jupyter_ydoc_observers: dict[str, callable[[str, Any], Any]]
+    _jupyter_ydoc_observers: dict[str, Callable[[str, Any], Any]]
     """
     Dictionary of JupyterYDoc observers added by consumers of this room.
 
@@ -167,10 +167,10 @@ class YRoom(LoggingConfigurable):
     `self._message_queue.put_nowait(None)`.
     """
 
-    _awareness_subscription: pycrdt.Subscription
+    _awareness_subscription: str
     """Subscription to awareness changes."""
 
-    _ydoc_subscription: pycrdt.Subscription
+    _ydoc_subscription: Subscription
     """Subscription to YDoc changes."""
 
     _stopped: bool
@@ -346,6 +346,8 @@ async def get_jupyter_ydoc(self) -> YBaseDoc:
             message = "There is no Jupyter ydoc for global awareness scenario"
             self.log.error(message)
             raise Exception(message)
+        if self._jupyter_ydoc is None:
+            raise RuntimeError("Jupyter YDoc is not available")
         if self.file_api:
             await self.file_api.until_content_loaded
         return self._jupyter_ydoc
@@ -428,7 +430,7 @@ def handle_message(self, client_id: str, message: bytes) -> None:
 
         # Determine message type & subtype from header
         message_type = message[0]
-        sync_message_subtype = "*"
+        sync_message_subtype = -1 # invalid sentinel value
         # message subtypes only exist on sync messages, hence this condition
         if message_type == YMessageType.SYNC and len(message) >= 2:
             sync_message_subtype = message[1]
@@ -585,7 +587,7 @@ def _on_ydoc_update(self, event: TransactionEvent) -> None:
         self._broadcast_message(message, message_type="SyncUpdate")
 
 
-    def observe_jupyter_ydoc(self, observer: callable[[str, Any], Any]) -> str:
+    def observe_jupyter_ydoc(self, observer: Callable[[str, Any], Any]) -> str:
         """
         Adds an observer callback to the JupyterYDoc that fires on change.
         The callback should accept 2 arguments:
@@ -604,8 +606,9 @@ def observe_jupyter_ydoc(self, observer: callable[[str, Any], Any]) -> str:
         Returns an `observer_id: str` that can be passed to
         `unobserve_jupyter_ydoc()` to remove the observer.
         """
-        observer_id = uuid.uuid4()
+        observer_id = str(uuid.uuid4())
         self._jupyter_ydoc_observers[observer_id] = observer
+        return observer_id
 
 
     def unobserve_jupyter_ydoc(self, observer_id: str):
@@ -745,7 +748,10 @@ def _on_awareness_update(self, type: str, changes: tuple[dict[str, Any], Any]) -
         self.log.debug(f"awareness update, updated_clients={updated_clients}")
         state = self._awareness.encode_awareness_update(updated_clients)
         message = pycrdt.create_awareness_message(state)
-        self.log.debug(f"awareness update, message={message}")
+        # !r ensures binary messages show as `b'...'`  instead of being decoded
+        # into jargon in log statements.
+        # https://docs.python.org/3/library/string.html#format-string-syntax
+        self.log.debug(f"awareness update, message={message!r}")
         self._broadcast_message(message, "AwarenessUpdate")
 
 
@@ -827,8 +833,11 @@ def stop(self, close_code: int = 1001, immediately: bool = False) -> None:
                 self._message_queue.get_nowait()
                 self._message_queue.task_done()
             else:
-                client_id, message = self._message_queue.get_nowait()
-                self.handle_message(client_id, message)
+                queue_item = self._message_queue.get_nowait()
+                if queue_item is not None:
+                    client_id, message = queue_item
+                    self.handle_message(client_id, message)
+                self._message_queue.task_done()
 
         # Stop the `_process_message_queue` task by enqueueing `None`
         self._message_queue.put_nowait(None)
@@ -924,8 +933,9 @@ def restart(self, close_code: int = 1001, immediately: bool = False) -> None:
         self.clients.restart()
 
         # Restart `YRoomFileAPI` & reload the document
-        self.file_api.restart()
-        self.file_api.load_content_into(self._jupyter_ydoc)
+        if self.file_api is not None and self._jupyter_ydoc is not None:
+            self.file_api.restart()
+            self.file_api.load_content_into(self._jupyter_ydoc)
 
         # Restart `_process_message_queue()` task
         asyncio.create_task(self._process_message_queue())
@@ -952,16 +962,16 @@ def should_ignore_state_update(event: pycrdt.MapEvent) -> bool:
     # `False` immediately if:
     # - a key was updated to a value different from the previous value
     # - a key was added with a value different from the previous value
-    for key in event.keys.keys():
-        update_info = event.keys[key]
+    for key in getattr(event, 'keys', {}).keys():
+        update_info = getattr(event, 'keys', {})[key]
         action = update_info.get('action', None)
         if action == 'update':
             old_value = update_info.get('oldValue', None)
             new_value = update_info.get('newValue', None)
             if old_value != new_value:
                 return False
         elif action == "add":
-            old_value = event.target.get(key, None)
+            old_value = getattr(event, 'target', {}).get(key, None)
             new_value = update_info.get('newValue', None)
             if old_value != new_value:
                 return False

jupyter_server_documents/rooms/yroom_events_api.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 from jupyter_events import EventLogger
-from jupyter_server_fileid.manager import BaseFileIdManager
+from jupyter_server_fileid.manager import BaseFileIdManager  # type: ignore
 from traitlets.config import LoggingConfigurable
 from typing import TYPE_CHECKING
 

jupyter_server_documents/rooms/yroom_file_api.py

@@ -12,7 +12,7 @@
 if TYPE_CHECKING:
     from typing import Any, Coroutine, Literal
     from .yroom import YRoom
-    from jupyter_server_fileid.manager import BaseFileIdManager
+    from jupyter_server_fileid.manager import BaseFileIdManager  # type: ignore
     from jupyter_server.services.contents.manager import ContentsManager
     from ..outputs.manager import OutputsManager
 
@@ -56,7 +56,6 @@ class YRoomFileAPI(LoggingConfigurable):
 
     # See `filemanager.py` in `jupyter_server` for references on supported file
     # formats & file types.
-    room_id: str
     file_format: Literal["text", "base64"]
     file_type: Literal["file", "notebook"]
     file_id: str

jupyter_server_documents/rooms/yroom_manager.py

@@ -1,16 +1,17 @@
 from __future__ import annotations
 
 from .yroom import YRoom
-from typing import TYPE_CHECKING
+from typing import cast, TYPE_CHECKING
 import asyncio
 import traitlets
 from traitlets.config import LoggingConfigurable
-from jupyter_server_fileid.manager import BaseFileIdManager
+from jupyter_server_fileid.manager import BaseFileIdManager  # type: ignore
 
 from ..outputs.manager import OutputsManager
 
 if TYPE_CHECKING:
     import logging
+    from typing import Set
     from jupyter_server.extension.application import ExtensionApp
     from jupyter_server.services.contents.manager import ContentsManager
     from jupyter_events import EventLogger
@@ -53,7 +54,7 @@ class YRoomManager(LoggingConfigurable):
     this declaration only hints the type for type checkers.
     """
 
-    _rooms_by_id: dict[str, YRoom] = traitlets.Dict(default_value={})
+    _rooms_by_id: traitlets.Dict[str, YRoom] = traitlets.Dict(default_value={})
     """
     Dictionary of active `YRoom` instances, keyed by room ID. Rooms are never
     deleted from this dictionary.
@@ -62,11 +63,12 @@ class YRoomManager(LoggingConfigurable):
     out-of-band. See #116.
     """
 
-    _inactive_rooms: set[str] = traitlets.Set()
+    _inactive_rooms = traitlets.Set()
     """
-    Set of room IDs that were marked inactive on the last iteration of
-    `_watch_rooms()`. If a room is inactive and its ID is present in this set,
-    then the room should be restarted as it has been inactive for >10 seconds.
+    Set of room IDs (as strings) that were marked inactive on the last iteration
+    of `_watch_rooms()`. If a room is inactive and its ID is present in this
+    set, then the room should be restarted as it has been inactive for >10
+    seconds.
     """
 
     _watch_rooms_task: asyncio.Task | None
@@ -84,23 +86,34 @@ def __init__(self, *args, **kwargs):
 
     @property
     def fileid_manager(self) -> BaseFileIdManager:
+        if self.parent.serverapp is None:
+            raise RuntimeError("ServerApp is not available")
         manager = self.parent.serverapp.web_app.settings.get("file_id_manager", None)
         assert isinstance(manager, BaseFileIdManager)
         return manager
 
 
     @property
     def contents_manager(self) -> ContentsManager:
+        if self.parent.serverapp is None:
+            raise RuntimeError("ServerApp is not available")
         return self.parent.serverapp.contents_manager
 
 
     @property
     def event_logger(self) -> EventLogger:
-        return self.parent.serverapp.event_logger
+        if self.parent.serverapp is None:
+            raise RuntimeError("ServerApp is not available")
+        event_logger = self.parent.serverapp.event_logger
+        if event_logger is None:
+            raise RuntimeError("Event logger is not available")
+        return event_logger
 
 
     @property
     def outputs_manager(self) -> OutputsManager:
+        if not hasattr(self.parent, 'outputs_manager'):
+            raise RuntimeError("Outputs manager is not available")
         return self.parent.outputs_manager
 
 
@@ -156,17 +169,17 @@ def delete_room(self, room_id: str) -> None:
         """
         yroom = self._rooms_by_id.pop(room_id, None)
         if not yroom:
-            return
+            return None
 
         self.log.info(f"Stopping YRoom '{room_id}'.")
         try:
             yroom.stop()
-            return True
+            return None
         except Exception as e:
             self.log.exception(
                 f"Exception raised when stopping YRoom '{room_id}: "
             )
-            return False
+            return None
 
 
     async def _watch_rooms(self) -> None:
@@ -274,9 +287,9 @@ async def stop(self) -> None:
 
         # Define task that deletes the room and waits until the content is saved
         async def delete_then_save(room_id: str, room: YRoom):
-            ret = self.delete_room(room_id)
+            self.delete_room(room_id)
             await room.until_saved
-            return ret
+            return None
 
         # Delete all rooms concurrently using `delete_then_save()`
         for room_id, room in self._rooms_by_id.items():