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/docs/conf.py b/docs/conf.py
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -103,6 +103,7 @@
     ("py:class", "litestar.utils.signature.ParsedParameter"),
     ("py:class", "litestar.utils.sync.AsyncCallable"),
     ("py:class", "BacklogStrategy"),
+    ("py:class", "ExceptionT"),
 ]
 nitpick_ignore_regex = [
     (r"py:.*", r"litestar\.types.*"),
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
@@ -1,11 +1,14 @@
 from __future__ import annotations
 
 import inspect
-from typing import TYPE_CHECKING, Any, Callable, Coroutine, cast
+from functools import wraps
+from typing import TYPE_CHECKING, Any, Callable, Coroutine, Dict, cast
 
+from litestar.di import Provide
 from litestar.dto.interface import ConnectionContext
 from litestar.serialization import decode_json
 from litestar.utils import AsyncCallable
+from litestar.utils.helpers import unwrap_partial
 
 if TYPE_CHECKING:
     from contextlib import AbstractAsyncContextManager
@@ -112,11 +115,13 @@ async def handle_send(socket: WebSocket, data_to_send: Any, dto: DTOInterface |
 
 def create_handler_function(
     listener_context: ListenerContext,
-    lifespan_manager: Callable[[WebSocket], AbstractAsyncContextManager],
+    lifespan_manager: Callable[..., AbstractAsyncContextManager],
 ) -> Callable[..., Coroutine[None, None, None]]:
     listener_callback = AsyncCallable(listener_context.listener_callback)
 
-    async def handler_fn(*args: Any, socket: WebSocket, **kwargs: Any) -> None:
+    async def handler_fn(
+        *args: Any, socket: WebSocket, connection_lifespan_dependencies: Dict[str, Any], **kwargs: Any  # noqa: UP006
+    ) -> 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
@@ -126,7 +131,7 @@ async def handler_fn(*args: Any, socket: WebSocket, **kwargs: Any) -> None:
         if listener_context.pass_socket:
             kwargs["socket"] = socket
 
-        async with lifespan_manager(socket):
+        async with lifespan_manager(**connection_lifespan_dependencies):
             while True:
                 received_data = await handle_receive(socket, data_dto)
                 data_to_send = await listener_callback(*args, data=received_data, **kwargs)
@@ -155,4 +160,24 @@ def create_handler_signature(callback_signature: inspect.Signature) -> inspect.S
     new_params = [p for p in callback_signature.parameters.values() if p.name not in {"data"}]
     if "socket" not in callback_signature.parameters:
         new_params.append(inspect.Parameter(name="socket", kind=inspect.Parameter.KEYWORD_ONLY, annotation="WebSocket"))
+
+    new_params.append(
+        inspect.Parameter(
+            name="connection_lifespan_dependencies", kind=inspect.Parameter.KEYWORD_ONLY, annotation="Dict[str, Any]"
+        )
+    )
+
     return callback_signature.replace(parameters=new_params)
+
+
+def create_stub_dependency(src: AnyCallable) -> Provide:
+    """Create a stub dependency, accepting any kwargs defined in ``src``, and
+    wrap it in ``Provide``
+    """
+    src = unwrap_partial(src)
+
+    @wraps(src)
+    async def stub(**kwargs: Any) -> Dict[str, Any]:  # noqa: UP006
+        return kwargs
+
+    return Provide(stub)
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,6 +1,5 @@
 from __future__ import annotations
 
-import inspect
 from abc import ABC, abstractmethod
 from contextlib import AbstractAsyncContextManager, asynccontextmanager
 from functools import partial
@@ -9,8 +8,11 @@
     Any,
     AsyncGenerator,
     Callable,
+    Dict,
     Mapping,
+    Optional,
     cast,
+    overload,
 )
 
 from msgspec.json import Encoder as JsonEncoder
@@ -28,7 +30,6 @@
     ExceptionHandler,
     Guard,
     Middleware,
-    SyncOrAsyncUnion,
     TypeEncodersMap,
 )
 from litestar.types.builtin_types import NoneType
@@ -64,14 +65,61 @@ class websocket_listener(WebsocketRouteHandler):
         "_send_mode": None,
         "_listener_context": None,
         "_connection_lifespan": None,
+        "_dependency_stubs": None,
     }
 
+    @overload
+    def __init__(
+        self,
+        path: str | None | list[str] | None = None,
+        *,
+        connection_lifespan: Callable[..., AbstractAsyncContextManager[Any]] | None = None,
+        dependencies: Dependencies | None = None,
+        dto: type[DTOInterface] | None | EmptyType = Empty,
+        exception_handlers: dict[int | type[Exception], ExceptionHandler] | None = None,
+        guards: list[Guard] | None = None,
+        middleware: list[Middleware] | None = None,
+        receive_mode: WebSocketMode = "text",
+        send_mode: WebSocketMode = "text",
+        name: str | None = None,
+        opt: dict[str, Any] | None = None,
+        return_dto: type[DTOInterface] | None | EmptyType = Empty,
+        signature_namespace: Mapping[str, Any] | None = None,
+        type_encoders: TypeEncodersMap | None = None,
+        **kwargs: Any,
+    ) -> None:
+        ...
+
+    @overload
+    def __init__(
+        self,
+        path: str | None | list[str] | None = None,
+        *,
+        connection_accept_handler: Callable[[WebSocket], Coroutine[Any, Any, None]] = WebSocket.accept,
+        dependencies: Dependencies | None = None,
+        dto: type[DTOInterface] | None | EmptyType = Empty,
+        exception_handlers: dict[int | type[Exception], ExceptionHandler] | None = None,
+        guards: list[Guard] | None = None,
+        middleware: list[Middleware] | None = None,
+        receive_mode: WebSocketMode = "text",
+        send_mode: WebSocketMode = "text",
+        name: str | None = None,
+        on_accept: AnyCallable | None = None,
+        on_disconnect: AnyCallable | None = None,
+        opt: dict[str, Any] | None = None,
+        return_dto: type[DTOInterface] | None | EmptyType = Empty,
+        signature_namespace: Mapping[str, Any] | None = None,
+        type_encoders: TypeEncodersMap | None = None,
+        **kwargs: Any,
+    ) -> 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,
+        connection_lifespan: Callable[..., AbstractAsyncContextManager[Any]] | None = None,
         dependencies: Dependencies | None = None,
         dto: type[DTOInterface] | None | EmptyType = Empty,
         exception_handlers: dict[int | type[Exception], ExceptionHandler] | None = None,
