Repository: litestar-org/litestar

Issue:
StaticFilesConfig and virtual directories
I'm trying to write a ``FileSystemProtocol`` to load files from the package data using [importlib_resources](https://importlib-resources.readthedocs.io/en/latest/using.html#). But because ``directories`` is defined as ``DirectoryPath``, pydantic checks if the given directories exist in the local filesystem. 

This is not generally true, especially in any kind of virtual filesystem (e.g. a zipped package). I think this condition should be relaxed to support virtual filesystems.

https://github.com/starlite-api/starlite/blob/9bb6dcd57c10a591377cf8e3a537e9292566d5b9/starlite/config/static_files.py#L32


Discussion:
I agree. Do you want to submit a PR?

Resolving patch:
diff --git a/litestar/handlers/websocket_handlers/_utils.py b/litestar/handlers/websocket_handlers/_utils.py
--- a/litestar/handlers/websocket_handlers/_utils.py
+++ b/litestar/handlers/websocket_handlers/_utils.py
@@ -4,11 +4,12 @@
 from typing import TYPE_CHECKING, Any, Callable, Coroutine, cast
 
 from litestar.dto.interface import ConnectionContext
-from litestar.exceptions import WebSocketDisconnect
 from litestar.serialization import decode_json
 from litestar.utils import AsyncCallable
 
 if TYPE_CHECKING:
+    from contextlib import AbstractAsyncContextManager
+
     from msgspec.json import Encoder as JsonEncoder
 
     from litestar import WebSocket
@@ -111,34 +112,26 @@ async def handle_send(socket: WebSocket, data_to_send: Any, dto: DTOInterface |
 
 def create_handler_function(
     listener_context: ListenerContext,
-    on_accept: AsyncCallable | None,
-    on_disconnect: AsyncCallable | None,
-    accept_connection_handler: Callable[[WebSocket], Coroutine[Any, Any, None]],
+    lifespan_manager: Callable[[WebSocket], AbstractAsyncContextManager],
 ) -> Callable[..., Coroutine[None, None, None]]:
-    async def handler_fn(socket: WebSocket, **kwargs: Any) -> None:
-        await accept_connection_handler(socket)
+    listener_callback = AsyncCallable(listener_context.listener_callback)
 
-        listener_callback = AsyncCallable(listener_context.listener_callback)
+    async def handler_fn(socket: WebSocket, **kwargs: Any) -> None:
         ctx = ConnectionContext.from_connection(socket)
         data_dto = listener_context.resolved_data_dto(ctx) if listener_context.resolved_data_dto else None
         return_dto = listener_context.resolved_return_dto(ctx) if listener_context.resolved_return_dto else None
-
-        if on_accept:
-            await on_accept(socket)
+        handle_receive = listener_context.handle_receive
+        handle_send = listener_context.handle_send if listener_context.can_send_data else None
 
         if listener_context.pass_socket:
             kwargs["socket"] = socket
 
-        while True:
-            try:
-                received_data = await listener_context.handle_receive(socket, data_dto)
+        async with lifespan_manager(socket):
+            while True:
+                received_data = await handle_receive(socket, data_dto)
                 data_to_send = await listener_callback(data=received_data, **kwargs)
-                if listener_context.can_send_data:
-                    await listener_context.handle_send(socket, data_to_send, return_dto)
-            except WebSocketDisconnect:
-                if on_disconnect:
-                    await on_disconnect(socket)
-                break
+                if handle_send:
+                    await handle_send(socket, data_to_send, return_dto)
 
     return handler_fn
 
diff --git a/litestar/handlers/websocket_handlers/listener.py b/litestar/handlers/websocket_handlers/listener.py
--- a/litestar/handlers/websocket_handlers/listener.py
+++ b/litestar/handlers/websocket_handlers/listener.py
@@ -1,8 +1,17 @@
 from __future__ import annotations
 
+import inspect
 from abc import ABC, abstractmethod
+from contextlib import AbstractAsyncContextManager, asynccontextmanager
 from functools import partial
-from typing import TYPE_CHECKING, Any, Callable, Mapping, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncGenerator,
+    Callable,
+    Mapping,
+    cast,
+)
 
 from msgspec.json import Encoder as JsonEncoder
 
@@ -36,10 +45,8 @@
     from litestar.dto.interface import DTOInterface
     from litestar.types.asgi_types import WebSocketMode
 
-__all__ = (
-    "WebsocketListener",
-    "websocket_listener",
-)
+
+__all__ = ("WebsocketListener", "websocket_listener")
 
 
 class websocket_listener(WebsocketRouteHandler):
