From a0c023d5cb6bcc24bf6d9021d99657ffbe5aea12 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Wed, 14 Jun 2023 08:47:18 -1000
Subject: [PATCH] Reduce overhead to lookup items in the entity and device
 registry (#94568)

---
 homeassistant/helpers/device_registry.py | 10 ++++++++--
 homeassistant/helpers/entity_registry.py | 10 ++++++++--
 tests/common.py                          |  2 ++
 tests/components/homekit/test_homekit.py | 19 +++++++++++++++++--
 4 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/homeassistant/helpers/device_registry.py b/homeassistant/helpers/device_registry.py
index 29e64639722..7df01fc8fd2 100644
--- a/homeassistant/helpers/device_registry.py
+++ b/homeassistant/helpers/device_registry.py
@@ -292,6 +292,7 @@ class DeviceRegistry:
 
     devices: DeviceRegistryItems[DeviceEntry]
     deleted_devices: DeviceRegistryItems[DeletedDeviceEntry]
+    _device_data: dict[str, DeviceEntry]
 
     def __init__(self, hass: HomeAssistant) -> None:
         """Initialize the device registry."""
@@ -306,8 +307,12 @@ class DeviceRegistry:
 
     @callback
     def async_get(self, device_id: str) -> DeviceEntry | None:
-        """Get device."""
-        return self.devices.get(device_id)
+        """Get device.
+
+        We retrieve the DeviceEntry from the underlying dict to avoid
+        the overhead of the UserDict __getitem__.
+        """
+        return self._device_data.get(device_id)
 
     @callback
     def async_get_device(
@@ -641,6 +646,7 @@ class DeviceRegistry:
 
         self.devices = devices
         self.deleted_devices = deleted_devices
+        self._device_data = devices.data
 
     @callback
     def async_schedule_save(self) -> None:
diff --git a/homeassistant/helpers/entity_registry.py b/homeassistant/helpers/entity_registry.py
index b6fc84b9627..15f05e2bd42 100644
--- a/homeassistant/helpers/entity_registry.py
+++ b/homeassistant/helpers/entity_registry.py
@@ -425,6 +425,7 @@ class EntityRegistry:
     """Class to hold a registry of entities."""
 
     entities: EntityRegistryItems
+    _entities_data: dict[str, RegistryEntry]
 
     def __init__(self, hass: HomeAssistant) -> None:
         """Initialize the registry."""
@@ -470,8 +471,12 @@ class EntityRegistry:
 
     @callback
     def async_get(self, entity_id: str) -> RegistryEntry | None:
-        """Get EntityEntry for an entity_id."""
-        return self.entities.get(entity_id)
+        """Get EntityEntry for an entity_id.
+
+        We retrieve the RegistryEntry from the underlying dict to avoid
+        the overhead of the UserDict __getitem__.
+        """
+        return self._entities_data.get(entity_id)
 
     @callback
     def async_get_entity_id(
@@ -991,6 +996,7 @@ class EntityRegistry:
                 )
 
         self.entities = entities
+        self._entities_data = entities.data
 
     @callback
     def async_schedule_save(self) -> None:
diff --git a/tests/common.py b/tests/common.py
index ca164fcaaf8..38b7cf79b75 100644
--- a/tests/common.py
+++ b/tests/common.py
@@ -509,6 +509,7 @@ def mock_registry(
     if mock_entries is None:
         mock_entries = {}
     registry.entities = er.EntityRegistryItems()
+    registry._entities_data = registry.entities.data
     for key, entry in mock_entries.items():
         registry.entities[key] = entry
 
@@ -554,6 +555,7 @@ def mock_device_registry(
     """
     registry = dr.DeviceRegistry(hass)
     registry.devices = dr.DeviceRegistryItems()
+    registry._device_data = registry.devices.data
     if mock_entries is None:
         mock_entries = {}
     for key, entry in mock_entries.items():
diff --git a/tests/components/homekit/test_homekit.py b/tests/components/homekit/test_homekit.py
index 0b74763c6a7..690bb7fef37 100644
--- a/tests/components/homekit/test_homekit.py
+++ b/tests/components/homekit/test_homekit.py
@@ -1644,7 +1644,6 @@ async def test_homekit_ignored_missing_devices(
     light = entity_registry.async_get_or_create(
         "light", "powerwall", "demo", device_id=device_entry.id
     )
-    before_removal = entity_registry.entities.copy()
     # Delete the device to make sure we fallback
     # to using the platform
     device_registry.async_remove_device(device_entry.id)
@@ -1652,7 +1651,23 @@ async def test_homekit_ignored_missing_devices(
     await asyncio.sleep(0)
     await asyncio.sleep(0)
     # Restore the registry
-    entity_registry.entities = before_removal
+    entity_registry.async_get_or_create(
+        "binary_sensor",
+        "powerwall",
+        "battery_charging",
+        device_id=device_entry.id,
+        original_device_class=BinarySensorDeviceClass.BATTERY_CHARGING,
+    )
+    entity_registry.async_get_or_create(
+        "sensor",
+        "powerwall",
+        "battery",
+        device_id=device_entry.id,
+        original_device_class=SensorDeviceClass.BATTERY,
+    )
+    light = entity_registry.async_get_or_create(
+        "light", "powerwall", "demo", device_id=device_entry.id
+    )
 
     hass.states.async_set(light.entity_id, STATE_ON)
     hass.states.async_set("light.two", STATE_ON)