Add subdevices support to ESPHome (home-assistant#147343)

Author: bdraco

homeassistant/components/esphome/entity.py

@@ -4,6 +4,7 @@
 
 from collections.abc import Awaitable, Callable, Coroutine
 import functools
+import logging
 import math
 from typing import TYPE_CHECKING, Any, Concatenate, Generic, TypeVar, cast
 
@@ -13,7 +14,6 @@
     EntityCategory as EsphomeEntityCategory,
     EntityInfo,
     EntityState,
-    build_unique_id,
 )
 import voluptuous as vol
 
@@ -24,6 +24,7 @@
     config_validation as cv,
     device_registry as dr,
     entity_platform,
+    entity_registry as er,
 )
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity import Entity
@@ -32,9 +33,11 @@
 from .const import DOMAIN
 
 # Import config flow so that it's added to the registry
-from .entry_data import ESPHomeConfigEntry, RuntimeEntryData
+from .entry_data import ESPHomeConfigEntry, RuntimeEntryData, build_device_unique_id
 from .enum_mapper import EsphomeEnumMapper
 
+_LOGGER = logging.getLogger(__name__)
+
 _InfoT = TypeVar("_InfoT", bound=EntityInfo)
 _EntityT = TypeVar("_EntityT", bound="EsphomeEntity[Any,Any]")
 _StateT = TypeVar("_StateT", bound=EntityState)
@@ -53,21 +56,74 @@ def async_static_info_updated(
 ) -> None:
     """Update entities of this platform when entities are listed."""
     current_infos = entry_data.info[info_type]
+    device_info = entry_data.device_info
+    if TYPE_CHECKING:
+        assert device_info is not None
     new_infos: dict[int, EntityInfo] = {}
     add_entities: list[_EntityT] = []
 
+    ent_reg = er.async_get(hass)
+    dev_reg = dr.async_get(hass)
+
     for info in infos:
-        if not current_infos.pop(info.key, None):
-            # Create new entity
+        new_infos[info.key] = info
+
+        # Create new entity if it doesn't exist
+        if not (old_info := current_infos.pop(info.key, None)):
             entity = entity_type(entry_data, platform.domain, info, state_type)
             add_entities.append(entity)
-        new_infos[info.key] = info
+            continue
+
+        # Entity exists - check if device_id has changed
+        if old_info.device_id == info.device_id:
+            continue
+
+        # Entity has switched devices, need to migrate unique_id
+        old_unique_id = build_device_unique_id(device_info.mac_address, old_info)
+        entity_id = ent_reg.async_get_entity_id(platform.domain, DOMAIN, old_unique_id)
+
+        # If entity not found in registry, re-add it
+        # This happens when the device_id changed and the old device was deleted
+        if entity_id is None:
+            _LOGGER.info(
+                "Entity with old unique_id %s not found in registry after device_id "
+                "changed from %s to %s, re-adding entity",
+                old_unique_id,
+                old_info.device_id,
+                info.device_id,
+            )
+            entity = entity_type(entry_data, platform.domain, info, state_type)
+            add_entities.append(entity)
+            continue
+
+        updates: dict[str, Any] = {}
+        new_unique_id = build_device_unique_id(device_info.mac_address, info)
+
+        # Update unique_id if it changed
+        if old_unique_id != new_unique_id:
+            updates["new_unique_id"] = new_unique_id
+
+        # Update device assignment
+        if info.device_id:
+            # Entity now belongs to a sub device
+            new_device = dev_reg.async_get_device(
+                identifiers={(DOMAIN, f"{device_info.mac_address}_{info.device_id}")}
+            )
+        else:
+            # Entity now belongs to the main device
+            new_device = dev_reg.async_get_device(
+                connections={(dr.CONNECTION_NETWORK_MAC, device_info.mac_address)}
+            )
+
+        if new_device:
+            updates["device_id"] = new_device.id
+
+        # Apply all updates at once
+        if updates:
+            ent_reg.async_update_entity(entity_id, **updates)
 
     # Anything still in current_infos is now gone
     if current_infos:
-        device_info = entry_data.device_info
-        if TYPE_CHECKING:
-            assert device_info is not None
         entry_data.async_remove_entities(
             hass, current_infos.values(), device_info.mac_address
         )
@@ -244,19 +300,36 @@ def __init__(
         self._key = entity_info.key
         self._state_type = state_type
         self._on_static_info_update(entity_info)
-        self._attr_device_info = DeviceInfo(
-            connections={(dr.CONNECTION_NETWORK_MAC, device_info.mac_address)}
-        )
+
+        device_name = device_info.name
+        # Determine the device connection based on whether this entity belongs to a sub device
+        if entity_info.device_id:
+            # Entity belongs to a sub device
+            self._attr_device_info = DeviceInfo(
+                identifiers={
+                    (DOMAIN, f"{device_info.mac_address}_{entity_info.device_id}")
+                }
+            )
+            # Use the pre-computed device_id_to_name mapping for O(1) lookup
+            device_name = entry_data.device_id_to_name.get(
+                entity_info.device_id, device_info.name
+            )
+        else:
+            # Entity belongs to the main device
+            self._attr_device_info = DeviceInfo(
+                connections={(dr.CONNECTION_NETWORK_MAC, device_info.mac_address)}
+            )
+
         if entity_info.name:
-            self.entity_id = f"{domain}.{device_info.name}_{entity_info.object_id}"
+            self.entity_id = f"{domain}.{device_name}_{entity_info.object_id}"
         else:
             # https://github.com/home-assistant/core/issues/132532
             # If name is not set, ESPHome will use the sanitized friendly name
             # as the name, however we want to use the original object_id
             # as the entity_id before it is sanitized since the sanitizer
             # is not utf-8 aware. In this case, its always going to be
             # an empty string so we drop the object_id.
-            self.entity_id = f"{domain}.{device_info.name}"
+            self.entity_id = f"{domain}.{device_name}"
 
     async def async_added_to_hass(self) -> None:
         """Register callbacks."""
@@ -290,7 +363,9 @@ def _on_static_info_update(self, static_info: EntityInfo) -> None:
             static_info = cast(_InfoT, static_info)
             assert device_info
         self._static_info = static_info
-        self._attr_unique_id = build_unique_id(device_info.mac_address, static_info)
+        self._attr_unique_id = build_device_unique_id(
+            device_info.mac_address, static_info
+        )
         self._attr_entity_registry_enabled_default = not static_info.disabled_by_default
         # https://github.com/home-assistant/core/issues/132532
         # If the name is "", we need to set it to None since otherwise

homeassistant/components/esphome/entry_data.py

@@ -95,6 +95,22 @@
 }
 
 
+def build_device_unique_id(mac: str, entity_info: EntityInfo) -> str:
+    """Build unique ID for entity, appending @device_id if it belongs to a sub-device.
+
+    This wrapper around build_unique_id ensures that entities belonging to sub-devices
+    have their device_id appended to the unique_id to handle proper migration when
+    entities move between devices.
+    """
+    base_unique_id = build_unique_id(mac, entity_info)
+
+    # If entity belongs to a sub-device, append @device_id
+    if entity_info.device_id:
+        return f"{base_unique_id}@{entity_info.device_id}"
+
+    return base_unique_id
+
+
 class StoreData(TypedDict, total=False):
     """ESPHome storage data."""
 
@@ -160,6 +176,7 @@ class RuntimeEntryData:
     assist_satellite_set_wake_word_callbacks: list[Callable[[str], None]] = field(
         default_factory=list
     )
+    device_id_to_name: dict[int, str] = field(default_factory=dict)
 
     @property
     def name(self) -> str:
