feat(channels): refine type hints and improve async handling in SQLSpecChannelsBackend

Author: cofin

sqlspec/extensions/events/channel.py

@@ -114,7 +114,7 @@ def __init__(self, config: "DatabaseConfigProtocol[Any, Any, Any]") -> None:
         self._config = config
         self._backend_name = backend_label
         self._is_async = bool(config.is_async)
-        self._portal_bridge = bool(extension_settings.get("portal_bridge", False)) and self._is_async
+        self._portal_bridge = bool(extension_settings.get("portal_bridge")) and self._is_async
         self._portal = None
         self._runtime = config.get_observability_runtime()
         self._listeners_async: dict[str, AsyncEventListener] = {}
@@ -626,7 +626,7 @@ async def shutdown_async(self) -> None:
             if backend_shutdown is not None and callable(backend_shutdown):
                 result = backend_shutdown()
                 if result is not None:
-                    await result  # type: ignore[misc]
+                    await result
         except Exception as error:
             self._end_event_span(span, error=error)
             raise

sqlspec/extensions/litestar/__init__.py

@@ -1,5 +1,5 @@
-from sqlspec.extensions.litestar.cli import database_group
 from sqlspec.extensions.litestar.channels import SQLSpecChannelsBackend
+from sqlspec.extensions.litestar.cli import database_group
 from sqlspec.extensions.litestar.config import LitestarConfig
 from sqlspec.extensions.litestar.plugin import (
     DEFAULT_COMMIT_MODE,

sqlspec/extensions/litestar/channels.py

@@ -4,14 +4,17 @@
 import base64
 import hashlib
 import re
-from collections.abc import AsyncGenerator, Iterable
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from litestar.channels.backends.base import ChannelsBackend
 
-from sqlspec.extensions.events.channel import EventChannel
 from sqlspec.utils.logging import get_logger
 
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator, Iterable
+
+    from sqlspec.extensions.events.channel import EventChannel
+
 logger = get_logger("extensions.litestar.channels")
 
 _IDENTIFIER_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
@@ -42,11 +45,7 @@ class SQLSpecChannelsBackend(ChannelsBackend):
     )
 
     def __init__(
-        self,
-        event_channel: "EventChannel",
-        *,
-        channel_prefix: str = "litestar",
-        poll_interval: float = 0.2,
+        self, event_channel: "EventChannel", *, channel_prefix: str = "litestar", poll_interval: float = 0.2
     ) -> None:
         if not _IDENTIFIER_PATTERN.match(channel_prefix):
             msg = f"channel_prefix must be a valid identifier, got: {channel_prefix!r}"
@@ -57,9 +56,9 @@ def __init__(
         self._event_channel = event_channel
         self._channel_prefix = channel_prefix
         self._poll_interval = poll_interval
-        self._output_queue: "asyncio.Queue[tuple[str, bytes]] | None" = None
+        self._output_queue: asyncio.Queue[tuple[str, bytes]] | None = None
         self._shutdown = asyncio.Event()
-        self._tasks: dict[str, "asyncio.Task[None]"] = {}
+        self._tasks: dict[str, asyncio.Task[None]] = {}
         self._to_db_channel: dict[str, str] = {}
         self._to_litestar_channel: dict[str, str] = {}
 
@@ -95,7 +94,7 @@ async def subscribe(self, channels: "Iterable[str]") -> None:
             self._tasks[channel] = task
 
     async def unsubscribe(self, channels: "Iterable[str]") -> None:
-        cancelled: list["asyncio.Task[None]"] = []
+        cancelled: list[asyncio.Task[None]] = []
         for channel in channels:
             task = self._tasks.pop(channel, None)
             if task is None:

tests/integration/test_extensions/test_litestar/test_channels_backend.py

@@ -1,6 +1,6 @@
 import asyncio
 import tempfile
-from typing import Any
+from typing import Any, cast
 
 import msgspec.json
 import pytest
@@ -14,7 +14,7 @@
 
 async def _next_event(subscriber: "Any") -> bytes:
     async for event in subscriber.iter_events():
-        return event
+        return cast("bytes", event)
     msg = "Subscriber stopped without yielding an event"
     raise RuntimeError(msg)