from __future__ import annotations from enum import StrEnum import logging from homeassistant.components.sensor import SensorDeviceClass from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_UNIQUE_ID, EVENT_HOMEASSISTANT_STARTED from homeassistant.core import Event, HomeAssistant, callback from homeassistant.helpers.entity import Entity import homeassistant.helpers.entity_registry as er from homeassistant.helpers.entity_registry import EVENT_ENTITY_REGISTRY_UPDATED, EventEntityRegistryUpdatedData from homeassistant.helpers.typing import ConfigType from custom_components.powercalc.const import ( CONF_CREATE_ENERGY_SENSOR, CONF_DISABLE_EXTENDED_ATTRIBUTES, CONF_ENERGY_SENSOR_UNIT_PREFIX, CONF_EXCLUDE_ENTITIES, CONF_GROUP_TRACKED_AUTO, CONF_GROUP_TRACKED_POWER_ENTITIES, CONF_MAIN_POWER_SENSOR, CONF_UTILITY_METER_NET_CONSUMPTION, GroupType, UnitPrefix, ) from custom_components.powercalc.group_include.filter import LambdaFilter from custom_components.powercalc.group_include.include import find_entities from custom_components.powercalc.sensors.abstract import ( generate_energy_sensor_entity_id, generate_energy_sensor_name, generate_power_sensor_entity_id, generate_power_sensor_name, ) from custom_components.powercalc.sensors.energy import VirtualEnergySensor from custom_components.powercalc.sensors.energy_related import create_energy_related_sensors from custom_components.powercalc.sensors.group.custom import GroupedPowerSensor, GroupedSensor from custom_components.powercalc.sensors.group.subtract import SubtractGroupSensor from custom_components.powercalc.sensors.power import PowerSensor _LOGGER = logging.getLogger(__name__) class SensorType(StrEnum): TRACKED = "tracked" UNTRACKED = "untracked" async def find_auto_tracked_power_entities(hass: HomeAssistant, exclude_entities: set[str] | None = None) -> set[str]: """Find tracked power entities.""" entity_filter = None if exclude_entities: entity_filter = LambdaFilter(lambda entity: entity.entity_id not in exclude_entities) result = await find_entities(hass, entity_filter) return { entity.entity_id for entity in result.resolved if isinstance(entity, PowerSensor) and not isinstance(entity, GroupedSensor) } class TrackedPowerSensorFactory: def __init__(self, hass: HomeAssistant, config_entry: ConfigEntry, config: ConfigType) -> None: self.hass = hass self.tracked_entities: set[str] = set() self.config_entry = config_entry self.config = config async def create_tracked_untracked_group_sensors(self) -> list[Entity]: """Create tracked/untracked group sensors.""" unique_id = str(self.config.get(CONF_UNIQUE_ID)) main_power_sensor = ( str(self.config.get(CONF_MAIN_POWER_SENSOR)) if self.config.get(CONF_MAIN_POWER_SENSOR) else None ) self.config[CONF_DISABLE_EXTENDED_ATTRIBUTES] = True # prevent adding all entities in the state attributes self.tracked_entities = await self.get_tracked_power_entities() if main_power_sensor and main_power_sensor in self.tracked_entities: self.tracked_entities.remove(main_power_sensor) should_create_energy_sensor = bool(self.config.get(CONF_CREATE_ENERGY_SENSOR, False)) entities: list[Entity] = [] tracked_sensor = await self.create_tracked_power_sensor(SensorType.TRACKED, unique_id, self.tracked_entities) await self._add_power_sensor_with_energy( entities, SensorType.TRACKED, tracked_sensor, should_create_energy_sensor, ) if main_power_sensor: untracked_sensor = await self.create_untracked_power_sensor( SensorType.UNTRACKED, unique_id, main_power_sensor, tracked_sensor.entity_id, ) await self._add_power_sensor_with_energy( entities, SensorType.UNTRACKED, untracked_sensor, should_create_energy_sensor, ) return entities async def _add_power_sensor_with_energy( self, entities: list[Entity], sensor_type: SensorType, power_sensor: GroupedPowerSensor, should_create_energy_sensor: bool, ) -> None: """Append the power sensor and, when enabled, its energy and energy-related sensors.""" entities.append(power_sensor) if not should_create_energy_sensor: return energy_sensor = await self.create_energy_sensor(sensor_type, power_sensor) entities.append(energy_sensor) entities.extend( create_energy_related_sensors( self.hass, self.config, energy_sensor, utility_meter_config={CONF_UTILITY_METER_NET_CONSUMPTION: True, **self.config}, cost_name=str(sensor_type), ), ) async def get_tracked_power_entities(self) -> set[str]: """ Get all power entities which are part of the tracked sensor group """ if not bool(self.config.get(CONF_GROUP_TRACKED_AUTO, False)): tracked_entities: list[str] | None = self.config.get(CONF_GROUP_TRACKED_POWER_ENTITIES) if not isinstance(tracked_entities, list): return set() return set(tracked_entities) # For auto mode, we also want to listen for any changes in the entity registry # Dynamically add/remove power sensors from the tracked group @callback def _start_entity_registry_listener(_: Event) -> None: self.hass.bus.async_listen(EVENT_ENTITY_REGISTRY_UPDATED, self._handle_entity_registry_updated) self.hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STARTED, _start_entity_registry_listener) exclude_entities = self.config.get(CONF_EXCLUDE_ENTITIES) return await find_auto_tracked_power_entities(self.hass, set(exclude_entities) if exclude_entities else None) async def _handle_entity_registry_updated( self, event: Event[EventEntityRegistryUpdatedData], ) -> None: """Listen to all entity registry updates and reload the config entry if a power sensor is added/removed.""" entity_id = event.data["entity_id"] action = event.data["action"] old_entity_id = event.data.get("old_entity_id") if action == "update" and old_entity_id is not None: if old_entity_id in self.tracked_entities: return await self.reload() return None # pragma: no cover if action == "remove" and entity_id in self.tracked_entities: return await self.reload() if action == "create": registry = er.async_get(self.hass) entity_entry = registry.async_get(entity_id) if entity_entry and entity_entry.original_device_class == SensorDeviceClass.POWER: return await self.reload() return None async def reload(self) -> None: """Reload the config entry.""" await self.hass.config_entries.async_reload(self.config_entry.entry_id) async def create_tracked_power_sensor( self, sensor_type: SensorType, unique_id: str, tracked_entities: set[str], ) -> GroupedPowerSensor: _LOGGER.debug("Creating tracked grouped power sensor, entities: %s", tracked_entities) unique_id = f"{unique_id}_{sensor_type}_power" entity_id = generate_power_sensor_entity_id(self.hass, self.config, name=sensor_type, unique_id=unique_id) name = generate_power_sensor_name(self.config, name=sensor_type) return GroupedPowerSensor( self.hass, sensor_config=self.config, group_type=GroupType.TRACKED_UNTRACKED, entities=tracked_entities, entity_id=entity_id, name=name, unique_id=unique_id, ) async def create_untracked_power_sensor( self, sensor_type: SensorType, unique_id: str, main_power_entity_id: str, tracked_entity_id: str, ) -> GroupedPowerSensor: _LOGGER.debug("Creating untracked grouped power sensor") unique_id = f"{unique_id}_{sensor_type}_power" entity_id = generate_power_sensor_entity_id(self.hass, self.config, name=sensor_type, unique_id=unique_id) name = generate_power_sensor_name(self.config, name=sensor_type) return SubtractGroupSensor( self.hass, entity_id=entity_id, name=name, sensor_config=self.config, base_entity_id=main_power_entity_id, subtract_entities=[tracked_entity_id], unique_id=unique_id, ) async def create_energy_sensor( self, sensor_type: SensorType, power_sensor: GroupedPowerSensor, ) -> VirtualEnergySensor: """Create an energy sensor for a power sensor.""" _LOGGER.debug("Creating %s grouped energy sensor", sensor_type) unique_id = f"{power_sensor.unique_id}_{sensor_type}_energy" name = generate_energy_sensor_name(self.config, sensor_type) entity_id = generate_energy_sensor_entity_id(self.hass, self.config, name=sensor_type, unique_id=unique_id) return VirtualEnergySensor( hass=self.hass, source_entity=power_sensor.entity_id, entity_id=entity_id, name=name, unique_id=unique_id, sensor_config=self.config, unit_prefix=self.config.get(CONF_ENERGY_SENSOR_UNIT_PREFIX, UnitPrefix.KILO), )