Rework device binding code to be more consise (#2058)

custom_components/powercalc/device_binding.py

@@ -0,0 +1,78 @@
+import logging
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_DEVICE
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import device_registry
+from homeassistant.helpers.device_registry import DeviceInfo
+from homeassistant.helpers.entity import Entity
+from homeassistant.helpers.typing import ConfigType
+
+from custom_components.powercalc.common import SourceEntity
+from custom_components.powercalc.sensors.abstract import BaseEntity
+
+_LOGGER = logging.getLogger(__name__)
+
+
+async def attach_entities_to_source_device(
+    config_entry: ConfigEntry | None,
+    entities_to_add: list[Entity],
+    hass: HomeAssistant,
+    source_entity: SourceEntity,
+) -> None:
+    """Set the entity to same device as the source entity, if any available."""
+    if not source_entity.entity_entry or not source_entity.device_entry:
+        return
+
+    for entity in (
+        entity for entity in entities_to_add if isinstance(entity, BaseEntity)
+    ):
+        try:
+            entity.source_device_id = source_entity.device_entry.id  # type: ignore
+        except AttributeError:  # pragma: no cover
+            _LOGGER.error("%s: Cannot set device id on entity", entity.entity_id)
+
+    if config_entry:
+        bind_config_entry_to_device(hass, config_entry)
+
+
+def bind_config_entry_to_device(hass: HomeAssistant, config_entry: ConfigEntry) -> None:
+    """
+    When the user selected a specific device in the config flow, bind the config entry to that device
+    This will let HA bind all the powercalc entities for that config entry to the concerning device
+    """
+    device_id = config_entry.data.get(CONF_DEVICE)
+    if not device_id:
+        return
+
+    device_reg = device_registry.async_get(hass)
+    device_entry = device_reg.async_get(device_id)
+    if device_entry and config_entry.entry_id not in device_entry.config_entries:
+        device_reg.async_update_device(
+            device_id,
+            add_config_entry_id=config_entry.entry_id,
+        )
+
+
+def get_device_info(hass: HomeAssistant, sensor_config: ConfigType, source_entity: SourceEntity) -> DeviceInfo | None:
+    """
+    Get device info for a given powercalc entity configuration.
+    Prefer user configured device, when it is not set fallback to the same device as the source entity
+    """
+    device_id = sensor_config.get(CONF_DEVICE)
+    if device_id is not None:
+        device_reg = device_registry.async_get(hass)
+        device = device_reg.async_get(device_id)
+    else:
+        device = source_entity.device_entry
+
+    if device is None:
+        return None
+
+    if not device.identifiers and not device.connections:
+        return None
+
+    return DeviceInfo(
+        identifiers=device.identifiers,
+        connections=device.connections,
+    )

custom_components/powercalc/sensor.py

@@ -10,7 +10,6 @@
 from typing import Any
 
 import homeassistant.helpers.config_validation as cv
-import homeassistant.helpers.device_registry as dr
 import homeassistant.helpers.entity_registry as er
 import voluptuous as vol
 from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
@@ -20,15 +19,14 @@
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     CONF_CONDITION,
-    CONF_DEVICE,
     CONF_DOMAIN,
     CONF_ENTITIES,
     CONF_ENTITY_ID,
     CONF_NAME,
     CONF_UNIQUE_ID,
 )
 from homeassistant.core import Event, HomeAssistant, callback
-from homeassistant.helpers import device_registry, entity_platform
+from homeassistant.helpers import entity_platform
 from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.entity_registry import (
@@ -123,13 +121,13 @@
     SensorType,
     UnitPrefix,
 )
+from .device_binding import attach_entities_to_source_device, bind_config_entry_to_device
 from .errors import (
     PowercalcSetupError,
     SensorAlreadyConfiguredError,
     SensorConfigurationError,
 )
 from .group_include.include import resolve_include_entities
