Merge pull request #37 from modern-python/28-feature-websockets-in-litestar

add ws support for litestar

Author: lesnik512

packages/modern-di-litestar/modern_di_litestar/main.py

@@ -18,7 +18,7 @@ def setup_di(app: litestar.Litestar, scope: enum.IntEnum = DIScope.APP) -> Conta
 
 
 def prepare_di_dependencies() -> dict[str, Provide]:
-    return {"request_di_container": Provide(build_di_container)}
+    return {"di_container": Provide(build_di_container)}
 
 
 def fetch_di_container(app: litestar.Litestar) -> Container:
@@ -28,8 +28,14 @@ def fetch_di_container(app: litestar.Litestar) -> Container:
 async def build_di_container(
     request: litestar.Request[typing.Any, typing.Any, typing.Any],
 ) -> typing.AsyncIterator[Container]:
-    scope = DIScope.REQUEST
-    context = {"request": request}
+    context: dict[str, typing.Any] = {}
+    scope: DIScope | None
+    if isinstance(request, litestar.WebSocket):
+        context["websocket"] = request
+        scope = DIScope.SESSION
+    else:
+        context["request"] = request
+        scope = DIScope.REQUEST
     container: Container = fetch_di_container(request.app)
     async with container.build_child_container(context=context, scope=scope) as request_container:
         yield request_container
@@ -39,11 +45,9 @@ async def build_di_container(
 class _Dependency(typing.Generic[T_co]):
     dependency: providers.AbstractProvider[T_co]
 
-    async def __call__(
-        self, request_di_container: typing.Annotated[Container | None, Dependency()] = None
-    ) -> T_co | None:
-        assert request_di_container
-        return await self.dependency.async_resolve(request_di_container)
+    async def __call__(self, di_container: typing.Annotated[Container | None, Dependency()] = None) -> T_co | None:
+        assert di_container
+        return await self.dependency.async_resolve(di_container)
 
 
 def FromDI(dependency: providers.AbstractProvider[T_co]) -> Provide:  # noqa: N802

packages/modern-di-litestar/tests_litestar/test_routes.py

@@ -55,8 +55,8 @@ async def read_root(method: str) -> None:
 
 def test_factories_action_scope(client: TestClient[litestar.Litestar], app: litestar.Litestar) -> None:
     @litestar.get("/")
-    async def read_root(request_di_container: modern_di.Container) -> None:
-        with request_di_container.build_child_container() as action_container:
+    async def read_root(di_container: modern_di.Container) -> None:
+        with di_container.build_child_container() as action_container:
             action_factory_instance = action_factory.sync_resolve(action_container)
             assert isinstance(action_factory_instance, DependentCreator)
 

packages/modern-di-litestar/tests_litestar/test_websockets.py

@@ -1,12 +1,66 @@
+import typing
+
 import litestar
 from litestar.testing import TestClient
+from modern_di import Container, Scope, providers
+from modern_di_litestar import FromDI
+
+from tests_litestar.dependencies import DependentCreator, SimpleCreator
+
+
+def context_adapter_function(*, websocket: litestar.Request[typing.Any, typing.Any, typing.Any], **_: object) -> str:
+    assert isinstance(websocket, litestar.WebSocket)
+    return websocket.url.path
+
+
+app_factory = providers.Factory(Scope.APP, SimpleCreator, dep1="original")
+session_factory = providers.Factory(Scope.SESSION, DependentCreator, dep1=app_factory.cast)
+request_factory = providers.Factory(Scope.REQUEST, DependentCreator, dep1=app_factory.cast)
+context_adapter = providers.ContextAdapter(Scope.SESSION, context_adapter_function)
+
+
+async def test_factories(client: TestClient[litestar.Litestar], app: litestar.Litestar) -> None:
+    @litestar.websocket_listener(
+        "/ws",
+        dependencies={"app_factory_instance": FromDI(app_factory), "session_factory_instance": FromDI(session_factory)},
+    )
+    async def websocket_handler(
+        data: str,
+        app_factory_instance: SimpleCreator,
+        session_factory_instance: DependentCreator,
+    ) -> None:
+        assert data == "test"
+        assert isinstance(app_factory_instance, SimpleCreator)
+        assert isinstance(session_factory_instance, DependentCreator)
+        assert session_factory_instance.dep1 is not app_factory_instance
+
+    app.register(websocket_handler)
+
+    with client.websocket_connect("/ws") as websocket:
+        websocket.send("test")
+
+
+async def test_factories_request_scope(client: TestClient[litestar.Litestar], app: litestar.Litestar) -> None:
+    @litestar.websocket_listener("/ws")
+    async def websocket_handler(data: str, di_container: Container) -> None:
+        assert data == "test"
+        with di_container.build_child_container() as request_container:
+            request_factory_instance = request_factory.sync_resolve(request_container)
+            assert isinstance(request_factory_instance, DependentCreator)
+
+    app.register(websocket_handler)
+
+    with client.websocket_connect("/ws") as websocket:
+        websocket.send("test")
 
 
-async def test_websocket_not_supported(client: TestClient[litestar.Litestar], app: litestar.Litestar) -> None:
-    async def websocket_handler(data: str) -> None:
-        pass
+async def test_context_adapter(client: TestClient[litestar.Litestar], app: litestar.Litestar) -> None:
+    @litestar.websocket_listener("/ws", dependencies={"path": FromDI(context_adapter)})
+    async def websocket_handler(data: str, path: str) -> None:
+        assert data == "test"
+        assert path == "/ws"
 
-    app.register(litestar.websocket_listener("/ws")(websocket_handler))
+    app.register(websocket_handler)
 
     with client.websocket_connect("/ws") as websocket:
         websocket.send("test")

packages/modern-di/modern_di/providers/abstract.py

@@ -58,8 +58,8 @@ def __init__(
         super().__init__(scope)
         self._check_providers_scope(itertools.chain(args, kwargs.values()))
         self._creator: typing.Final = creator
-        self._args: typing.Final = args
-        self._kwargs: typing.Final = kwargs
+        self._args: typing.Final[P.args] = args
+        self._kwargs: typing.Final[P.kwargs] = kwargs
 
     def _sync_resolve_args(self, container: Container) -> list[typing.Any]:
         return [x.sync_resolve(container) if isinstance(x, AbstractProvider) else x for x in self._args]