Typing all self parameters

Author: blast-hardcheese

src/replit_river/session.py

@@ -1,6 +1,6 @@
 import asyncio
 import logging
-from typing import Any, Awaitable, Callable, Coroutine
+from typing import Any, Awaitable, Callable, Coroutine, TypeAlias
 
 import nanoid  # type: ignore
 import websockets
@@ -36,27 +36,55 @@
 trace_propagator = TraceContextTextMapPropagator()
 trace_setter = TransportMessageTracingSetter()
 
+CloseSessionCallback: TypeAlias = Callable[["Session"], Coroutine[Any, Any, Any]]
+RetryConnectionCallback: TypeAlias = (
+            Callable[
+                [],
+                Coroutine[Any, Any, Any],
+            ]
+        )
+
 
 class Session:
     """Common functionality shared between client_session and server_session"""
 
+    _transport_id: str
+    _to_id: str
+    session_id: str
+    _transport_options: TransportOptions
+
+    # session state, only modified during closing
+    _state: SessionState
+    _state_lock: asyncio.Lock
+    _close_session_callback: CloseSessionCallback
+    _close_session_after_time_secs: float | None
+
+    # ws state
+    _ws_lock: asyncio.Lock
+    _ws_wrapper: WebsocketWrapper
+    _heartbeat_misses: int
+    _retry_connection_callback: RetryConnectionCallback | None
+
+    # stream for tasks
+    _stream_lock: asyncio.Lock
     _streams: dict[str, Channel[Any]]
 
+    # book keeping
+    _seq_manager: SeqManager
+    _msg_lock: asyncio.Lock
+    _buffer: MessageBuffer
+    _task_manager: BackgroundTaskManager
+
+
     def __init__(
         self,
         transport_id: str,
         to_id: str,
         session_id: str,
         websocket: websockets.WebSocketCommonProtocol,
         transport_options: TransportOptions,
-        close_session_callback: Callable[["Session"], Coroutine[Any, Any, Any]],
-        retry_connection_callback: (
-            Callable[
-                [],
-                Coroutine[Any, Any, Any],
-            ]
-            | None
-        ) = None,
+        close_session_callback: CloseSessionCallback,
+        retry_connection_callback: RetryConnectionCallback | None = None,
     ) -> None:
         self._transport_id = transport_id
         self._to_id = to_id

src/replit_river/v2/client_session.py

@@ -34,7 +34,7 @@
     InvalidMessageException,
     OutOfOrderMessageException,
 )
-from replit_river.session import Session
+from replit_river.session import CloseSessionCallback, RetryConnectionCallback, Session
 from replit_river.transport_options import MAX_MESSAGE_BUFFER_SIZE, TransportOptions
 
 STREAM_CANCEL_BIT_TYPE = Literal[0b00100]
@@ -54,14 +54,8 @@ def __init__(
         session_id: str,
         websocket: websockets.WebSocketCommonProtocol,
         transport_options: TransportOptions,
-        close_session_callback: Callable[[Session], Coroutine[Any, Any, Any]],
-        retry_connection_callback: (
-            Callable[
-                [],
-                Coroutine[Any, Any, Any],
-            ]
-            | None
-        ) = None,
+        close_session_callback: CloseSessionCallback,
+        retry_connection_callback: RetryConnectionCallback | None = None,
     ) -> None:
         super().__init__(
             transport_id=transport_id,