@@ -222,7 +239,9 @@ def async_remove_entities(
         ent_reg = er.async_get(hass)
         for info in static_infos:
             if entry := ent_reg.async_get_entity_id(
-                INFO_TYPE_TO_PLATFORM[type(info)], DOMAIN, build_unique_id(mac, info)
+                INFO_TYPE_TO_PLATFORM[type(info)],
+                DOMAIN,
+                build_device_unique_id(mac, info),
             ):
                 ent_reg.async_remove(entry)
 
@@ -278,7 +297,8 @@ async def async_update_static_infos(
             if (
                 (old_unique_id := info.unique_id)
                 and (old_entry := registry_get_entity(platform, DOMAIN, old_unique_id))
-                and (new_unique_id := build_unique_id(mac, info)) != old_unique_id
+                and (new_unique_id := build_device_unique_id(mac, info))
+                != old_unique_id
                 and not registry_get_entity(platform, DOMAIN, new_unique_id)
             ):
                 ent_reg.async_update_entity(old_entry, new_unique_id=new_unique_id)

homeassistant/components/esphome/manager.py

@@ -527,6 +527,11 @@ async def _on_connect(self) -> None:
                 device_info.name,
                 device_mac,
             )
+        # Build device_id_to_name mapping for efficient lookup
+        entry_data.device_id_to_name = {
+            sub_device.device_id: sub_device.name or device_info.name
+            for sub_device in device_info.devices
+        }
         self.device_id = _async_setup_device_registry(hass, entry, entry_data)
 
         entry_data.async_update_device_state()
@@ -751,6 +756,28 @@ def _async_setup_device_registry(
     device_info = entry_data.device_info
     if TYPE_CHECKING:
         assert device_info is not None
+
+    device_registry = dr.async_get(hass)
+    # Build sets of valid device identifiers and connections
+    valid_connections = {
+        (dr.CONNECTION_NETWORK_MAC, format_mac(device_info.mac_address))
+    }
+    valid_identifiers = {
+        (DOMAIN, f"{device_info.mac_address}_{sub_device.device_id}")
+        for sub_device in device_info.devices
+    }
+
+    # Remove devices that no longer exist
+    for device in dr.async_entries_for_config_entry(device_registry, entry.entry_id):
+        # Skip devices we want to keep
+        if (
+            device.connections & valid_connections
+            or device.identifiers & valid_identifiers
+        ):
+            continue
+        # Remove everything else
+        device_registry.async_remove_device(device.id)
+
     sw_version = device_info.esphome_version
     if device_info.compilation_time:
         sw_version += f" ({device_info.compilation_time})"
@@ -779,11 +806,14 @@ def _async_setup_device_registry(
             f"{device_info.project_version} (ESPHome {device_info.esphome_version})"
         )
 
-    suggested_area = None
-    if device_info.suggested_area:
+    suggested_area: str | None = None
+    if device_info.area and device_info.area.name:
+        # Prefer device_info.area over suggested_area when area name is not empty
+        suggested_area = device_info.area.name
+    elif device_info.suggested_area:
         suggested_area = device_info.suggested_area
 
-    device_registry = dr.async_get(hass)
+    # Create/update main device
     device_entry = device_registry.async_get_or_create(
         config_entry_id=entry.entry_id,
         configuration_url=configuration_url,
@@ -794,6 +824,36 @@ def _async_setup_device_registry(
         sw_version=sw_version,
         suggested_area=suggested_area,
     )
+
+    # Handle sub devices
+    # Find available areas from device_info
+    areas_by_id = {area.area_id: area for area in device_info.areas}
+    # Add the main device's area if it exists
+    if device_info.area:
+        areas_by_id[device_info.area.area_id] = device_info.area
+    # Create/update sub devices that should exist
+    for sub_device in device_info.devices:
+        # Determine the area for this sub device
+        sub_device_suggested_area: str | None = None
+        if sub_device.area_id is not None and sub_device.area_id in areas_by_id:
+            sub_device_suggested_area = areas_by_id[sub_device.area_id].name
+
+        sub_device_entry = device_registry.async_get_or_create(
+            config_entry_id=entry.entry_id,
+            identifiers={(DOMAIN, f"{device_info.mac_address}_{sub_device.device_id}")},
+            name=sub_device.name or device_entry.name,
+            manufacturer=manufacturer,
+            model=model,
+            sw_version=sw_version,
+            suggested_area=sub_device_suggested_area,
+        )
+
+        # Update the sub device to set via_device_id
+        device_registry.async_update_device(
+            sub_device_entry.id,
+            via_device_id=device_entry.id,
+        )
+
     return device_entry.id