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