Speed up registry indices (#117897)

* Use defaultdict for registry indices

defaultdict is faster and does not have to create an empty
dict that gets throw away when the key is already present

* Use defaultdict for registry indices

defaultdict is faster and does not have to create an empty
dict that gets throw away when the key is already present

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(

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

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

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.