diff --git a/homeassistant/components/frontend/storage.py b/homeassistant/components/frontend/storage.py index a33a9de7ac5..c75fa360db7 100644 --- a/homeassistant/components/frontend/storage.py +++ b/homeassistant/components/frontend/storage.py @@ -10,53 +10,63 @@ import voluptuous as vol from homeassistant.components import websocket_api from homeassistant.components.websocket_api import ActiveConnection -from homeassistant.core import HomeAssistant, callback +from homeassistant.core import HomeAssistant from homeassistant.helpers.storage import Store from homeassistant.util.hass_dict import HassKey -DATA_STORAGE: HassKey[tuple[dict[str, Store], dict[str, dict]]] = HassKey( - "frontend_storage" -) +DATA_STORAGE: HassKey[dict[str, UserStore]] = HassKey("frontend_storage") STORAGE_VERSION_USER_DATA = 1 -@callback -def _initialize_frontend_storage(hass: HomeAssistant) -> None: - """Set up frontend storage.""" - if DATA_STORAGE in hass.data: - return - hass.data[DATA_STORAGE] = ({}, {}) - - async def async_setup_frontend_storage(hass: HomeAssistant) -> None: """Set up frontend storage.""" - _initialize_frontend_storage(hass) websocket_api.async_register_command(hass, websocket_set_user_data) websocket_api.async_register_command(hass, websocket_get_user_data) -async def async_user_store( - hass: HomeAssistant, user_id: str -) -> tuple[Store, dict[str, Any]]: +async def async_user_store(hass: HomeAssistant, user_id: str) -> UserStore: """Access a user store.""" - _initialize_frontend_storage(hass) - stores, data = hass.data[DATA_STORAGE] + stores = hass.data.setdefault(DATA_STORAGE, {}) if (store := stores.get(user_id)) is None: - store = stores[user_id] = Store( + store = stores[user_id] = UserStore(hass, user_id) + await store.async_load() + + return store + + +class UserStore: + """User store for frontend data.""" + + def __init__(self, hass: HomeAssistant, user_id: str) -> None: + """Initialize the user store.""" + self._store = _UserStore(hass, user_id) + self.data: dict[str, Any] = {} + + async def async_load(self) -> None: + """Load the data from the store.""" + self.data = await self._store.async_load() or {} + + async def async_set_item(self, key: str, value: Any) -> None: + """Set an item item and save the store.""" + self.data[key] = value + await self._store.async_save(self.data) + + +class _UserStore(Store[dict[str, Any]]): + """User store for frontend data.""" + + def __init__(self, hass: HomeAssistant, user_id: str) -> None: + """Initialize the user store.""" + super().__init__( hass, STORAGE_VERSION_USER_DATA, f"frontend.user_data_{user_id}", ) - if user_id not in data: - data[user_id] = await store.async_load() or {} - return store, data[user_id] - - -def with_store( +def with_user_store( orig_func: Callable[ - [HomeAssistant, ActiveConnection, dict[str, Any], Store, dict[str, Any]], + [HomeAssistant, ActiveConnection, dict[str, Any], UserStore], Coroutine[Any, Any, None], ], ) -> Callable[ @@ -65,17 +75,17 @@ def with_store( """Decorate function to provide data.""" @wraps(orig_func) - async def with_store_func( + async def with_user_store_func( hass: HomeAssistant, connection: ActiveConnection, msg: dict[str, Any] ) -> None: """Provide user specific data and store to function.""" user_id = connection.user.id - store, user_data = await async_user_store(hass, user_id) + store = await async_user_store(hass, user_id) - await orig_func(hass, connection, msg, store, user_data) + await orig_func(hass, connection, msg, store) - return with_store_func + return with_user_store_func @websocket_api.websocket_command( @@ -86,41 +96,31 @@ def with_store( } ) @websocket_api.async_response -@with_store +@with_user_store async def websocket_set_user_data( hass: HomeAssistant, connection: ActiveConnection, msg: dict[str, Any], - store: Store, - data: dict[str, Any], + store: UserStore, ) -> None: - """Handle set global data command. - - Async friendly. - """ - data[msg["key"]] = msg["value"] - await store.async_save(data) - connection.send_message(websocket_api.result_message(msg["id"])) + """Handle set user data command.""" + await store.async_set_item(msg["key"], msg["value"]) + connection.send_result(msg["id"]) @websocket_api.websocket_command( {vol.Required("type"): "frontend/get_user_data", vol.Optional("key"): str} ) @websocket_api.async_response -@with_store +@with_user_store async def websocket_get_user_data( hass: HomeAssistant, connection: ActiveConnection, msg: dict[str, Any], - store: Store, - data: dict[str, Any], + store: UserStore, ) -> None: - """Handle get global data command. - - Async friendly. - """ - connection.send_message( - websocket_api.result_message( - msg["id"], {"value": data.get(msg["key"]) if "key" in msg else data} - ) + """Handle get user data command.""" + data = store.data + connection.send_result( + msg["id"], {"value": data.get(msg["key"]) if "key" in msg else data} ) diff --git a/homeassistant/core_config.py b/homeassistant/core_config.py index ccae24a907b..f1ba96daae4 100644 --- a/homeassistant/core_config.py +++ b/homeassistant/core_config.py @@ -866,17 +866,17 @@ class Config: # pylint: disable-next=import-outside-toplevel from .components.frontend import storage as frontend_store - _, owner_data = await frontend_store.async_user_store( + owner_store = await frontend_store.async_user_store( self.hass, owner.id ) if ( - "language" in owner_data - and "language" in owner_data["language"] + "language" in owner_store.data + and "language" in owner_store.data["language"] ): with suppress(vol.InInvalid): data["language"] = cv.language( - owner_data["language"]["language"] + owner_store.data["language"]["language"] ) # pylint: disable-next=broad-except except Exception: