From 46a8325556dde5f14224f06bdf09ea7e0fd2ef6c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ab=C3=ADlio=20Costa?= Date: Mon, 31 Mar 2025 11:32:30 +0100 Subject: [PATCH] Simplify Energy cost sensor update method (#138961) --- homeassistant/components/energy/sensor.py | 112 +++++++++++++--------- 1 file changed, 67 insertions(+), 45 deletions(-) diff --git a/homeassistant/components/energy/sensor.py b/homeassistant/components/energy/sensor.py index eec92c32f98..062601eb4c5 100644 --- a/homeassistant/components/energy/sensor.py +++ b/homeassistant/components/energy/sensor.py @@ -25,6 +25,7 @@ from homeassistant.core import ( split_entity_id, valid_entity_id, ) +from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers import entity_registry as er from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.event import async_track_state_change_event @@ -122,6 +123,10 @@ SOURCE_ADAPTERS: Final = ( ) +class EntityNotFoundError(HomeAssistantError): + """When a referenced entity was not found.""" + + class SensorManager: """Class to handle creation/removal of sensor data.""" @@ -311,43 +316,25 @@ class EnergyCostSensor(SensorEntity): except ValueError: return - # Determine energy price - if self._config["entity_energy_price"] is not None: - energy_price_state = self.hass.states.get( - self._config["entity_energy_price"] + try: + energy_price, energy_price_unit = self._get_energy_price( + valid_units, default_price_unit ) - - if energy_price_state is None: - return - - try: - energy_price = float(energy_price_state.state) - except ValueError: - if self._last_energy_sensor_state is None: - # Initialize as it's the first time all required entities except - # price are in place. This means that the cost will update the first - # time the energy is updated after the price entity is in place. - self._reset(energy_state) - return - - energy_price_unit: str | None = energy_price_state.attributes.get( - ATTR_UNIT_OF_MEASUREMENT, "" - ).partition("/")[2] - - # For backwards compatibility we don't validate the unit of the price - # If it is not valid, we assume it's our default price unit. - if energy_price_unit not in valid_units: - energy_price_unit = default_price_unit - - else: - energy_price = cast(float, self._config["number_energy_price"]) - energy_price_unit = default_price_unit + except EntityNotFoundError: + return + except ValueError: + energy_price = None if self._last_energy_sensor_state is None: - # Initialize as it's the first time all required entities are in place. + # Initialize as it's the first time all required entities are in place or + # only the price is missing. In the later case, cost will update the first + # time the energy is updated after the price entity is in place. self._reset(energy_state) return + if energy_price is None: + return + energy_unit: str | None = energy_state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) if energy_unit is None or energy_unit not in valid_units: @@ -383,20 +370,9 @@ class EnergyCostSensor(SensorEntity): old_energy_value = float(self._last_energy_sensor_state.state) cur_value = cast(float, self._attr_native_value) - if energy_price_unit is None: - converted_energy_price = energy_price - else: - converter: Callable[[float, str, str], float] - if energy_unit in VALID_ENERGY_UNITS: - converter = unit_conversion.EnergyConverter.convert - else: - converter = unit_conversion.VolumeConverter.convert - - converted_energy_price = converter( - energy_price, - energy_unit, - energy_price_unit, - ) + converted_energy_price = self._convert_energy_price( + energy_price, energy_price_unit, energy_unit + ) self._attr_native_value = ( cur_value + (energy - old_energy_value) * converted_energy_price @@ -404,6 +380,52 @@ class EnergyCostSensor(SensorEntity): self._last_energy_sensor_state = energy_state + def _get_energy_price( + self, valid_units: set[str], default_unit: str | None + ) -> tuple[float, str | None]: + """Get the energy price. + + Raises: + EntityNotFoundError: When the energy price entity is not found. + ValueError: When the entity state is not a valid float. + + """ + + if self._config["entity_energy_price"] is None: + return cast(float, self._config["number_energy_price"]), default_unit + + energy_price_state = self.hass.states.get(self._config["entity_energy_price"]) + if energy_price_state is None: + raise EntityNotFoundError + + energy_price = float(energy_price_state.state) + + energy_price_unit: str | None = energy_price_state.attributes.get( + ATTR_UNIT_OF_MEASUREMENT, "" + ).partition("/")[2] + + # For backwards compatibility we don't validate the unit of the price + # If it is not valid, we assume it's our default price unit. + if energy_price_unit not in valid_units: + energy_price_unit = default_unit + + return energy_price, energy_price_unit + + def _convert_energy_price( + self, energy_price: float, energy_price_unit: str | None, energy_unit: str + ) -> float: + """Convert the energy price to the correct unit.""" + if energy_price_unit is None: + return energy_price + + converter: Callable[[float, str, str], float] + if energy_unit in VALID_ENERGY_UNITS: + converter = unit_conversion.EnergyConverter.convert + else: + converter = unit_conversion.VolumeConverter.convert + + return converter(energy_price, energy_unit, energy_price_unit) + async def async_added_to_hass(self) -> None: """Register callbacks.""" energy_state = self.hass.states.get(self._config[self._adapter.stat_energy_key])