From b933a53aa3aea4d97c26fdf384e93b234aae1a6a Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Mon, 9 Jan 2023 16:52:52 +0100 Subject: [PATCH] Refactor entity registry JSON cache (#85085) * Refactor entity registry JSON cache * Fix generator * Tweak * Improve string building * Improve test coverage * Override EntityRegistryItems.values to avoid __iter__ overhead --- .../components/config/entity_registry.py | 69 ++++++++----------- homeassistant/helpers/entity_registry.py | 48 ++++++++++++- .../components/config/test_entity_registry.py | 14 ++-- 3 files changed, 84 insertions(+), 47 deletions(-) diff --git a/homeassistant/components/config/entity_registry.py b/homeassistant/components/config/entity_registry.py index da3d8c7e2b1..90ec415703f 100644 --- a/homeassistant/components/config/entity_registry.py +++ b/homeassistant/components/config/entity_registry.py @@ -9,12 +9,7 @@ from homeassistant import config_entries from homeassistant.components import websocket_api from homeassistant.components.websocket_api import ERR_NOT_FOUND from homeassistant.components.websocket_api.decorators import require_admin -from homeassistant.components.websocket_api.messages import ( - IDEN_JSON_TEMPLATE, - IDEN_TEMPLATE, - message_to_json, -) -from homeassistant.core import Event, HomeAssistant, callback +from homeassistant.core import HomeAssistant, callback from homeassistant.helpers import ( config_validation as cv, device_registry as dr, @@ -25,41 +20,6 @@ from homeassistant.helpers import ( async def async_setup(hass: HomeAssistant) -> bool: """Enable the Entity Registry views.""" - cached_list_entities: str | None = None - - @callback - def _async_clear_list_entities_cache(event: Event) -> None: - nonlocal cached_list_entities - cached_list_entities = None - - @websocket_api.websocket_command( - {vol.Required("type"): "config/entity_registry/list"} - ) - @callback - def websocket_list_entities( - hass: HomeAssistant, - connection: websocket_api.ActiveConnection, - msg: dict[str, Any], - ) -> None: - """Handle list registry entries command.""" - nonlocal cached_list_entities - if not cached_list_entities: - registry = er.async_get(hass) - cached_list_entities = message_to_json( - websocket_api.result_message( - IDEN_TEMPLATE, # type: ignore[arg-type] - [_entry_dict(entry) for entry in registry.entities.values()], - ) - ) - connection.send_message( - cached_list_entities.replace(IDEN_JSON_TEMPLATE, str(msg["id"]), 1) - ) - - hass.bus.async_listen( - er.EVENT_ENTITY_REGISTRY_UPDATED, - _async_clear_list_entities_cache, - run_immediately=True, - ) websocket_api.async_register_command(hass, websocket_list_entities) websocket_api.async_register_command(hass, websocket_get_entity) websocket_api.async_register_command(hass, websocket_get_entities) @@ -68,6 +28,33 @@ async def async_setup(hass: HomeAssistant) -> bool: return True +@websocket_api.websocket_command({vol.Required("type"): "config/entity_registry/list"}) +@callback +def websocket_list_entities( + hass: HomeAssistant, + connection: websocket_api.ActiveConnection, + msg: dict[str, Any], +) -> None: + """Handle list registry entries command.""" + registry = er.async_get(hass) + # Build start of response message + msg_json_prefix = ( + f'{{"id":{msg["id"]},"type": "{websocket_api.const.TYPE_RESULT}",' + f'"success":true,"result": [' + ) + # Concatenate cached entity registry item JSON serializations + msg_json = ( + msg_json_prefix + + ",".join( + entry.json_repr + for entry in registry.entities.values() + if entry.json_repr is not None + ) + + "]}" + ) + connection.send_message(msg_json) + + @websocket_api.websocket_command( { vol.Required("type"): "config/entity_registry/get", diff --git a/homeassistant/helpers/entity_registry.py b/homeassistant/helpers/entity_registry.py index 54ed93aebb9..4222d38ecd4 100644 --- a/homeassistant/helpers/entity_registry.py +++ b/homeassistant/helpers/entity_registry.py @@ -10,7 +10,7 @@ timer. from __future__ import annotations from collections import UserDict -from collections.abc import Callable, Iterable, Mapping +from collections.abc import Callable, Iterable, Mapping, ValuesView import logging from typing import TYPE_CHECKING, Any, TypeVar, cast @@ -42,10 +42,15 @@ from homeassistant.core import ( from homeassistant.exceptions import MaxLengthExceeded from homeassistant.loader import bind_hass from homeassistant.util import slugify, uuid as uuid_util +from homeassistant.util.json import ( + find_paths_unserializable_data, + format_unserializable_data, +) from . import device_registry as dr, storage from .device_registry import EVENT_DEVICE_REGISTRY_UPDATED from .frame import report +from .json import JSON_DUMP from .typing import UNDEFINED, UndefinedType if TYPE_CHECKING: @@ -130,6 +135,8 @@ class RegistryEntry: translation_key: str | None = attr.ib(default=None) unit_of_measurement: str | None = attr.ib(default=None) + _json_repr: str | None = attr.ib(cmp=False, default=None, init=False, repr=False) + @domain.default def _domain_default(self) -> str: """Compute domain value.""" @@ -145,6 +152,41 @@ class RegistryEntry: """Return if entry is hidden.""" return self.hidden_by is not None + @property + def json_repr(self) -> str | None: + """Return a cached JSON representation of the entry.""" + if self._json_repr is not None: + return self._json_repr + + try: + dict_repr = { + "area_id": self.area_id, + "config_entry_id": self.config_entry_id, + "device_id": self.device_id, + "disabled_by": self.disabled_by, + "entity_category": self.entity_category, + "entity_id": self.entity_id, + "has_entity_name": self.has_entity_name, + "hidden_by": self.hidden_by, + "icon": self.icon, + "id": self.id, + "name": self.name, + "original_name": self.original_name, + "platform": self.platform, + "translation_key": self.translation_key, + "unique_id": self.unique_id, + } + object.__setattr__(self, "_json_repr", JSON_DUMP(dict_repr)) + except (ValueError, TypeError): + _LOGGER.error( + "Unable to serialize entry %s to JSON. Bad data found at %s", + self.entity_id, + format_unserializable_data( + find_paths_unserializable_data(dict_repr, dump=JSON_DUMP) + ), + ) + return self._json_repr + @callback def write_unavailable_state(self, hass: HomeAssistant) -> None: """Write the unavailable state to the state machine.""" @@ -268,6 +310,10 @@ class EntityRegistryItems(UserDict[str, "RegistryEntry"]): self._entry_ids: dict[str, RegistryEntry] = {} self._index: dict[tuple[str, str, str], str] = {} + def values(self) -> ValuesView[RegistryEntry]: + """Return the underlying values to avoid __iter__ overhead.""" + return self.data.values() + def __setitem__(self, key: str, entry: RegistryEntry) -> None: """Add an item.""" if key in self: diff --git a/tests/components/config/test_entity_registry.py b/tests/components/config/test_entity_registry.py index 84426a3d791..38984d74057 100644 --- a/tests/components/config/test_entity_registry.py +++ b/tests/components/config/test_entity_registry.py @@ -6,7 +6,6 @@ from homeassistant.components.config import entity_registry from homeassistant.const import ATTR_ICON from homeassistant.helpers.device_registry import DeviceEntryDisabler from homeassistant.helpers.entity_registry import ( - EVENT_ENTITY_REGISTRY_UPDATED, RegistryEntry, RegistryEntryDisabler, RegistryEntryHider, @@ -95,6 +94,9 @@ async def test_list_entities(hass, client): }, ] + class Unserializable: + """Good luck serializing me.""" + mock_registry( hass, { @@ -104,13 +106,15 @@ async def test_list_entities(hass, client): platform="test_platform", name="Hello World", ), + "test_domain.name_2": RegistryEntry( + entity_id="test_domain.name_2", + unique_id="6789", + platform="test_platform", + name=Unserializable(), + ), }, ) - hass.bus.async_fire( - EVENT_ENTITY_REGISTRY_UPDATED, - {"action": "create", "entity_id": "test_domain.no_name"}, - ) await client.send_json({"id": 6, "type": "config/entity_registry/list"}) msg = await client.receive_json()