@@ -48,21 +55,23 @@ class websocket_listener(WebsocketRouteHandler):
     returned
     """
 
-    __slots__ = (
-        "_on_accept",
-        "_on_disconnect",
-        "_pass_socket",
-        "_receive_mode",
-        "_send_mode",
-        "_listener_context",
-        "accept_connection_handler",
-    )
+    __slots__ = {
+        "connection_accept_handler": "Callback to accept a WebSocket connection. By default, calls WebSocket.accept",
+        "on_accept": "Callback invoked after a WebSocket connection has been accepted",
+        "on_disconnect": "Callback invoked after a WebSocket connection has been closed",
+        "_pass_socket": None,
+        "_receive_mode": None,
+        "_send_mode": None,
+        "_listener_context": None,
+        "_connection_lifespan": None,
+    }
 
     def __init__(
         self,
         path: str | None | list[str] | None = None,
         *,
         connection_accept_handler: Callable[[WebSocket], Coroutine[Any, Any, None]] = WebSocket.accept,
+        connection_lifespan: Callable[[WebSocket], AbstractAsyncContextManager[Any]] | None = None,
         dependencies: Dependencies | None = None,
         dto: type[DTOInterface] | None | EmptyType = Empty,
         exception_handlers: dict[int | type[Exception], ExceptionHandler] | None = None,
@@ -86,6 +95,8 @@ def __init__(
                 to ``/``
             connection_accept_handler: A callable that accepts a :class:`WebSocket <.connection.WebSocket>` instance
                 and returns a coroutine that when awaited, will accept the connection. Defaults to ``WebSocket.accept``.
+            connection_lifespan: An asynchronous context manager, handling the lifespan of the connection. By default,
+                it calls the ``connection_accept_handler``, ``on_connect`` and ``on_disconnect``.
             dependencies: A string keyed mapping of dependency :class:`Provider <.di.Provide>` instances.
             dto: :class:`DTOInterface <.dto.interface.DTOInterface>` to use for (de)serializing and
                 validation of request data.
@@ -112,9 +123,11 @@ def __init__(
         self._listener_context = _utils.ListenerContext()
         self._receive_mode: WebSocketMode = receive_mode
         self._send_mode: WebSocketMode = send_mode
-        self._on_accept = AsyncCallable(on_accept) if on_accept else None
-        self._on_disconnect = AsyncCallable(on_disconnect) if on_disconnect else None
-        self.accept_connection_handler = connection_accept_handler
+        self._connection_lifespan = connection_lifespan
+
+        self.connection_accept_handler = connection_accept_handler
+        self.on_accept = AsyncCallable(on_accept) if on_accept else None
+        self.on_disconnect = AsyncCallable(on_disconnect) if on_disconnect else None
         self.type_encoders = type_encoders
 
         super().__init__(
@@ -133,6 +146,26 @@ def __init__(
         self.dto = dto
         self.return_dto = return_dto
 
+    @asynccontextmanager
+    async def default_connection_lifespan(self, socket: WebSocket) -> AsyncGenerator[None, None]:
+        """Handle the connection lifespan of a WebSocket.
+
+        By, default this will
+
+            - Call :attr:`connection_accept_handler` to accept a connection
+            - Call :attr:`on_accept` if defined after a connection has been accepted
+            - Call :attr:`on_disconnect` upon leaving the context
+        """
+        await self.connection_accept_handler(socket)
+
+        if self.on_accept:
+            await self.on_accept(socket)
+        try:
+            yield
+        finally:
+            if self.on_disconnect:
+                await self.on_disconnect(socket)
+
     def _validate_handler_function(self) -> None:
         """Validate the route handler function once it's set by inspecting its return annotations."""
         # since none of the validation rules of WebsocketRouteHandler apply here, this is let empty. Validation of the
@@ -152,9 +185,7 @@ def __call__(self, listener_callback: AnyCallable) -> websocket_listener:
         self._listener_context.listener_callback = listener_callback
         self._listener_context.handler_function = handler_function = _utils.create_handler_function(
             listener_context=self._listener_context,
-            on_accept=self._on_accept,
-            on_disconnect=self._on_disconnect,
-            accept_connection_handler=self.accept_connection_handler,
+            lifespan_manager=self._connection_lifespan or self.default_connection_lifespan,
         )
         return super().__call__(handler_function)
 
@@ -165,6 +196,12 @@ def on_registration(self, app: Litestar) -> None:
     def _create_signature_model(self, app: Litestar) -> None:
         """Create signature model for handler function."""
         if not self.signature_model:
+            extra_signatures = []
+            if self.on_accept:
+                extra_signatures.append(inspect.signature(self.on_accept))
+            if self.on_disconnect:
+                extra_signatures.append(inspect.signature(self.on_disconnect))
+
             new_signature = _utils.create_handler_signature(
                 self._listener_context.listener_callback_signature.original_signature
             )


Test patch:
diff --git a/tests/handlers/websocket/test_listeners.py b/tests/handlers/websocket/test_listeners.py
--- a/tests/handlers/websocket/test_listeners.py
+++ b/tests/handlers/websocket/test_listeners.py
@@ -1,5 +1,6 @@
+from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
-from typing import Dict, List, Optional, Type, Union, cast
+from typing import AsyncGenerator, Dict, List, Optional, Type, Union, cast
 from unittest.mock import MagicMock
 
 import pytest
@@ -275,3 +276,43 @@ def handler(data: bytes) -> None:
     client = create_test_client([handler])
     with client.websocket_connect("/") as ws:
         assert ws.extra_headers == [(b"cookie", b"custom-cookie")]
+
+
+def test_connection_callbacks() -> None:
+    on_accept = MagicMock()
+    on_disconnect = MagicMock()
+
+    @websocket_listener("/", on_accept=on_accept, on_disconnect=on_disconnect)
+    def handler(data: bytes) -> None:
+        pass
+
+    client = create_test_client([handler])
+    with client.websocket_connect("/"):
+        pass
+
+    on_accept.assert_called_once()
+    on_disconnect.assert_called_once()
+
+
+def test_connection_lifespan() -> None:
+    on_accept = MagicMock()
+    on_disconnect = MagicMock()
+
+    @asynccontextmanager
+    async def lifespan(socket: WebSocket) -> AsyncGenerator[None, None]:
+        on_accept(socket)
+        try:
+            yield
+        finally:
+            on_disconnect(socket)
+
+    @websocket_listener("/", connection_lifespan=lifespan)
+    def handler(data: bytes) -> None:
+        pass
+
+    client = create_test_client([handler])
+    with client.websocket_connect("/", timeout=1):
+        pass
+
+    on_accept.assert_called_once()
+    on_disconnect.assert_called_once()