Inlining no-longer-invariant _sessions access

Author: blast-hardcheese

src/replit_river/client_transport.py

@@ -1,7 +1,7 @@
 import asyncio
 import logging
 from collections.abc import Awaitable, Callable
-from typing import Generic, Mapping, assert_never
+from typing import Generic, assert_never
 
 import nanoid
 import websockets
@@ -71,11 +71,8 @@ def __init__(
         # We want to make sure there's only one session creation at a time
         self._create_session_lock = asyncio.Lock()
 
-    async def _close_all_sessions(
-        self,
-        get_all_sessions: Callable[[], Mapping[str, Session]],
-    ) -> None:
-        sessions = get_all_sessions().values()
+    async def _close_all_sessions(self) -> None:
+        sessions = self._sessions.values()
         logger.info(
             f"start closing sessions {self._transport_id}, number sessions : "
             f"{len(sessions)}"
@@ -94,7 +91,7 @@ def generate_nanoid(self) -> str:
 
     async def close(self) -> None:
         self._rate_limiter.close()
-        await self._close_all_sessions(self._get_all_sessions)
+        await self._close_all_sessions()
 
     async def get_or_create_session(self) -> ClientSession:
         async with self._create_session_lock:
@@ -235,7 +232,7 @@ async def _create_new_session(
 
     async def _retry_connection(self) -> ClientSession:
         if not self._transport_options.transparent_reconnect:
-            await self._close_all_sessions(self._get_all_sessions)
+            await self._close_all_sessions()
         return await self.get_or_create_session()
 
     async def _send_handshake_request(
@@ -377,9 +374,6 @@ async def _establish_handshake(
             )
         return handshake_request, handshake_response
 
-    def _get_all_sessions(self) -> Mapping[str, Session]:
-        return self._sessions
-
     async def _delete_session(self, session: Session) -> None:
         async with self._session_lock:
             if session._to_id in self._sessions:

src/replit_river/server_transport.py

@@ -1,6 +1,6 @@
 import asyncio
 import logging
-from typing import Any, Callable, Mapping
+from typing import Any
 
 import nanoid  # type: ignore  # type: ignore
 from pydantic import ValidationError
@@ -51,11 +51,8 @@ def __init__(
         self._handlers: dict[tuple[str, str], tuple[str, GenericRpcHandlerBuilder]] = {}
         self._session_lock = asyncio.Lock()
 
-    async def _close_all_sessions(
-        self,
-        get_all_sessions: Callable[[], Mapping[str, Session]],
-    ) -> None:
-        sessions = get_all_sessions().values()
+    async def _close_all_sessions(self) -> None:
+        sessions = self._sessions.values()
         logger.info(
             f"start closing sessions {self._transport_id}, number sessions : "
             f"{len(sessions)}"
@@ -111,7 +108,7 @@ async def handshake_to_get_session(
         raise WebsocketClosedException("No handshake message received")
 
     async def close(self) -> None:
-        await self._close_all_sessions(self._get_all_sessions)
+        await self._close_all_sessions()
 
     async def _get_existing_session(self, to_id: str) -> ServerSession | None:
         async with self._session_lock:
@@ -315,9 +312,6 @@ async def _establish_handshake(
 
         return handshake_request, handshake_response
 
-    def _get_all_sessions(self) -> Mapping[str, Session]:
-        return self._sessions
-
     async def _delete_session(self, session: Session) -> None:
         async with self._session_lock:
             if session._to_id in self._sessions: