1from __future__ import annotations
 2
 3from typing import TYPE_CHECKING
 4from weakref import WeakKeyDictionary
 5
 6from .exceptions import SessionNotAvailable
 7
 8if TYPE_CHECKING:
 9    from plain.http import Request
10
11    from .core import SessionStore
12
13_request_sessions: WeakKeyDictionary[Request, SessionStore] = WeakKeyDictionary()
14
15
16def set_request_session(request: Request, session: SessionStore) -> None:
17    """Store the session for this request."""
18    _request_sessions[request] = session
19
20
21def get_request_session(request: Request) -> SessionStore:
22    """
23    Get the session for this request.
24
25    Raises SessionNotAvailable if no session is set (typically because
26    SessionMiddleware hasn't been called yet or an error occurred before it could run).
27    """
28    try:
29        return _request_sessions[request]
30    except KeyError as e:
31        raise SessionNotAvailable(
32            "Session is not available for this request. "
33            "Ensure SessionMiddleware is installed and has been called."
34        ) from e