From 0c4c95394ea8d061996dcd66b8f02f1aa90f6add Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Tue, 21 Feb 2023 20:40:39 +0100 Subject: [PATCH] Add WS command config/entity_registry/list_for_display (#87787) * Add WS command config/entity_registry/list_for_display * Make more keys in the display dict optional * Move disabled_by check to ws command handler * Hide hidden_by if not hidden * Use send_json_auto_id in the new test * Don't include entities which have no data needed for display * Include platform for entries with translation_key --- .../components/config/entity_registry.py | 40 ++++- homeassistant/helpers/entity_registry.py | 87 ++++++++- .../components/config/test_entity_registry.py | 165 +++++++++++++++++- 3 files changed, 279 insertions(+), 13 deletions(-) diff --git a/homeassistant/components/config/entity_registry.py b/homeassistant/components/config/entity_registry.py index 0f48c64ae4b..20e00ec11ed 100644 --- a/homeassistant/components/config/entity_registry.py +++ b/homeassistant/components/config/entity_registry.py @@ -15,16 +15,18 @@ from homeassistant.helpers import ( device_registry as dr, entity_registry as er, ) +from homeassistant.helpers.json import json_dumps async def async_setup(hass: HomeAssistant) -> bool: """Enable the Entity Registry views.""" - 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) - websocket_api.async_register_command(hass, websocket_update_entity) + websocket_api.async_register_command(hass, websocket_get_entity) + websocket_api.async_register_command(hass, websocket_list_entities_for_display) + websocket_api.async_register_command(hass, websocket_list_entities) websocket_api.async_register_command(hass, websocket_remove_entity) + websocket_api.async_register_command(hass, websocket_update_entity) return True @@ -40,7 +42,7 @@ def websocket_list_entities( # Build start of response message msg_json_prefix = ( f'{{"id":{msg["id"]},"type": "{websocket_api.const.TYPE_RESULT}",' - f'"success":true,"result": [' + '"success":true,"result": [' ) # Concatenate cached entity registry item JSON serializations msg_json = ( @@ -55,6 +57,36 @@ def websocket_list_entities( connection.send_message(msg_json) +@websocket_api.websocket_command( + {vol.Required("type"): "config/entity_registry/list_for_display"} +) +@callback +def websocket_list_entities_for_display( + 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 + entity_categories = json_dumps(er.ENTITY_CATEGORY_INDEX_TO_VALUE) + msg_json_prefix = ( + f'{{"id":{msg["id"]},"type":"{websocket_api.const.TYPE_RESULT}","success":true,' + f'"result":{{"entity_categories":{entity_categories},"entities":[' + ) + # Concatenate cached entity registry item JSON serializations + msg_json = ( + msg_json_prefix + + ",".join( + entry.display_json_repr + for entry in registry.entities.values() + if entry.disabled_by is None and entry.display_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 f779fe00ed1..78ccdcd12f5 100644 --- a/homeassistant/helpers/entity_registry.py +++ b/homeassistant/helpers/entity_registry.py @@ -65,6 +65,13 @@ STORAGE_VERSION_MAJOR = 1 STORAGE_VERSION_MINOR = 10 STORAGE_KEY = "core.entity_registry" +ENTITY_CATEGORY_VALUE_TO_INDEX: dict[EntityCategory | None, int] = { + # mypy does not understand strenum + val: idx # type: ignore[misc] + for idx, val in enumerate(EntityCategory) +} +ENTITY_CATEGORY_INDEX_TO_VALUE = dict(enumerate(EntityCategory)) + # Attributes relevant to describing entity # to external services. ENTITY_DESCRIBING_ATTRIBUTES = { @@ -97,6 +104,12 @@ class RegistryEntryHider(StrEnum): EntityOptionsType = Mapping[str, Mapping[str, Any]] +DISLAY_DICT_OPTIONAL = ( + ("ai", "area_id"), + ("di", "device_id"), + ("tk", "translation_key"), +) + @attr.s(slots=True, frozen=True) class RegistryEntry: @@ -131,7 +144,12 @@ 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) + _partial_repr: str | None | UndefinedType = attr.ib( + cmp=False, default=UNDEFINED, init=False, repr=False + ) + _display_repr: str | None | UndefinedType = attr.ib( + cmp=False, default=UNDEFINED, init=False, repr=False + ) @domain.default def _domain_default(self) -> str: @@ -148,6 +166,63 @@ class RegistryEntry: """Return if entry is hidden.""" return self.hidden_by is not None + @property + def _as_display_dict(self) -> dict[str, Any] | None: + """Return a partial dict representation of the entry. + + This version only includes what's needed for display. + Returns None if there's no data needed for display. + """ + display_dict: dict[str, Any] = {} + for key, attr_name in DISLAY_DICT_OPTIONAL: + if (attr_val := getattr(self, attr_name)) is not None: + display_dict[key] = attr_val + if "tk" in display_dict: + display_dict["pl"] = self.platform + if (category := self.entity_category) is not None: + display_dict["ec"] = ENTITY_CATEGORY_VALUE_TO_INDEX[category] + if self.hidden_by is not None: + display_dict["hb"] = True + if not self.name and self.has_entity_name: + display_dict["en"] = self.original_name + if self.domain == "sensor" and (sensor_options := self.options.get("sensor")): + if (precision := sensor_options.get("display_precision")) is not None: + display_dict["dp"] = precision + elif ( + precision := sensor_options.get("suggested_display_precision") + ) is not None: + display_dict["dp"] = precision + if not display_dict: + # We didn't gather any data needed for display + return None + display_dict["ei"] = self.entity_id + return display_dict + + @property + def display_json_repr(self) -> str | None: + """Return a cached partial JSON representation of the entry. + + This version only includes what's needed for display. + """ + if self._display_repr is not UNDEFINED: + return self._display_repr + + try: + dict_repr = self._as_display_dict + json_repr: str | None = JSON_DUMP(dict_repr) if dict_repr else None + object.__setattr__(self, "_display_repr", json_repr) + except (ValueError, TypeError): + object.__setattr__(self, "_display_repr", None) + _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) + ), + ) + # Mypy doesn't understand the __setattr__ business + return self._display_repr # type: ignore[return-value] + @property def as_partial_dict(self) -> dict[str, Any]: """Return a partial dict representation of the entry.""" @@ -173,13 +248,14 @@ class RegistryEntry: @property def partial_json_repr(self) -> str | None: """Return a cached partial JSON representation of the entry.""" - if self._json_repr is not None: - return self._json_repr + if self._partial_repr is not UNDEFINED: + return self._partial_repr try: dict_repr = self.as_partial_dict - object.__setattr__(self, "_json_repr", JSON_DUMP(dict_repr)) + object.__setattr__(self, "_partial_repr", JSON_DUMP(dict_repr)) except (ValueError, TypeError): + object.__setattr__(self, "_partial_repr", None) _LOGGER.error( "Unable to serialize entry %s to JSON. Bad data found at %s", self.entity_id, @@ -187,7 +263,8 @@ class RegistryEntry: find_paths_unserializable_data(dict_repr, dump=JSON_DUMP) ), ) - return self._json_repr + # Mypy doesn't understand the __setattr__ business + return self._partial_repr # type: ignore[return-value] @callback def write_unavailable_state(self, hass: HomeAssistant) -> None: diff --git a/tests/components/config/test_entity_registry.py b/tests/components/config/test_entity_registry.py index 8873c85dcf1..66f1567ddbe 100644 --- a/tests/components/config/test_entity_registry.py +++ b/tests/components/config/test_entity_registry.py @@ -3,7 +3,7 @@ import pytest from pytest_unordered import unordered from homeassistant.components.config import entity_registry -from homeassistant.const import ATTR_ICON +from homeassistant.const import ATTR_ICON, EntityCategory from homeassistant.core import HomeAssistant from homeassistant.helpers import device_registry as dr from homeassistant.helpers.device_registry import DeviceEntryDisabler @@ -22,13 +22,16 @@ from tests.common import ( mock_device_registry, mock_registry, ) +from tests.typing import MockHAClientWebSocket, WebSocketGenerator @pytest.fixture -def client(hass, hass_ws_client): +async def client( + hass: HomeAssistant, hass_ws_client: WebSocketGenerator +) -> MockHAClientWebSocket: """Fixture that can interact with the config manager API.""" - hass.loop.run_until_complete(entity_registry.async_setup(hass)) - return hass.loop.run_until_complete(hass_ws_client(hass)) + await entity_registry.async_setup(hass) + return await hass_ws_client(hass) @pytest.fixture @@ -144,6 +147,160 @@ async def test_list_entities(hass: HomeAssistant, client) -> None: ] +async def test_list_entities_for_display( + hass: HomeAssistant, client: MockHAClientWebSocket +) -> None: + """Test list entries.""" + mock_registry( + hass, + { + "test_domain.test": RegistryEntry( + area_id="area52", + device_id="device123", + entity_category=EntityCategory.DIAGNOSTIC, + entity_id="test_domain.test", + has_entity_name=True, + original_name="Hello World", + platform="test_platform", + translation_key="translations_galore", + unique_id="1234", + ), + "test_domain.nameless": RegistryEntry( + area_id="area52", + device_id="device123", + entity_id="test_domain.nameless", + has_entity_name=True, + original_name=None, + platform="test_platform", + unique_id="2345", + ), + "test_domain.renamed": RegistryEntry( + area_id="area52", + device_id="device123", + entity_id="test_domain.renamed", + has_entity_name=True, + name="User name", + original_name="Hello World", + platform="test_platform", + unique_id="3456", + ), + "test_domain.boring": RegistryEntry( + entity_id="test_domain.boring", + platform="test_platform", + unique_id="4567", + ), + "test_domain.disabled": RegistryEntry( + disabled_by=RegistryEntryDisabler.USER, + entity_id="test_domain.disabled", + hidden_by=RegistryEntryHider.USER, + platform="test_platform", + unique_id="789A", + ), + "test_domain.hidden": RegistryEntry( + entity_id="test_domain.hidden", + hidden_by=RegistryEntryHider.USER, + platform="test_platform", + unique_id="89AB", + ), + "sensor.default_precision": RegistryEntry( + entity_id="sensor.default_precision", + options={"sensor": {"suggested_display_precision": 0}}, + platform="test_platform", + unique_id="9ABC", + ), + "sensor.user_precision": RegistryEntry( + entity_id="sensor.user_precision", + options={ + "sensor": {"display_precision": 0, "suggested_display_precision": 1} + }, + platform="test_platform", + unique_id="ABCD", + ), + }, + ) + + await client.send_json_auto_id({"type": "config/entity_registry/list_for_display"}) + msg = await client.receive_json() + + assert msg["result"] == { + "entity_categories": {"0": "config", "1": "diagnostic"}, + "entities": [ + { + "ai": "area52", + "di": "device123", + "ec": 1, + "ei": "test_domain.test", + "en": "Hello World", + "pl": "test_platform", + "tk": "translations_galore", + }, + { + "ai": "area52", + "di": "device123", + "ei": "test_domain.nameless", + "en": None, + }, + { + "ai": "area52", + "di": "device123", + "ei": "test_domain.renamed", + }, + { + "ei": "test_domain.hidden", + "hb": True, + }, + { + "dp": 0, + "ei": "sensor.default_precision", + }, + { + "dp": 0, + "ei": "sensor.user_precision", + }, + ], + } + + class Unserializable: + """Good luck serializing me.""" + + mock_registry( + hass, + { + "test_domain.test": RegistryEntry( + area_id="area52", + device_id="device123", + entity_id="test_domain.test", + has_entity_name=True, + original_name="Hello World", + platform="test_platform", + unique_id="1234", + ), + "test_domain.name_2": RegistryEntry( + entity_id="test_domain.name_2", + has_entity_name=True, + original_name=Unserializable(), + platform="test_platform", + unique_id="6789", + ), + }, + ) + + await client.send_json_auto_id({"type": "config/entity_registry/list_for_display"}) + msg = await client.receive_json() + + assert msg["result"] == { + "entity_categories": {"0": "config", "1": "diagnostic"}, + "entities": [ + { + "ai": "area52", + "di": "device123", + "ei": "test_domain.test", + "en": "Hello World", + }, + ], + } + + async def test_get_entity(hass: HomeAssistant, client) -> None: """Test get entry.""" mock_registry(