import inspect import logging from typing import Any from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry, ConfigFlow from homeassistant.const import CONF_NAME from homeassistant.core import HomeAssistant, callback from custom_components.powercalc.const import ( CONF_GROUP, CONF_GROUP_MEMBER_SENSORS, CONF_GROUP_TYPE, CONF_SENSOR_TYPE, CONF_SUB_GROUPS, DOMAIN, ENTRY_GLOBAL_CONFIG_UNIQUE_ID, GroupType, SensorType, ) _LOGGER = logging.getLogger(__name__) def remove_power_sensor_from_associated_groups( hass: HomeAssistant, config_entry: ConfigEntry, ) -> list[ConfigEntry]: """When the user removes a virtual power config entry, update all groups this sensor belongs to.""" group_entries = get_groups_having_member(hass, config_entry) for group_entry in group_entries: # Copy the list so we don't mutate the config entry's data in place, which would make # async_update_entry's change detection see no change and skip persisting/reloading. member_sensors = list(group_entry.data.get(CONF_GROUP_MEMBER_SENSORS) or []) member_sensors.remove(config_entry.entry_id) hass.config_entries.async_update_entry( group_entry, data={**group_entry.data, CONF_GROUP_MEMBER_SENSORS: member_sensors}, ) return group_entries async def add_to_associated_groups(hass: HomeAssistant, config_entry: ConfigEntry) -> None: """ When the user has set a group on a virtual power config entry, we need to add this config entry to the group members sensors and update the group. """ sensor_type = config_entry.data.get(CONF_SENSOR_TYPE) if sensor_type not in [SensorType.VIRTUAL_POWER, SensorType.DAILY_ENERGY]: return raw_groups = config_entry.data.get(CONF_GROUP) if not raw_groups: return group_ids = raw_groups if isinstance(raw_groups, list) else [raw_groups] for group_entry_id in group_ids: group_entry = await add_to_associated_group(hass, config_entry, group_entry_id) if group_entry: _LOGGER.debug( "ConfigEntry %s: Added to group %s.", config_entry.title, group_entry.title, ) # After processed correctly we can want to unset the group, to prevent is being processed again new_data = {k: v for k, v in config_entry.data.items() if k != CONF_GROUP} hass.config_entries.async_update_entry(config_entry, data=new_data) async def add_to_associated_group( hass: HomeAssistant, config_entry: ConfigEntry, group_entry_id: str, ) -> ConfigEntry | None: """When the user has set a group on a virtual power config entry, we need to add this config entry to the group members sensors and update the group. """ group_entry = hass.config_entries.async_get_entry(group_entry_id) # When we are not dealing with a uuid, the user has set a group name manually # Create a new group entry for this group if not group_entry and len(group_entry_id) != 32: group_entry = hass.config_entries.async_entry_for_domain_unique_id(DOMAIN, group_entry_id) if not group_entry: additional_args: dict[str, Any] = {} signature = inspect.signature(ConfigEntry.__init__) if "discovery_keys" in signature.parameters: additional_args["discovery_keys"] = {} if "subentries_data" in signature.parameters: additional_args["subentries_data"] = None group_entry = ConfigEntry( version=ConfigFlow.VERSION, minor_version=ConfigFlow.MINOR_VERSION, domain=DOMAIN, source=SOURCE_IMPORT, title=group_entry_id, data={ CONF_SENSOR_TYPE: SensorType.GROUP, CONF_NAME: group_entry_id, }, options={}, unique_id=group_entry_id, **additional_args, ) await hass.config_entries.async_add(group_entry) if not group_entry: _LOGGER.warning( "ConfigEntry %s: Cannot add/remove to group %s. It does not exist.", config_entry.title, group_entry_id, ) return None member_sensors = set(group_entry.data.get(CONF_GROUP_MEMBER_SENSORS) or []) # Config entry has already been added to associated group. just skip adding it again if config_entry.entry_id in member_sensors: return None member_sensors.add(config_entry.entry_id) hass.config_entries.async_update_entry( group_entry, data={**group_entry.data, CONF_GROUP_MEMBER_SENSORS: list(member_sensors)}, ) return group_entry def get_entries_having_subgroup(hass: HomeAssistant, subgroup_entry: ConfigEntry) -> list[ConfigEntry]: """Get all virtual power entries which have the subgroup in their subgroups list.""" return [ entry for entry in get_group_entries(hass) if subgroup_entry.entry_id in (entry.data.get(CONF_SUB_GROUPS) or []) ] def get_groups_having_member(hass: HomeAssistant, member_entry: ConfigEntry) -> list[ConfigEntry]: """Get all group entries which have the member sensor in their member list.""" return [ entry for entry in hass.config_entries.async_entries(DOMAIN) if entry.data.get(CONF_SENSOR_TYPE) == SensorType.GROUP and member_entry.entry_id in (entry.data.get(CONF_GROUP_MEMBER_SENSORS) or []) ] @callback def get_group_entries(hass: HomeAssistant, group_type: GroupType | None = None) -> list[ConfigEntry]: return [ entry for entry in hass.config_entries.async_entries(DOMAIN) if entry.data.get(CONF_SENSOR_TYPE) == SensorType.GROUP and (group_type is None or entry.data.get(CONF_GROUP_TYPE, GroupType.CUSTOM) == group_type) ] @callback def get_entries_excluding_global_config(hass: HomeAssistant) -> list[ConfigEntry]: return [ entry for entry in hass.config_entries.async_entries(DOMAIN) if entry.unique_id != ENTRY_GLOBAL_CONFIG_UNIQUE_ID ]