From 3760e5073cae09a9c6464fe3310e9d70e90b5aa9 Mon Sep 17 00:00:00 2001 From: Bram Date: Thu, 26 Dec 2024 10:08:55 +0100 Subject: [PATCH] fix: cleanup --- custom_components/powercalc/config_flow.py | 4 +-- .../powercalc/group_include/include.py | 2 +- custom_components/powercalc/sensor.py | 4 +-- .../powercalc/sensors/group/custom.py | 4 +-- .../sensors/group/tracked_untracked.py | 31 ++++++++++--------- 5 files changed, 24 insertions(+), 21 deletions(-) diff --git a/custom_components/powercalc/config_flow.py b/custom_components/powercalc/config_flow.py index aa9ace8ea..bef124f82 100644 --- a/custom_components/powercalc/config_flow.py +++ b/custom_components/powercalc/config_flow.py @@ -120,7 +120,7 @@ ) from .discovery import get_power_profile_by_source_entity from .errors import ModelNotSupportedError, StrategyConfigurationError -from .group_include.include import resolve_include_entities +from .group_include.include import find_entities from .power_profile.factory import get_power_profile from .power_profile.library import ModelInfo, ProfileLibrary from .power_profile.power_profile import DEVICE_TYPE_DOMAIN, DeviceType, PowerProfile @@ -1474,7 +1474,7 @@ async def async_step_group_tracked_untracked_manual(self, user_input: dict[str, """Handle the flow for tracked/untracked group sensor.""" schema = SCHEMA_GROUP_TRACKED_UNTRACKED_MANUAL if not user_input: - entities, _ = await resolve_include_entities(self.hass) + entities, _ = await find_entities(self.hass) tracked_entities = [entity.entity_id for entity in entities if isinstance(entity, PowerSensor)] schema = self.fill_schema_defaults(schema, {CONF_GROUP_TRACKED_POWER_ENTITIES: tracked_entities}) diff --git a/custom_components/powercalc/group_include/include.py b/custom_components/powercalc/group_include/include.py index 618f7c57f..8a9c5a231 100644 --- a/custom_components/powercalc/group_include/include.py +++ b/custom_components/powercalc/group_include/include.py @@ -20,7 +20,7 @@ _LOGGER = logging.getLogger(__name__) -async def resolve_include_entities( +async def find_entities( hass: HomeAssistant, entity_filter: EntityFilter | None = None, include_non_powercalc: bool = True, diff --git a/custom_components/powercalc/sensor.py b/custom_components/powercalc/sensor.py index c4f41349b..5a08c6e70 100644 --- a/custom_components/powercalc/sensor.py +++ b/custom_components/powercalc/sensor.py @@ -141,7 +141,7 @@ SensorConfigurationError, ) from .group_include.filter import FilterOperator, create_composite_filter -from .group_include.include import resolve_include_entities +from .group_include.include import find_entities from .sensors.daily_energy import ( DAILY_FIXED_ENERGY_SCHEMA, create_daily_fixed_energy_power_sensor, @@ -762,7 +762,7 @@ async def add_discovered_entities( include_config: dict = cast(dict, config[CONF_INCLUDE]) include_non_powercalc: bool = include_config.get(CONF_INCLUDE_NON_POWERCALC_SENSORS, True) entity_filter = create_composite_filter(include_config, hass, FilterOperator.AND) - found_entities, discoverable_entities = await resolve_include_entities(hass, entity_filter, include_non_powercalc) + found_entities, discoverable_entities = await find_entities(hass, entity_filter, include_non_powercalc) entities_to_add.existing.extend(found_entities) for entity_id in discoverable_entities: sensor_configs[entity_id] = {CONF_ENTITY_ID: entity_id} diff --git a/custom_components/powercalc/sensors/group/custom.py b/custom_components/powercalc/sensors/group/custom.py index d21c56c56..950362c66 100644 --- a/custom_components/powercalc/sensors/group/custom.py +++ b/custom_components/powercalc/sensors/group/custom.py @@ -88,7 +88,7 @@ ) from custom_components.powercalc.device_binding import get_device_info from custom_components.powercalc.group_include.filter import AreaFilter -from custom_components.powercalc.group_include.include import resolve_include_entities +from custom_components.powercalc.group_include.include import find_entities from custom_components.powercalc.sensors.abstract import ( BaseEntity, generate_energy_sensor_entity_id, @@ -269,7 +269,7 @@ async def add_area_entities() -> None: return entity_filter = AreaFilter(hass, entry.data[CONF_AREA]) - resolved_area_entities, _ = await resolve_include_entities( + resolved_area_entities, _ = await find_entities( hass, entity_filter, bool(entry.data.get(CONF_INCLUDE_NON_POWERCALC_SENSORS)), diff --git a/custom_components/powercalc/sensors/group/tracked_untracked.py b/custom_components/powercalc/sensors/group/tracked_untracked.py index 55f439007..16d571335 100644 --- a/custom_components/powercalc/sensors/group/tracked_untracked.py +++ b/custom_components/powercalc/sensors/group/tracked_untracked.py @@ -24,7 +24,7 @@ GroupType, UnitPrefix, ) -from custom_components.powercalc.group_include.include import resolve_include_entities +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, @@ -101,19 +101,22 @@ async def create_tracked_untracked_group_sensors(self) -> list[Entity]: return entities async def get_tracked_power_entities(self) -> set[str]: - """Get all power entities which are part of the tracked sensor group""" - if bool(self.config.get(CONF_GROUP_TRACKED_AUTO, False)): - - @callback - def _start_entity_registry_listener(_: Any) -> None: # noqa ANN401 - 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) - - entities, _ = await resolve_include_entities(self.hass) - return {entity.entity_id for entity in entities if isinstance(entity, PowerSensor) and not isinstance(entity, GroupedSensor)} - - return set(self.config.get(CONF_GROUP_TRACKED_POWER_ENTITIES)) # type: ignore + """ + Get all power entities which are part of the tracked sensor group + """ + if not bool(self.config.get(CONF_GROUP_TRACKED_AUTO, False)): + return set(self.config.get(CONF_GROUP_TRACKED_POWER_ENTITIES)) # type: ignore + + # 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(_: Any) -> None: # noqa ANN401 + 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) + + entities, _ = await find_entities(self.hass) + return {entity.entity_id for entity in entities if isinstance(entity, PowerSensor) and not isinstance(entity, GroupedSensor)} async def _handle_entity_registry_updated( self,