-from .sensors.abstract import BaseEntity
 from .sensors.daily_energy import (
     DAILY_FIXED_ENERGY_SCHEMA,
     create_daily_fixed_energy_power_sensor,
@@ -808,33 +806,6 @@ async def create_individual_sensors(
     return EntitiesBucket(new=entities_to_add, existing=[])
 
 
-async def attach_entities_to_source_device(
-    config_entry: ConfigEntry | None,
-    entities_to_add: list[Entity],
-    hass: HomeAssistant,
-    source_entity: SourceEntity,
-) -> None:
-    """Set the entity to same device as the source entity, if any available."""
-    if source_entity.entity_entry and source_entity.device_entry:
-        device_id = source_entity.device_entry.id
-        device_registry = dr.async_get(hass)
-        for entity in (
-            entity for entity in entities_to_add if isinstance(entity, BaseEntity)
-        ):
-            try:
-                entity.source_device_id = source_entity.device_entry.id  # type: ignore
-            except AttributeError:  # pragma: no cover
-                _LOGGER.error("%s: Cannot set device id on entity", entity.entity_id)
-        if (
-            config_entry
-            and config_entry.entry_id not in source_entity.device_entry.config_entries
-        ):
-            device_registry.async_update_device(
-                device_id,
-                add_config_entry_id=config_entry.entry_id,
-            )
-
-
 async def check_entity_not_already_configured(
     sensor_config: dict,
     source_entity: SourceEntity,
@@ -859,24 +830,6 @@ async def check_entity_not_already_configured(
         raise SensorAlreadyConfiguredError(source_entity.entity_id, existing_entities)
 
 
-def bind_config_entry_to_device(hass: HomeAssistant, config_entry: ConfigEntry) -> None:
-    """
-    When the user selected a specific device in the config flow, bind the config entry to that device
-    This will let HA bind all the powercalc entities for that config entry to the concerning device
-    """
-    device_id = config_entry.data.get(CONF_DEVICE)
-    if not device_id:
-        return
-
-    device_reg = device_registry.async_get(hass)
-    device_entry = device_reg.async_get(device_id)
-    if device_entry and config_entry.entry_id not in device_entry.config_entries:
-        device_reg.async_update_device(
-            device_id,
-            add_config_entry_id=config_entry.entry_id,
-        )
-
-
 @dataclass
 class EntitiesBucket:
     new: list[Entity] = field(default_factory=list)

custom_components/powercalc/sensors/energy.py

@@ -9,14 +9,12 @@
 from homeassistant.components.sensor import SensorDeviceClass, SensorStateClass
 from homeassistant.const import (
     ATTR_UNIT_OF_MEASUREMENT,
-    CONF_DEVICE,
     CONF_NAME,
     UnitOfEnergy,
     UnitOfPower,
     UnitOfTime,
 )
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers import device_registry
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.typing import ConfigType
@@ -36,6 +34,7 @@
     DEFAULT_ENERGY_INTEGRATION_METHOD,
     UnitPrefix,
 )
+from custom_components.powercalc.device_binding import get_device_info
 from custom_components.powercalc.errors import SensorConfigurationError
 
 from .abstract import (
@@ -138,7 +137,7 @@ async def create_energy_sensor(
         powercalc_source_entity=source_entity.entity_id,
         powercalc_source_domain=source_entity.domain,
         sensor_config=sensor_config,
-        device_info=get_device_info(hass, sensor_config, power_sensor),
+        device_info=get_device_info(hass, sensor_config, source_entity),
     )
 
 
@@ -164,22 +163,6 @@ def get_unit_prefix(
     return unit_prefix
 
 
-def get_device_info(hass: HomeAssistant, sensor_config: ConfigType, power_sensor: PowerSensor) -> DeviceInfo | None:
-    device_id = sensor_config.get(CONF_DEVICE)
-    if device_id is None:
-        return power_sensor.device_info
-
-    device_reg = device_registry.async_get(hass)
-    device = device_reg.async_get(device_id)
-    if device is None:
-        return None
-
-    return DeviceInfo(
-        identifiers=device.identifiers,
-        connections=device.connections,
-    )
-
-
 @callback
 def find_related_real_energy_sensor(
     hass: HomeAssistant,

custom_components/powercalc/sensors/group.py

@@ -386,6 +386,7 @@ async def get_entries_having_subgroup(hass: HomeAssistant, subgroup_entry: Confi
            and subgroup_entry.entry_id in (entry.data.get(CONF_SUB_GROUPS) or [])
     ]
 
+
 @callback
 def create_grouped_power_sensor(
     hass: HomeAssistant,