diff --git a/homeassistant/helpers/area_registry.py b/homeassistant/helpers/area_registry.py index 598eff0f70c..975750ebbdd 100644 --- a/homeassistant/helpers/area_registry.py +++ b/homeassistant/helpers/area_registry.py @@ -2,6 +2,7 @@ from __future__ import annotations +from collections import defaultdict from collections.abc import Iterable import dataclasses from functools import cached_property @@ -19,7 +20,7 @@ from .normalized_name_base_registry import ( NormalizedNameBaseRegistryItems, normalize_name, ) -from .registry import BaseRegistry +from .registry import BaseRegistry, RegistryIndexType from .singleton import singleton from .storage import Store from .typing import UNDEFINED, UndefinedType @@ -135,15 +136,15 @@ class AreaRegistryItems(NormalizedNameBaseRegistryItems[AreaEntry]): def __init__(self) -> None: """Initialize the area registry items.""" super().__init__() - self._labels_index: dict[str, dict[str, Literal[True]]] = {} - self._floors_index: dict[str, dict[str, Literal[True]]] = {} + self._labels_index: RegistryIndexType = defaultdict(dict) + self._floors_index: RegistryIndexType = defaultdict(dict) def _index_entry(self, key: str, entry: AreaEntry) -> None: """Index an entry.""" if entry.floor_id is not None: - self._floors_index.setdefault(entry.floor_id, {})[key] = True + self._floors_index[entry.floor_id][key] = True for label in entry.labels: - self._labels_index.setdefault(label, {})[key] = True + self._labels_index[label][key] = True super()._index_entry(key, entry) def _unindex_entry( diff --git a/homeassistant/helpers/device_registry.py b/homeassistant/helpers/device_registry.py index e39676146d6..75fcda18eac 100644 --- a/homeassistant/helpers/device_registry.py +++ b/homeassistant/helpers/device_registry.py @@ -2,6 +2,7 @@ from __future__ import annotations +from collections import defaultdict from collections.abc import Mapping from enum import StrEnum from functools import cached_property, lru_cache, partial @@ -37,7 +38,7 @@ from .deprecation import ( ) from .frame import report from .json import JSON_DUMP, find_paths_unserializable_data, json_bytes, json_fragment -from .registry import BaseRegistry, BaseRegistryItems +from .registry import BaseRegistry, BaseRegistryItems, RegistryIndexType from .singleton import singleton from .typing import UNDEFINED, UndefinedType @@ -513,19 +514,19 @@ class ActiveDeviceRegistryItems(DeviceRegistryItems[DeviceEntry]): - label -> dict[key, True] """ super().__init__() - self._area_id_index: dict[str, dict[str, Literal[True]]] = {} - self._config_entry_id_index: dict[str, dict[str, Literal[True]]] = {} - self._labels_index: dict[str, dict[str, Literal[True]]] = {} + self._area_id_index: RegistryIndexType = defaultdict(dict) + self._config_entry_id_index: RegistryIndexType = defaultdict(dict) + self._labels_index: RegistryIndexType = defaultdict(dict) def _index_entry(self, key: str, entry: DeviceEntry) -> None: """Index an entry.""" super()._index_entry(key, entry) if (area_id := entry.area_id) is not None: - self._area_id_index.setdefault(area_id, {})[key] = True + self._area_id_index[area_id][key] = True for label in entry.labels: - self._labels_index.setdefault(label, {})[key] = True + self._labels_index[label][key] = True for config_entry_id in entry.config_entries: - self._config_entry_id_index.setdefault(config_entry_id, {})[key] = True + self._config_entry_id_index[config_entry_id][key] = True def _unindex_entry( self, key: str, replacement_entry: DeviceEntry | None = None diff --git a/homeassistant/helpers/entity_registry.py b/homeassistant/helpers/entity_registry.py index 1c43c8e7ec9..ebca6f17d43 100644 --- a/homeassistant/helpers/entity_registry.py +++ b/homeassistant/helpers/entity_registry.py @@ -10,6 +10,7 @@ timer. from __future__ import annotations +from collections import defaultdict from collections.abc import Callable, Container, Hashable, KeysView, Mapping from datetime import datetime, timedelta from enum import StrEnum @@ -58,7 +59,7 @@ from .device_registry import ( EventDeviceRegistryUpdatedData, ) from .json import JSON_DUMP, find_paths_unserializable_data, json_bytes, json_fragment -from .registry import BaseRegistry, BaseRegistryItems +from .registry import BaseRegistry, BaseRegistryItems, RegistryIndexType from .singleton import singleton from .typing import UNDEFINED, UndefinedType @@ -533,10 +534,10 @@ class EntityRegistryItems(BaseRegistryItems[RegistryEntry]): super().__init__() self._entry_ids: dict[str, RegistryEntry] = {} self._index: dict[tuple[str, str, str], str] = {} - self._config_entry_id_index: dict[str, dict[str, Literal[True]]] = {} - self._device_id_index: dict[str, dict[str, Literal[True]]] = {} - self._area_id_index: dict[str, dict[str, Literal[True]]] = {} - self._labels_index: dict[str, dict[str, Literal[True]]] = {} + self._config_entry_id_index: RegistryIndexType = defaultdict(dict) + self._device_id_index: RegistryIndexType = defaultdict(dict) + self._area_id_index: RegistryIndexType = defaultdict(dict) + self._labels_index: RegistryIndexType = defaultdict(dict) def _index_entry(self, key: str, entry: RegistryEntry) -> None: """Index an entry.""" @@ -545,13 +546,13 @@ class EntityRegistryItems(BaseRegistryItems[RegistryEntry]): # python has no ordered set, so we use a dict with True values # https://discuss.python.org/t/add-orderedset-to-stdlib/12730 if (config_entry_id := entry.config_entry_id) is not None: - self._config_entry_id_index.setdefault(config_entry_id, {})[key] = True + self._config_entry_id_index[config_entry_id][key] = True if (device_id := entry.device_id) is not None: - self._device_id_index.setdefault(device_id, {})[key] = True + self._device_id_index[device_id][key] = True if (area_id := entry.area_id) is not None: - self._area_id_index.setdefault(area_id, {})[key] = True + self._area_id_index[area_id][key] = True for label in entry.labels: - self._labels_index.setdefault(label, {})[key] = True + self._labels_index[label][key] = True def _unindex_entry( self, key: str, replacement_entry: RegistryEntry | None = None diff --git a/homeassistant/helpers/registry.py b/homeassistant/helpers/registry.py index 9791b03c5cb..21f2178554e 100644 --- a/homeassistant/helpers/registry.py +++ b/homeassistant/helpers/registry.py @@ -3,7 +3,7 @@ from __future__ import annotations from abc import ABC, abstractmethod -from collections import UserDict +from collections import UserDict, defaultdict from collections.abc import Mapping, Sequence, ValuesView from typing import TYPE_CHECKING, Any, Literal @@ -15,6 +15,8 @@ if TYPE_CHECKING: SAVE_DELAY = 10 SAVE_DELAY_LONG = 180 +type RegistryIndexType = defaultdict[str, dict[str, Literal[True]]] + class BaseRegistryItems[_DataT](UserDict[str, _DataT], ABC): """Base class for registry items.""" @@ -42,7 +44,7 @@ class BaseRegistryItems[_DataT](UserDict[str, _DataT], ABC): self._index_entry(key, entry) def _unindex_entry_value( - self, key: str, value: str, index: dict[str, dict[str, Literal[True]]] + self, key: str, value: str, index: RegistryIndexType ) -> None: """Unindex an entry value.