Resolving circular import

Author: blast-hardcheese

src/replit_river/common_session.py

@@ -1,14 +1,40 @@
 import asyncio
+import enum
 import logging
-from typing import Awaitable, Callable
+from typing import Any, Awaitable, Callable, Protocol
+
+from opentelemetry.trace import Span
 
 from replit_river.messages import FailedSendingMessageException
 from replit_river.rpc import ACK_BIT
-from replit_river.session import SendMessage, SessionState
 
 logger = logging.getLogger(__name__)
 
 
+class SendMessage(Protocol):
+    async def __call__(
+        self,
+        *,
+        stream_id: str,
+        payload: dict[Any, Any] | str,
+        control_flags: int,
+        service_name: str | None,
+        procedure_name: str | None,
+        span: Span | None,
+    ) -> None: ...
+
+
+class SessionState(enum.Enum):
+    """The state a session can be in.
+
+    Can only transition from ACTIVE to CLOSING to CLOSED.
+    """
+
+    ACTIVE = 0
+    CLOSING = 1
+    CLOSED = 2
+
+
 async def setup_heartbeat(
     session_id: str,
     heartbeat_ms: float,

src/replit_river/session.py

@@ -1,15 +1,14 @@
 import asyncio
-import enum
 import logging
-from typing import Any, Awaitable, Callable, Coroutine, Protocol
+from typing import Any, Awaitable, Callable, Coroutine
 
 import nanoid  # type: ignore
 import websockets
 from aiochannel import Channel, ChannelClosed
 from opentelemetry.trace import Span, use_span
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 
-from replit_river.common_session import setup_heartbeat
+from replit_river.common_session import SessionState, setup_heartbeat
 from replit_river.message_buffer import MessageBuffer, MessageBufferClosedError
 from replit_river.messages import (
     FailedSendingMessageException,
@@ -36,30 +35,6 @@
 trace_setter = TransportMessageTracingSetter()
 
 
-class SendMessage(Protocol):
-    async def __call__(
-        self,
-        *,
-        stream_id: str,
-        payload: dict[Any, Any] | str,
-        control_flags: int,
-        service_name: str | None,
-        procedure_name: str | None,
-        span: Span | None,
-    ) -> None: ...
-
-
-class SessionState(enum.Enum):
-    """The state a session can be in.
-
-    Can only transition from ACTIVE to CLOSING to CLOSED.
-    """
-
-    ACTIVE = 0
-    CLOSING = 1
-    CLOSED = 2
-
-
 class Session:
     """Common functionality shared between client_session and server_session"""