from dataclasses import dataclass import logging from homeassistant.components import sensor from homeassistant.components.sensor import SensorDeviceClass from homeassistant.core import HomeAssistant from homeassistant.helpers.entity import Entity from homeassistant.helpers.entity_registry import RegistryEntry from custom_components.powercalc.common import create_source_entity from custom_components.powercalc.const import ( CONF_SENSOR_TYPE, DATA_CONFIGURED_ENTITIES, DATA_ENTITIES, DOMAIN, ENTRY_DATA_ENERGY_ENTITY, ENTRY_DATA_POWER_ENTITY, SensorType, ) from custom_components.powercalc.discovery import get_power_profile_by_source_entity from custom_components.powercalc.power_profile.power_profile import SUPPORTED_DOMAINS from custom_components.powercalc.sensors.energy import RealEnergySensor from custom_components.powercalc.sensors.power import RealPowerSensor from custom_components.powercalc.sensors.utility_meter import VirtualUtilityMeter from .filter import CompositeFilter, DomainFilter, EntityFilter, LambdaFilter, get_filtered_entity_list _LOGGER = logging.getLogger(__name__) @dataclass(slots=True) class FindEntitiesResult: resolved: list[Entity] discoverable: list[str] async def find_entities( hass: HomeAssistant, entity_filter: EntityFilter | None = None, include_non_powercalc: bool = True, exclude_utility_meters: bool = True, ) -> FindEntitiesResult: """ Based on the given entity filter, fetch all power and energy sensors from the HA instance. """ domain_data = hass.data.get(DOMAIN, {}) source_entity_powercalc_entity_map: dict[str, list[tuple[Entity, bool]]] = domain_data.get( DATA_CONFIGURED_ENTITIES, {}, ) powercalc_entities: dict[str, Entity] = domain_data.get( DATA_ENTITIES, {}, ) resolved_entities: list[Entity] = [] discoverable_entities: list[str] = [] source_entities = get_filtered_entity_list(hass, _build_filter(hass, entity_filter, include_non_powercalc)) if _LOGGER.isEnabledFor(logging.DEBUG): # pragma: no cover _LOGGER.debug("Source entities: %s", [entity.entity_id for entity in source_entities]) for source_entity in source_entities: entity_id = source_entity.entity_id mapped = source_entity_powercalc_entity_map.get(entity_id) if mapped: resolved_entities.extend(entity for entity, _ in mapped) continue existing = powercalc_entities.get(entity_id) if existing: resolved_entities.append(existing) continue real_sensor = _create_real_sensor(source_entity) if real_sensor: resolved_entities.append(real_sensor) continue if await _is_discoverable_source_entity(hass, source_entity): discoverable_entities.append(entity_id) if exclude_utility_meters: resolved_entities = [entity for entity in resolved_entities if not isinstance(entity, VirtualUtilityMeter)] if _LOGGER.isEnabledFor(logging.DEBUG): # pragma: no cover _LOGGER.debug("Resolved entities: %s", [entity.entity_id for entity in resolved_entities]) _LOGGER.debug("Discoverable entities: %s", discoverable_entities) return FindEntitiesResult(resolved_entities, discoverable_entities) def _is_source_entity_eligible( hass: HomeAssistant, source_entity: RegistryEntry, include_non_powercalc: bool, ) -> bool: """Return whether a registry entity is eligible for Powercalc discovery.""" if source_entity.platform != DOMAIN: return include_non_powercalc or source_entity.domain != sensor.DOMAIN # YAML-created Powercalc entities have no config entry and are classified by their runtime type later. if source_entity.config_entry_id is None: return True config_entry = hass.config_entries.async_get_entry(source_entity.config_entry_id) if config_entry is None: return False if config_entry.data.get(CONF_SENSOR_TYPE) == SensorType.GROUP: return False main_entity_ids = { config_entry.data.get(ENTRY_DATA_POWER_ENTITY), config_entry.data.get(ENTRY_DATA_ENERGY_ENTITY), } return source_entity.entity_id in main_entity_ids def _create_real_sensor(source_entity: RegistryEntry) -> Entity | None: if source_entity.domain != sensor.DOMAIN: return None device_class = source_entity.device_class or source_entity.original_device_class if device_class == SensorDeviceClass.POWER: return RealPowerSensor(source_entity.entity_id, source_entity.unit_of_measurement) if device_class == SensorDeviceClass.ENERGY: return RealEnergySensor(source_entity.entity_id) return None # pragma: no cover async def _is_discoverable_source_entity(hass: HomeAssistant, source_entity: RegistryEntry) -> bool: power_profile = await get_power_profile_by_source_entity( hass, create_source_entity(source_entity.entity_id, hass), ) return bool( power_profile and not await power_profile.needs_user_configuration and power_profile.is_entity_domain_supported(source_entity), ) def _build_filter( hass: HomeAssistant, entity_filter: EntityFilter | None, include_non_powercalc: bool, ) -> EntityFilter: base_filter = CompositeFilter( [ DomainFilter(SUPPORTED_DOMAINS), LambdaFilter(lambda entity: _is_source_entity_eligible(hass, entity, include_non_powercalc)), LambdaFilter(lambda entity: entity.platform != "utility_meter"), LambdaFilter(lambda entity: not str(entity.unique_id).startswith("powercalc_standby_group")), LambdaFilter(lambda entity: "tracked_" not in str(entity.unique_id)), LambdaFilter( lambda entity: ( entity.platform != "tasmota" or not str(entity.entity_id).endswith(("_yesterday", "_today")) ), ), ], ) if not entity_filter: return base_filter return CompositeFilter([base_filter, entity_filter])