Refactor frontend user store (#144723)

* Refactor frontend user store

* Address review comments
This commit is contained in:
Erik Montnemery 2025-05-12 12:00:32 +02:00 committed by GitHub
parent 63e38b4d8d
commit cba12fb598
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 55 additions and 55 deletions

View File

@ -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}
)

View File

@ -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: