Merge pull request #4286 from bramstroker/fix/refactor-cleanup

Several cleanups removing duplication

Author: bramstroker

custom_components/powercalc/__init__.py

@@ -105,7 +105,7 @@
     SensorType,
     UnitPrefix,
 )
-from .discovery import DiscoveryManager, DiscoveryStatus
+from .discovery import DiscoveryManager, DiscoveryStatus, get_discovery_manager
 from .migrate import async_fix_legacy_profile_config_entry, async_migrate_config_entry
 from .power_profile.power_profile import DeviceType
 from .sensor import SENSOR_CONFIG
@@ -307,7 +307,7 @@ async def _handle_change_gui_service(call: ServiceCall) -> None:
 
     async def _handle_update_library_service(_: ServiceCall) -> None:
         _LOGGER.info("Updating library and rediscovering devices")
-        discovery_manager: DiscoveryManager = hass.data[DOMAIN][DATA_DISCOVERY_MANAGER]
+        discovery_manager = get_discovery_manager(hass)
         await discovery_manager.update_library_and_rediscover()
 
     hass.services.async_register(
@@ -461,7 +461,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     await async_fix_legacy_profile_config_entry(hass, entry)
     await hass.config_entries.async_forward_entry_setups(entry, [Platform.SENSOR, Platform.SELECT])
-    # await hass.config_entries.async_forward_entry_setups(entry, [Platform.SENSOR])
 
     entry.async_on_unload(entry.add_update_listener(async_update_entry))
 
@@ -474,7 +473,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             await apply_global_gui_configuration_changes(hass)
 
         discovery_enabled = bool(entry.data.get(CONF_DISCOVERY, {}).get(CONF_ENABLED, False))
-        discovery_manager: DiscoveryManager = hass.data[DOMAIN][DATA_DISCOVERY_MANAGER]
+        discovery_manager = get_discovery_manager(hass)
         if discovery_enabled and discovery_manager.status == DiscoveryStatus.DISABLED:
             _LOGGER.debug("Enabling discovery manager based on global configuration")
             discovery_manager.enable()
@@ -530,7 +529,7 @@ async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
 
 async def async_remove_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> None:
     """Called after a config entry is removed."""
-    discovery_manager: DiscoveryManager = hass.data[DOMAIN][DATA_DISCOVERY_MANAGER]
+    discovery_manager = get_discovery_manager(hass)
     discovery_manager.remove_initialized_flow(config_entry)
 
     updated_entries: list[ConfigEntry] = []

custom_components/powercalc/config_flow.py

@@ -103,14 +103,11 @@
     Step.WLED,
 ]
 
+# Order matters: async_step() delegates to the first handler that defines the requested step.
 FLOW_HANDLERS: dict[FlowType, dict] = {
-    FlowType.GROUP: {
-        "config": GroupConfigFlow,
-        "options": GroupOptionsFlow,
-    },
-    FlowType.DAILY_ENERGY: {
-        "config": DailyEnergyConfigFlow,
-        "options": DailyEnergyOptionsFlow,
+    FlowType.GLOBAL_CONFIGURATION: {
+        "config": GlobalConfigurationConfigFlow,
+        "options": GlobalConfigurationOptionsFlow,
     },
     FlowType.LIBRARY: {
         "config": LibraryConfigFlow,
@@ -120,9 +117,13 @@
         "config": VirtualPowerConfigFlow,
         "options": VirtualPowerOptionsFlow,
     },
-    FlowType.GLOBAL_CONFIGURATION: {
-        "config": GlobalConfigurationConfigFlow,
-        "options": GlobalConfigurationOptionsFlow,
+    FlowType.GROUP: {
+        "config": GroupConfigFlow,
+        "options": GroupOptionsFlow,
+    },
+    FlowType.DAILY_ENERGY: {
+        "config": DailyEnergyConfigFlow,
+        "options": DailyEnergyOptionsFlow,
     },
     FlowType.REAL_POWER: {
         "config": RealPowerConfigFlow,
@@ -153,17 +154,9 @@ def __init__(self) -> None:
         self.name: str | None = None
         self.handled_steps: list[Step] = []
 
-        # Initialize flow handlers
+        # Initialize flow handlers. Iteration order follows FLOW_HANDLERS, which async_step() relies on.
         flow_key = "options" if self.is_options_flow else "config"
-        self.flow_handlers = {
-            FlowType.GLOBAL_CONFIGURATION: FLOW_HANDLERS[FlowType.GLOBAL_CONFIGURATION][flow_key](self),
-            FlowType.LIBRARY: FLOW_HANDLERS[FlowType.LIBRARY][flow_key](self),
-            FlowType.VIRTUAL_POWER: FLOW_HANDLERS[FlowType.VIRTUAL_POWER][flow_key](self),
-            FlowType.GROUP: FLOW_HANDLERS[FlowType.GROUP][flow_key](self),
-            FlowType.DAILY_ENERGY: FLOW_HANDLERS[FlowType.DAILY_ENERGY][flow_key](self),
-            FlowType.REAL_POWER: FLOW_HANDLERS[FlowType.REAL_POWER][flow_key](self),
-            FlowType.COST: FLOW_HANDLERS[FlowType.COST][flow_key](self),
-        }
+        self.flow_handlers = {flow_type: handlers[flow_key](self) for flow_type, handlers in FLOW_HANDLERS.items()}
 
         for step in Step:
             step_method = f"async_step_{step}"

custom_components/powercalc/const.py

@@ -14,6 +14,7 @@
     STATE_OPEN,
     STATE_STANDBY,
     STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
     EntityCategory,
 )
 
@@ -272,6 +273,7 @@ class UnitPrefix(StrEnum):
 SIGNAL_POWER_SENSOR_STATE_CHANGE = "powercalc_power_sensor_state_change"
 
 OFF_STATES = {STATE_OFF, STATE_STANDBY, STATE_UNAVAILABLE}
+UNAVAILABLE_STATES = frozenset({STATE_UNAVAILABLE, STATE_UNKNOWN})
 OFF_STATES_BY_DOMAIN: dict[str, set[str]] = {
     cover.DOMAIN: {STATE_CLOSED, STATE_OPEN},
     device_tracker.DOMAIN: {STATE_NOT_HOME},

custom_components/powercalc/discovery.py

@@ -51,35 +51,38 @@
 _DiscoverySourceT = TypeVar("_DiscoverySourceT", er.RegistryEntry, dr.DeviceEntry)
 
 
-async def get_power_profile_by_source_entity(hass: HomeAssistant, source_entity: SourceEntity) -> PowerProfile | None:
-    """Given a certain entity, lookup the manufacturer and model and return the power profile."""
+def get_discovery_manager(hass: HomeAssistant) -> DiscoveryManager:
+    """Return the shared discovery manager, creating a throwaway one when not yet set up."""
     try:
-        discovery_manager: DiscoveryManager = hass.data[DOMAIN][DATA_DISCOVERY_MANAGER]
+        return hass.data[DOMAIN][DATA_DISCOVERY_MANAGER]  # type: ignore[no-any-return]
     except KeyError:
-        discovery_manager = DiscoveryManager(hass, {})
+        return DiscoveryManager(hass, {})
+
+
+async def _get_power_profile_by_source(
+    hass: HomeAssistant,
+    source_entity: SourceEntity,
+    discovery_by: DiscoveryBy,
+) -> PowerProfile | None:
+    """Look up a power profile for a source entity, discovered either by entity or by device."""
+    discovery_manager = get_discovery_manager(hass)
     model_info = await discovery_manager.extract_model_info_from_device_info(source_entity.entity_entry)
     if not model_info:
         return None
-    profiles = await discovery_manager.find_power_profiles(model_info, source_entity, DiscoveryBy.ENTITY)
+    profiles = await discovery_manager.find_power_profiles(model_info, source_entity, discovery_by)
     return profiles[0] if profiles else None
 
 
+async def get_power_profile_by_source_entity(hass: HomeAssistant, source_entity: SourceEntity) -> PowerProfile | None:
+    """Given a certain entity, lookup the manufacturer and model and return the power profile."""
+    return await _get_power_profile_by_source(hass, source_entity, DiscoveryBy.ENTITY)
+
+
 async def get_power_profile_by_source_device(hass: HomeAssistant, source_entity: SourceEntity) -> PowerProfile | None:
     """Look up a device-discovered power profile for a source entity's device."""
     if not source_entity.device_entry or not source_entity.entity_entry:
         return None
-
-    try:
-        discovery_manager: DiscoveryManager = hass.data[DOMAIN][DATA_DISCOVERY_MANAGER]
-    except KeyError:
-        discovery_manager = DiscoveryManager(hass, {})
-
-    model_info = await discovery_manager.extract_model_info_from_device_info(source_entity.entity_entry)
-    if not model_info:
-        return None
-
-    profiles = await discovery_manager.find_power_profiles(model_info, source_entity, DiscoveryBy.DEVICE)
-    return profiles[0] if profiles else None
+    return await _get_power_profile_by_source(hass, source_entity, DiscoveryBy.DEVICE)
 
 
 class DiscoveryStatus(StrEnum):

custom_components/powercalc/flow_helper/dynamic_field_builder.py

@@ -23,12 +23,21 @@ def build_dynamic_field_schema(
         else:
             key = vol.Required(field.key, description=field_description)
 
+        field_selector = field.selector
         if "entity" in field.selector and source_entity and source_entity.device_entry:
             entity_reg = er.async_get(hass)
-            field.selector["entity"]["include_entities"] = [
-                entity.entity_id
-                for entity in entity_reg.entities.get_entries_for_device_id(source_entity.device_entry.id)
-            ]
+            # Build a new selector dict instead of mutating field.selector, which is a reference
+            # into the (potentially cached) profile json_data.
+            field_selector = {
+                **field.selector,
+                "entity": {
+                    **field.selector["entity"],
+                    "include_entities": [
+                        entity.entity_id
+                        for entity in entity_reg.entities.get_entries_for_device_id(source_entity.device_entry.id)
+                    ],
+                },
+            }
 
-        schema[key] = selector(field.selector)
+        schema[key] = selector(field_selector)
     return vol.Schema(schema)

custom_components/powercalc/group_include/filter.py

@@ -11,6 +11,7 @@
 from homeassistant.core import HomeAssistant, split_entity_id
 from homeassistant.helpers import area_registry, device_registry, entity_registry, floor_registry, label_registry
 import homeassistant.helpers.config_validation as cv
+from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.entity_component import EntityComponent
 from homeassistant.helpers.entity_registry import RegistryEntry
 from homeassistant.helpers.template import Template
@@ -148,8 +149,6 @@ def is_valid(self, entity: RegistryEntry) -> bool:
 
 class StandardGroupFilter(EntityFilter):
     def __init__(self, hass: HomeAssistant, group_id: str) -> None:
-        entity_reg = entity_registry.async_get(hass)
-        entity_reg.async_get(group_id)
         group_state = hass.states.get(group_id)
         if group_state is None:
             raise SensorConfigurationError(f"Group state {group_id} not found")
@@ -161,14 +160,7 @@ def is_valid(self, entity: RegistryEntry) -> bool:
 
 class LightGroupFilter(EntityFilter):
     def __init__(self, hass: HomeAssistant, group_id: str) -> None:
-        light_component = cast(EntityComponent, hass.data.get(LIGHT_DOMAIN))
-        light_group = next(
-            filter(
-                lambda entity: entity.entity_id == group_id,
-                light_component.entities,
-            ),
-            None,
-        )
+        light_group = self._find_light_group(hass, group_id)
         if light_group is None or light_group.platform.platform_name != GROUP_DOMAIN:
             raise SensorConfigurationError(f"Light group {group_id} not found")
 
@@ -177,22 +169,26 @@ def __init__(self, hass: HomeAssistant, group_id: str) -> None:
     def is_valid(self, entity: RegistryEntry) -> bool:
         return entity.entity_id in self.entity_ids
 
-    def find_all_entity_ids_recursively(
-        self,
-        hass: HomeAssistant,
-        group_entity_id: str,
-        all_entity_ids: list[str],
-    ) -> list[str]:
-        entity_reg = entity_registry.async_get(hass)
+    @staticmethod
+    def _find_light_group(hass: HomeAssistant, group_entity_id: str) -> Entity | None:
         light_component = cast(EntityComponent, hass.data.get(LIGHT_DOMAIN))
-        light_group = next(
+        return next(
             filter(
                 lambda entity: entity.entity_id == group_entity_id,
                 light_component.entities,
             ),
             None,
         )
 
+    def find_all_entity_ids_recursively(
+        self,
+        hass: HomeAssistant,
+        group_entity_id: str,
+        all_entity_ids: list[str],
+    ) -> list[str]:
+        entity_reg = entity_registry.async_get(hass)
+        light_group = self._find_light_group(hass, group_entity_id)
+
         entity_ids: list[str] = light_group.extra_state_attributes.get(ATTR_ENTITY_ID)  # type: ignore
         for entity_id in entity_ids:
             registry_entry = entity_reg.async_get(entity_id)
@@ -373,7 +369,7 @@ def __init__(
         self.operator = operator
 
     def is_valid(self, entity: RegistryEntry) -> bool:
-        evaluations = [entity_filter.is_valid(entity) for entity_filter in self.filters]
+        evaluations = (entity_filter.is_valid(entity) for entity_filter in self.filters)
         if self.operator == FilterOperator.OR:
             return any(evaluations)
 

custom_components/powercalc/helpers.py

@@ -200,14 +200,6 @@ def _resolve_related_entity_by_device_class(
     return get_related_entity_by_device_class(hass, source_entity, device_class)
 
 
-def _resolve_related_entity_by_translation_key(
-    hass: HomeAssistant,
-    source_entity: SourceEntity,
-    translation_key: str,
-) -> str | None:
-    return get_related_entity_by_translation_key(hass, source_entity, translation_key)
-
-
 RELATED_ENTITY_PLACEHOLDER_DEFINITIONS = (
     RelatedEntityPlaceholderDefinition(
         PLACEHOLDER_ENTITY_BY_DEVICE_CLASS,
@@ -217,7 +209,11 @@ def _resolve_related_entity_by_translation_key(
     RelatedEntityPlaceholderDefinition(
         PLACEHOLDER_ENTITY_BY_TRANSLATION_KEY,
         "translation key",
-        _resolve_related_entity_by_translation_key,
+        lambda hass, source_entity, translation_key: get_related_entity_by_translation_key(
+            hass,
+            source_entity,
+            translation_key,
+        ),
     ),
 )
 

custom_components/powercalc/migrate.py

@@ -51,7 +51,7 @@ async def async_migrate_config_entry(hass: HomeAssistant, config_entry: ConfigEn
     if version <= 3:
         _migrate_playbook_trigger(data)
 
-    if version <= 4 and config_entry.entry_id == ENTRY_GLOBAL_CONFIG_UNIQUE_ID:
+    if version <= 4 and config_entry.unique_id == ENTRY_GLOBAL_CONFIG_UNIQUE_ID:
         _migrate_global_discovery_config(data)
 
     if version <= 5:
@@ -96,16 +96,22 @@ def _migrate_playbook_trigger(data: dict) -> None:
 
 
 def _migrate_global_discovery_config(data: dict) -> None:
+    deprecated_keys = [
+        CONF_ENABLE_AUTODISCOVERY_DEPRECATED,
+        CONF_DISCOVERY_EXCLUDE_DEVICE_TYPES_DEPRECATED,
+        CONF_DISCOVERY_EXCLUDE_SELF_USAGE_DEPRECATED,
+    ]
+    # Nothing to convert when there are no legacy keys and the new format is already present;
+    # avoid clobbering an existing discovery config.
+    if not any(key in data for key in deprecated_keys) and CONF_DISCOVERY in data:
+        return
+
     data[CONF_DISCOVERY] = {
         CONF_ENABLED: data.get(CONF_ENABLE_AUTODISCOVERY_DEPRECATED, True),
         CONF_EXCLUDE_DEVICE_TYPES: data.get(CONF_DISCOVERY_EXCLUDE_DEVICE_TYPES_DEPRECATED, []),
         CONF_EXCLUDE_SELF_USAGE: data.get(CONF_DISCOVERY_EXCLUDE_SELF_USAGE_DEPRECATED, False),
     }
-    for key in [
-        CONF_ENABLE_AUTODISCOVERY_DEPRECATED,
-        CONF_DISCOVERY_EXCLUDE_DEVICE_TYPES_DEPRECATED,
-        CONF_DISCOVERY_EXCLUDE_SELF_USAGE_DEPRECATED,
-    ]:
+    for key in deprecated_keys:
         data.pop(key, None)
 
 

custom_components/powercalc/power_profile/loader/local.py

@@ -2,7 +2,6 @@
 import json
 import logging
 import os
-import re
 from typing import Any, cast
 
 from homeassistant.core import HomeAssistant
@@ -154,8 +153,7 @@ def _load_custom_library(self) -> None:
 
             manufacturer = manufacturer_dir.lower()
             for model_dir in next(os.walk(manufacturer_path))[1]:
-                pattern = re.compile(r"^\..*")
-                if pattern.match(model_dir):
+                if model_dir.startswith("."):
                     continue
 
                 model_path = os.path.join(manufacturer_path, model_dir)

custom_components/powercalc/power_profile/loader/remote.py

@@ -279,7 +279,7 @@ def _get_library_model(self, manufacturer: str, model: str) -> LibraryModel:
         """Retrieve model info, or raise an error if not found."""
         model_info = self.model_infos.get(f"{manufacturer}/{model}")
         if not model_info:
-            raise LibraryLoadingError("Model not found in library: %s/%s", manufacturer, model)
+            raise LibraryLoadingError(f"Model not found in library: {manufacturer}/{model}")
         return model_info
 
     async def _needs_update(
@@ -424,10 +424,14 @@ def _save_file(data: bytes, directory: str) -> None:
         except (TimeoutError, aiohttp.ClientError) as e:
             raise ProfileDownloadError(f"Failed to download profile: {manufacturer}/{model}") from e
 
+    def _get_profile_hashes_path(self) -> str:
+        """Retrieve the local storage path for the profile hashes file."""
+        return str(self.hass.config.path(STORAGE_DIR, BUILT_IN_LIBRARY_DIR, ".profile_hashes"))
+
     def _load_profile_hashes(self) -> dict[str, str]:
         """Load profile hashes from local storage"""
 
-        path = self.hass.config.path(STORAGE_DIR, BUILT_IN_LIBRARY_DIR, ".profile_hashes")
+        path = self._get_profile_hashes_path()
         if not os.path.exists(path):
             return {}
 
@@ -437,6 +441,6 @@ def _load_profile_hashes(self) -> dict[str, str]:
     def _write_profile_hashes(self, hashes: dict[str, str]) -> None:
         """Write profile hashes to local storage"""
 
-        path = self.hass.config.path(STORAGE_DIR, BUILT_IN_LIBRARY_DIR, ".profile_hashes")
+        path = self._get_profile_hashes_path()
         with open(path, "w") as json_file:
             json.dump(hashes, json_file, indent=4)