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
This commit is contained in:
J. Nick Koston 2024-05-22 08:07:39 -10:00 committed by GitHub
parent f99ec87338
commit 6113b58e9c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 28 additions and 23 deletions

View File

@ -2,6 +2,7 @@
from __future__ import annotations from __future__ import annotations
from collections import defaultdict
from collections.abc import Iterable from collections.abc import Iterable
import dataclasses import dataclasses
from functools import cached_property from functools import cached_property
@ -19,7 +20,7 @@ from .normalized_name_base_registry import (
NormalizedNameBaseRegistryItems, NormalizedNameBaseRegistryItems,
normalize_name, normalize_name,
) )
from .registry import BaseRegistry from .registry import BaseRegistry, RegistryIndexType
from .singleton import singleton from .singleton import singleton
from .storage import Store from .storage import Store
from .typing import UNDEFINED, UndefinedType from .typing import UNDEFINED, UndefinedType
@ -135,15 +136,15 @@ class AreaRegistryItems(NormalizedNameBaseRegistryItems[AreaEntry]):
def __init__(self) -> None: def __init__(self) -> None:
"""Initialize the area registry items.""" """Initialize the area registry items."""
super().__init__() super().__init__()
self._labels_index: dict[str, dict[str, Literal[True]]] = {} self._labels_index: RegistryIndexType = defaultdict(dict)
self._floors_index: dict[str, dict[str, Literal[True]]] = {} self._floors_index: RegistryIndexType = defaultdict(dict)
def _index_entry(self, key: str, entry: AreaEntry) -> None: def _index_entry(self, key: str, entry: AreaEntry) -> None:
"""Index an entry.""" """Index an entry."""
if entry.floor_id is not None: 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: for label in entry.labels:
self._labels_index.setdefault(label, {})[key] = True self._labels_index[label][key] = True
super()._index_entry(key, entry) super()._index_entry(key, entry)
def _unindex_entry( def _unindex_entry(

View File

@ -2,6 +2,7 @@
from __future__ import annotations from __future__ import annotations
from collections import defaultdict
from collections.abc import Mapping from collections.abc import Mapping
from enum import StrEnum from enum import StrEnum
from functools import cached_property, lru_cache, partial from functools import cached_property, lru_cache, partial
@ -37,7 +38,7 @@ from .deprecation import (
) )
from .frame import report from .frame import report
from .json import JSON_DUMP, find_paths_unserializable_data, json_bytes, json_fragment 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 .singleton import singleton
from .typing import UNDEFINED, UndefinedType from .typing import UNDEFINED, UndefinedType
@ -513,19 +514,19 @@ class ActiveDeviceRegistryItems(DeviceRegistryItems[DeviceEntry]):
- label -> dict[key, True] - label -> dict[key, True]
""" """
super().__init__() super().__init__()
self._area_id_index: dict[str, dict[str, Literal[True]]] = {} self._area_id_index: RegistryIndexType = defaultdict(dict)
self._config_entry_id_index: dict[str, dict[str, Literal[True]]] = {} self._config_entry_id_index: RegistryIndexType = defaultdict(dict)
self._labels_index: dict[str, dict[str, Literal[True]]] = {} self._labels_index: RegistryIndexType = defaultdict(dict)
def _index_entry(self, key: str, entry: DeviceEntry) -> None: def _index_entry(self, key: str, entry: DeviceEntry) -> None:
"""Index an entry.""" """Index an entry."""
super()._index_entry(key, entry) super()._index_entry(key, entry)
if (area_id := entry.area_id) is not None: 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: 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: 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( def _unindex_entry(
self, key: str, replacement_entry: DeviceEntry | None = None self, key: str, replacement_entry: DeviceEntry | None = None

View File

@ -10,6 +10,7 @@ timer.
from __future__ import annotations from __future__ import annotations
from collections import defaultdict
from collections.abc import Callable, Container, Hashable, KeysView, Mapping from collections.abc import Callable, Container, Hashable, KeysView, Mapping
from datetime import datetime, timedelta from datetime import datetime, timedelta
from enum import StrEnum from enum import StrEnum
@ -58,7 +59,7 @@ from .device_registry import (
EventDeviceRegistryUpdatedData, EventDeviceRegistryUpdatedData,
) )
from .json import JSON_DUMP, find_paths_unserializable_data, json_bytes, json_fragment 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 .singleton import singleton
from .typing import UNDEFINED, UndefinedType from .typing import UNDEFINED, UndefinedType
@ -533,10 +534,10 @@ class EntityRegistryItems(BaseRegistryItems[RegistryEntry]):
super().__init__() super().__init__()
self._entry_ids: dict[str, RegistryEntry] = {} self._entry_ids: dict[str, RegistryEntry] = {}
self._index: dict[tuple[str, str, str], str] = {} self._index: dict[tuple[str, str, str], str] = {}
self._config_entry_id_index: dict[str, dict[str, Literal[True]]] = {} self._config_entry_id_index: RegistryIndexType = defaultdict(dict)
self._device_id_index: dict[str, dict[str, Literal[True]]] = {} self._device_id_index: RegistryIndexType = defaultdict(dict)
self._area_id_index: dict[str, dict[str, Literal[True]]] = {} self._area_id_index: RegistryIndexType = defaultdict(dict)
self._labels_index: dict[str, dict[str, Literal[True]]] = {} self._labels_index: RegistryIndexType = defaultdict(dict)
def _index_entry(self, key: str, entry: RegistryEntry) -> None: def _index_entry(self, key: str, entry: RegistryEntry) -> None:
"""Index an entry.""" """Index an entry."""
@ -545,13 +546,13 @@ class EntityRegistryItems(BaseRegistryItems[RegistryEntry]):
# python has no ordered set, so we use a dict with True values # python has no ordered set, so we use a dict with True values
# https://discuss.python.org/t/add-orderedset-to-stdlib/12730 # https://discuss.python.org/t/add-orderedset-to-stdlib/12730
if (config_entry_id := entry.config_entry_id) is not None: 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: 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: 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: for label in entry.labels:
self._labels_index.setdefault(label, {})[key] = True self._labels_index[label][key] = True
def _unindex_entry( def _unindex_entry(
self, key: str, replacement_entry: RegistryEntry | None = None self, key: str, replacement_entry: RegistryEntry | None = None

View File

@ -3,7 +3,7 @@
from __future__ import annotations from __future__ import annotations
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from collections import UserDict from collections import UserDict, defaultdict
from collections.abc import Mapping, Sequence, ValuesView from collections.abc import Mapping, Sequence, ValuesView
from typing import TYPE_CHECKING, Any, Literal from typing import TYPE_CHECKING, Any, Literal
@ -15,6 +15,8 @@ if TYPE_CHECKING:
SAVE_DELAY = 10 SAVE_DELAY = 10
SAVE_DELAY_LONG = 180 SAVE_DELAY_LONG = 180
type RegistryIndexType = defaultdict[str, dict[str, Literal[True]]]
class BaseRegistryItems[_DataT](UserDict[str, _DataT], ABC): class BaseRegistryItems[_DataT](UserDict[str, _DataT], ABC):
"""Base class for registry items.""" """Base class for registry items."""
@ -42,7 +44,7 @@ class BaseRegistryItems[_DataT](UserDict[str, _DataT], ABC):
self._index_entry(key, entry) self._index_entry(key, entry)
def _unindex_entry_value( def _unindex_entry_value(
self, key: str, value: str, index: dict[str, dict[str, Literal[True]]] self, key: str, value: str, index: RegistryIndexType
) -> None: ) -> None:
"""Unindex an entry value. """Unindex an entry value.