@@ -80,8 +128,8 @@ def __init__(
         receive_mode: WebSocketMode = "text",
         send_mode: WebSocketMode = "text",
         name: str | None = None,
-        on_accept: Callable[[WebSocket], SyncOrAsyncUnion[None]] | None = None,
-        on_disconnect: Callable[[WebSocket], SyncOrAsyncUnion[None]] | None = None,
+        on_accept: AnyCallable | None = None,
+        on_disconnect: AnyCallable | None = None,
         opt: dict[str, Any] | None = None,
         return_dto: type[DTOInterface] | None | EmptyType = Empty,
         signature_namespace: Mapping[str, Any] | None = None,
@@ -96,7 +144,8 @@ def __init__(
             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``.
+                it calls the ``connection_accept_handler``, ``on_connect`` and ``on_disconnect``. Can request any
+                dependencies, for example the :class:`WebSocket <.connection.WebSocket>` connection
             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.
@@ -106,10 +155,10 @@ def __init__(
             receive_mode: Websocket mode to receive data in, either `text` or `binary`.
             send_mode: Websocket mode to receive data in, either `text` or `binary`.
             name: A string identifying the route handler.
-            on_accept: Callback invoked after a connection has been accepted, receiving the
-                :class:`WebSocket <.connection.WebSocket>` instance as its only argument
-            on_disconnect: Callback invoked after a connection has been closed, receiving the
-                :class:`WebSocket <.connection.WebSocket>` instance as its only argument
+            on_accept: Callback invoked after a connection has been accepted. Can request any dependencies, for example
+                the :class:`WebSocket <.connection.WebSocket>` connection
+            on_disconnect: Callback invoked after a connection has been closed. Can request any dependencies, for
+                example the :class:`WebSocket <.connection.WebSocket>` connection
             opt: A string keyed mapping of arbitrary values that can be accessed in :class:`Guards <.types.Guard>` or
                 wherever you have access to :class:`Request <.connection.Request>` or
                 :class:`ASGI Scope <.types.Scope>`.
@@ -120,6 +169,12 @@ def __init__(
             type_encoders: A mapping of types to callables that transform them into types supported for serialization.
             **kwargs: Any additional kwarg - will be set in the opt dictionary.
         """
+        if connection_lifespan and any([on_accept, on_disconnect, connection_accept_handler is not WebSocket.accept]):
+            raise ImproperlyConfiguredException(
+                "connection_lifespan can not be used with connection hooks "
+                "(on_accept, on_disconnect, connection_accept_handler)",
+            )
+
         self._listener_context = _utils.ListenerContext()
         self._receive_mode: WebSocketMode = receive_mode
         self._send_mode: WebSocketMode = send_mode
@@ -146,9 +201,32 @@ def __init__(
         self.dto = dto
         self.return_dto = return_dto
 
+        if not self.dependencies:
+            self.dependencies = {}
+        self.dependencies = dict(self.dependencies)
+        self.dependencies["connection_lifespan_dependencies"] = _utils.create_stub_dependency(
+            self._connection_lifespan or self.default_connection_lifespan
+        )
+        if self.on_accept:
+            self.dependencies["on_accept_dependencies"] = _utils.create_stub_dependency(self.on_accept.ref.value)
+        if self.on_disconnect:
+            self.dependencies["on_disconnect_dependencies"] = _utils.create_stub_dependency(
+                self.on_disconnect.ref.value
+            )
+
     @asynccontextmanager
-    async def default_connection_lifespan(self, socket: WebSocket) -> AsyncGenerator[None, None]:
-        """Handle the connection lifespan of a WebSocket.
+    async def default_connection_lifespan(
+        self,
+        socket: WebSocket,
+        on_accept_dependencies: Optional[Dict[str, Any]] = None,  # noqa: UP007, UP006
+        on_disconnect_dependencies: Optional[Dict[str, Any]] = None,  # noqa: UP007, UP006
+    ) -> AsyncGenerator[None, None]:
+        """Handle the connection lifespan of a :class:`WebSocket <.connection.WebSocket>`.
+
+        Args:
+            socket: The :class:`WebSocket <.connection.WebSocket>` connection
+            on_accept_dependencies: Dependencies requested by the :attr:`on_accept` hook
+            on_disconnect_dependencies: Dependencies requested by the :attr:`on_disconnect` hook
 
         By, default this will
 
@@ -159,14 +237,14 @@ async def default_connection_lifespan(self, socket: WebSocket) -> AsyncGenerator
         await self.connection_accept_handler(socket)
 
         if self.on_accept:
-            await self.on_accept(socket)
+            await self.on_accept(**(on_accept_dependencies or {}))
         try:
             yield
         except WebSocketDisconnect:
             pass
         finally:
             if self.on_disconnect:
-                await self.on_disconnect(socket)
+                await self.on_disconnect(**(on_disconnect_dependencies or {}))
 
     def _validate_handler_function(self) -> None:
         """Validate the route handler function once it's set by inspecting its return annotations."""
@@ -198,12 +276,6 @@ 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
             )
@@ -260,8 +332,12 @@ class WebsocketListener(ABC):
     """A sequence of :class:`Guard <.types.Guard>` callables."""
     middleware: list[Middleware] | None = None
     """A sequence of :class:`Middleware <.types.Middleware>`."""
+    on_accept: AnyCallable | None = None
+    """Called after a :class:`WebSocket <.connection.WebSocket>` connection has been accepted. Can receive any dependencies"""
+    on_disconnect: AnyCallable | None = None
+    """Called after a :class:`WebSocket <.connection.WebSocket>` connection has been disconnected. Can receive any dependencies"""
     receive_mode: WebSocketMode = "text"
-    """Websocket mode to receive data in, either `text` or `binary`."""
+    """:class:`WebSocket <.connection.WebSocket>` mode to receive data in, either ``text`` or ``binary``."""
     send_mode: WebSocketMode = "text"
     """Websocket mode to send data in, either `text` or `binary`."""
     name: str | None = None
@@ -301,9 +377,6 @@ def __init__(self) -> None:
             type_encoders=self.type_encoders,
         )(self.on_receive)
 
-    def on_accept(self, socket: WebSocket) -> SyncOrAsyncUnion[None]:  # noqa: B027
-        """Called after a WebSocket connection has been accepted"""
-
     @abstractmethod
     def on_receive(self, *args: Any, **kwargs: Any) -> Any:
         """Called after data has been received from the WebSocket.
@@ -316,6 +389,3 @@ def on_receive(self, *args: Any, **kwargs: Any) -> Any:
         according to handler configuration.
         """
         raise NotImplementedError
-
-    def on_disconnect(self, socket: WebSocket) -> SyncOrAsyncUnion[None]:  # noqa: B027
-        """Called after a WebSocket connection has been disconnected"""


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,7 +1,7 @@
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from typing import AsyncGenerator, Dict, List, Optional, Type, Union, cast
-from unittest.mock import MagicMock
+from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 from pytest_lazyfixture import lazy_fixture
@@ -279,8 +279,14 @@ def handler(data: bytes) -> None:
 
 
 def test_connection_callbacks() -> None:
-    on_accept = MagicMock()
-    on_disconnect = MagicMock()
+    on_accept_mock = MagicMock()
+    on_disconnect_mock = MagicMock()
+
+    def on_accept(socket: WebSocket) -> None:
+        on_accept_mock()
+
+    def on_disconnect(socket: WebSocket) -> None:
+        on_disconnect_mock()
 
     @websocket_listener("/", on_accept=on_accept, on_disconnect=on_disconnect)
     def handler(data: bytes) -> None:
@@ -290,8 +296,8 @@ def handler(data: bytes) -> None:
     with client.websocket_connect("/"):
         pass
 
-    on_accept.assert_called_once()
-    on_disconnect.assert_called_once()
+    on_accept_mock.assert_called_once()
+    on_disconnect_mock.assert_called_once()
 
 
 def test_connection_lifespan() -> None:
@@ -328,7 +334,67 @@ class ClientController(Controller):
         async def websocket_handler(self, data: str, socket: WebSocket) -> str:
             return data
 
-    with create_test_client(ClientController, debug=True) as client, client.websocket_connect("/ws") as ws:
+    with create_test_client(ClientController) as client, client.websocket_connect("/ws") as ws:
         ws.send_text("foo")
         data = ws.receive_text(timeout=1)
         assert data == "foo"
+
+
+def test_lifespan_dependencies() -> None:
+    mock = MagicMock()
+
+    @asynccontextmanager
+    async def lifespan(name: str, state: State, query: dict) -> AsyncGenerator[None, None]:
+        mock(name=name, state=state, query=query)
+        yield
+
+    @websocket_listener("/{name:str}", connection_lifespan=lifespan)
+    async def handler(data: str) -> None:
+        pass
+
+    with create_test_client([handler], debug=True) as client, client.websocket_connect("/foo") as ws:
+        ws.send_text("")
+
+    assert mock.call_args_list[0].kwargs["name"] == "foo"
+    assert isinstance(mock.call_args_list[0].kwargs["state"], State)
+    assert isinstance(mock.call_args_list[0].kwargs["query"], dict)
+
+
+def test_hook_dependencies() -> None:
+    on_accept_mock = MagicMock()
+    on_disconnect_mock = MagicMock()
+
+    def on_accept(name: str, state: State, query: dict) -> None:
+        on_accept_mock(name=name, state=state, query=query)
+
+    def on_disconnect(name: str, state: State, query: dict) -> None:
+        on_disconnect_mock(name=name, state=state, query=query)
+
+    @websocket_listener("/{name: str}", on_accept=on_accept, on_disconnect=on_disconnect)
+    def handler(data: bytes) -> None:
+        pass
+
+    with create_test_client([handler], debug=True) as client, client.websocket_connect("/foo") as ws:
+        ws.send_text("")
+
+    assert on_accept_mock.call_args_list[0].kwargs["name"] == "foo"
+    assert isinstance(on_accept_mock.call_args_list[0].kwargs["state"], State)
+    assert isinstance(on_accept_mock.call_args_list[0].kwargs["query"], dict)
+    assert on_disconnect_mock.call_args_list[0].kwargs["name"] == "foo"
+    assert isinstance(on_disconnect_mock.call_args_list[0].kwargs["state"], State)
+    assert isinstance(on_disconnect_mock.call_args_list[0].kwargs["query"], dict)
+
+
+@pytest.mark.parametrize("hook_name", ["on_accept", "on_disconnect", "connection_accept_handler"])
+def test_listeners_lifespan_hooks_and_manager_raises(hook_name: str) -> None:
+    @asynccontextmanager
+    async def lifespan() -> AsyncGenerator[None, None]:
+        yield
+
+    hook_callback = AsyncMock()
+
+    with pytest.raises(ImproperlyConfiguredException):
+
+        @websocket_listener("/", **{hook_name: hook_callback}, connection_lifespan=lifespan)
+        def handler(data: bytes) -> None:
+            pass