Refactor frontend user store (#144723)

* Refactor frontend user store

* Address review comments

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(
+DATA_STORAGE: HassKey[dict[str, UserStore]] = HassKey("frontend_storage")
-    "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(
+async def async_user_store(hass: HomeAssistant, user_id: str) -> UserStore:
-    hass: HomeAssistant, user_id: str
-) -> tuple[Store, dict[str, Any]]:
     """Access a user store."""
-    _initialize_frontend_storage(hass)
+    stores = hass.data.setdefault(DATA_STORAGE, {})
-    stores, data = hass.data[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_user_store(
-
-
-def with_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,
+    store: UserStore,
-    data: dict[str, Any],
 ) -> None:
-    """Handle set global data command.
+    """Handle set user data command."""
-
+    await store.async_set_item(msg["key"], msg["value"])
-    Async friendly.
+    connection.send_result(msg["id"])
-    """
-    data[msg["key"]] = msg["value"]
-    await store.async_save(data)
-    connection.send_message(websocket_api.result_message(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,
+    store: UserStore,
-    data: dict[str, Any],
 ) -> None:
-    """Handle get global data command.
+    """Handle get user data command."""
-
+    data = store.data
-    Async friendly.
+    connection.send_result(
-    """
-    connection.send_message(
-        websocket_api.result_message(
         msg["id"], {"value": data.get(msg["key"]) if "key" in msg else data}
     )
-    )

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
+                            "language" in owner_store.data
-                            and "language" in owner_data["language"]
+                            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: