1
mirror of https://github.com/home-assistant/core synced 2024-09-28 03:04:04 +02:00
ha-core/homeassistant/components/frontend/storage.py
Erik Montnemery e1338adf1a
Allow configuring country and language in core config (#81734)
* Allow configuring country and language in core config

* Add script for updating list of countries

* Use black for formatting

* Fix quoting

* Move country codes to a separate file

* Address review comments

* Add generated/countries.py

* Get default language from owner account

* Remove unused variable

* Add script to generate list of supported languages

* Add tests

* Fix stale docsring

* Use format_python_namespace

* Correct async_user_store

* Improve typing

* Fix with_store decorator

* Initialize language in core store migration

* Fix startup

* Tweak

* Apply suggestions from code review

Co-authored-by: Franck Nijhof <git@frenck.dev>

* Update storage.py

Co-authored-by: Franck Nijhof <git@frenck.dev>
2022-11-24 23:25:50 +01:00

116 lines
3.1 KiB
Python

"""API for persistent storage for the frontend."""
from __future__ import annotations
from collections.abc import Callable
from functools import wraps
from typing import Any
import voluptuous as vol
from homeassistant.components import websocket_api
from homeassistant.components.websocket_api.connection import ActiveConnection
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.storage import Store
DATA_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(
hass: HomeAssistant, user_id: str
) -> tuple[Store, dict[str, Any]]:
"""Access a user store."""
_initialize_frontend_storage(hass)
stores, data = hass.data[DATA_STORAGE]
if (store := stores.get(user_id)) is None:
store = stores[user_id] = Store(
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(orig_func: Callable) -> Callable:
"""Decorate function to provide data."""
@wraps(orig_func)
async def with_store_func(
hass: HomeAssistant, connection: ActiveConnection, msg: dict
) -> None:
"""Provide user specific data and store to function."""
user_id = connection.user.id
store, user_data = await async_user_store(hass, user_id)
await orig_func(hass, connection, msg, store, user_data)
return with_store_func
@websocket_api.websocket_command(
{
vol.Required("type"): "frontend/set_user_data",
vol.Required("key"): str,
vol.Required("value"): vol.Any(bool, str, int, float, dict, list, None),
}
)
@websocket_api.async_response
@with_store
async def websocket_set_user_data(
hass: HomeAssistant,
connection: ActiveConnection,
msg: dict[str, Any],
store: Store,
data: dict[str, Any],
) -> 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"]))
@websocket_api.websocket_command(
{vol.Required("type"): "frontend/get_user_data", vol.Optional("key"): str}
)
@websocket_api.async_response
@with_store
async def websocket_get_user_data(
hass: HomeAssistant,
connection: ActiveConnection,
msg: dict[str, Any],
store: Store,
data: dict[str, Any],
) -> 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}
)
)