From df608b56a5947d64f4cff2d22f5ba970155277a6 Mon Sep 17 00:00:00 2001 From: epenet <6771947+epenet@users.noreply.github.com> Date: Fri, 10 Dec 2021 08:40:53 +0100 Subject: [PATCH] Use new enums in energy (#61386) Co-authored-by: epenet --- homeassistant/components/energy/sensor.py | 34 ++++++++++----------- homeassistant/components/energy/validate.py | 29 ++++++++++-------- 2 files changed, 33 insertions(+), 30 deletions(-) diff --git a/homeassistant/components/energy/sensor.py b/homeassistant/components/energy/sensor.py index 0800b02330e..636c6dc32ae 100644 --- a/homeassistant/components/energy/sensor.py +++ b/homeassistant/components/energy/sensor.py @@ -10,11 +10,9 @@ from typing import Any, Final, Literal, TypeVar, cast from homeassistant.components.sensor import ( ATTR_LAST_RESET, ATTR_STATE_CLASS, - DEVICE_CLASS_MONETARY, - STATE_CLASS_MEASUREMENT, - STATE_CLASS_TOTAL, - STATE_CLASS_TOTAL_INCREASING, + SensorDeviceClass, SensorEntity, + SensorStateClass, ) from homeassistant.components.sensor.recorder import reset_detected from homeassistant.const import ( @@ -22,7 +20,6 @@ from homeassistant.const import ( ENERGY_KILO_WATT_HOUR, ENERGY_MEGA_WATT_HOUR, ENERGY_WATT_HOUR, - ENTITY_CATEGORY_SYSTEM, VOLUME_CUBIC_METERS, ) from homeassistant.core import ( @@ -32,6 +29,7 @@ from homeassistant.core import ( split_entity_id, valid_entity_id, ) +from homeassistant.helpers.entity import EntityCategory from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.event import async_track_state_change_event from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType @@ -41,9 +39,9 @@ from .const import DOMAIN from .data import EnergyManager, async_get_manager SUPPORTED_STATE_CLASSES = [ - STATE_CLASS_MEASUREMENT, - STATE_CLASS_TOTAL, - STATE_CLASS_TOTAL_INCREASING, + SensorStateClass.MEASUREMENT, + SensorStateClass.TOTAL, + SensorStateClass.TOTAL_INCREASING, ] VALID_ENERGY_UNITS = [ENERGY_WATT_HOUR, ENERGY_KILO_WATT_HOUR, ENERGY_MEGA_WATT_HOUR] VALID_ENERGY_UNITS_GAS = [VOLUME_CUBIC_METERS] + VALID_ENERGY_UNITS @@ -215,7 +213,7 @@ class EnergyCostSensor(SensorEntity): utility. """ - _attr_entity_category = ENTITY_CATEGORY_SYSTEM + _attr_entity_category = EntityCategory.SYSTEM _wrong_state_class_reported = False _wrong_unit_reported = False @@ -231,8 +229,8 @@ class EnergyCostSensor(SensorEntity): self.entity_id = ( f"{config[adapter.entity_energy_key]}_{adapter.entity_id_suffix}" ) - self._attr_device_class = DEVICE_CLASS_MONETARY - self._attr_state_class = STATE_CLASS_TOTAL + self._attr_device_class = SensorDeviceClass.MONETARY + self._attr_state_class = SensorStateClass.TOTAL self._config = config self._last_energy_sensor_state: State | None = None # add_finished is set when either of async_added_to_hass or add_to_platform_abort @@ -267,9 +265,9 @@ class EnergyCostSensor(SensorEntity): ) return - # last_reset must be set if the sensor is STATE_CLASS_MEASUREMENT + # last_reset must be set if the sensor is SensorStateClass.MEASUREMENT if ( - state_class == STATE_CLASS_MEASUREMENT + state_class == SensorStateClass.MEASUREMENT and ATTR_LAST_RESET not in energy_state.attributes ): return @@ -337,14 +335,16 @@ class EnergyCostSensor(SensorEntity): ) return - if state_class != STATE_CLASS_TOTAL_INCREASING and energy_state.attributes.get( - ATTR_LAST_RESET - ) != self._last_energy_sensor_state.attributes.get(ATTR_LAST_RESET): + if ( + state_class != SensorStateClass.TOTAL_INCREASING + and energy_state.attributes.get(ATTR_LAST_RESET) + != self._last_energy_sensor_state.attributes.get(ATTR_LAST_RESET) + ): # Energy meter was reset, reset cost sensor too energy_state_copy = copy.copy(energy_state) energy_state_copy.state = "0.0" self._reset(energy_state_copy) - elif state_class == STATE_CLASS_TOTAL_INCREASING and reset_detected( + elif state_class == SensorStateClass.TOTAL_INCREASING and reset_detected( self.hass, cast(str, self._config[self._adapter.entity_energy_key]), energy, diff --git a/homeassistant/components/energy/validate.py b/homeassistant/components/energy/validate.py index d77ea75d36c..d9cd9a73aa0 100644 --- a/homeassistant/components/energy/validate.py +++ b/homeassistant/components/energy/validate.py @@ -21,19 +21,22 @@ from homeassistant.core import HomeAssistant, callback, valid_entity_id from . import data from .const import DOMAIN -ENERGY_USAGE_DEVICE_CLASSES = (sensor.DEVICE_CLASS_ENERGY,) +ENERGY_USAGE_DEVICE_CLASSES = (sensor.SensorDeviceClass.ENERGY,) ENERGY_USAGE_UNITS = { - sensor.DEVICE_CLASS_ENERGY: (ENERGY_KILO_WATT_HOUR, ENERGY_WATT_HOUR) + sensor.SensorDeviceClass.ENERGY: (ENERGY_KILO_WATT_HOUR, ENERGY_WATT_HOUR) } ENERGY_PRICE_UNITS = tuple( f"/{unit}" for units in ENERGY_USAGE_UNITS.values() for unit in units ) ENERGY_UNIT_ERROR = "entity_unexpected_unit_energy" ENERGY_PRICE_UNIT_ERROR = "entity_unexpected_unit_energy_price" -GAS_USAGE_DEVICE_CLASSES = (sensor.DEVICE_CLASS_ENERGY, sensor.DEVICE_CLASS_GAS) +GAS_USAGE_DEVICE_CLASSES = ( + sensor.SensorDeviceClass.ENERGY, + sensor.SensorDeviceClass.GAS, +) GAS_USAGE_UNITS = { - sensor.DEVICE_CLASS_ENERGY: (ENERGY_WATT_HOUR, ENERGY_KILO_WATT_HOUR), - sensor.DEVICE_CLASS_GAS: (VOLUME_CUBIC_METERS, VOLUME_CUBIC_FEET), + sensor.SensorDeviceClass.ENERGY: (ENERGY_WATT_HOUR, ENERGY_KILO_WATT_HOUR), + sensor.SensorDeviceClass.GAS: (VOLUME_CUBIC_METERS, VOLUME_CUBIC_FEET), } GAS_PRICE_UNITS = tuple( f"/{unit}" for units in GAS_USAGE_UNITS.values() for unit in units @@ -141,9 +144,9 @@ def _async_validate_usage_stat( state_class = state.attributes.get(sensor.ATTR_STATE_CLASS) allowed_state_classes = [ - sensor.STATE_CLASS_MEASUREMENT, - sensor.STATE_CLASS_TOTAL, - sensor.STATE_CLASS_TOTAL_INCREASING, + sensor.SensorStateClass.MEASUREMENT, + sensor.SensorStateClass.TOTAL, + sensor.SensorStateClass.TOTAL_INCREASING, ] if state_class not in allowed_state_classes: result.append( @@ -155,7 +158,7 @@ def _async_validate_usage_stat( ) if ( - state_class == sensor.STATE_CLASS_MEASUREMENT + state_class == sensor.SensorStateClass.MEASUREMENT and sensor.ATTR_LAST_RESET not in state.attributes ): result.append( @@ -221,9 +224,9 @@ def _async_validate_cost_stat( state_class = state.attributes.get("state_class") supported_state_classes = [ - sensor.STATE_CLASS_MEASUREMENT, - sensor.STATE_CLASS_TOTAL, - sensor.STATE_CLASS_TOTAL_INCREASING, + sensor.SensorStateClass.MEASUREMENT, + sensor.SensorStateClass.TOTAL, + sensor.SensorStateClass.TOTAL_INCREASING, ] if state_class not in supported_state_classes: result.append( @@ -231,7 +234,7 @@ def _async_validate_cost_stat( ) if ( - state_class == sensor.STATE_CLASS_MEASUREMENT + state_class == sensor.SensorStateClass.MEASUREMENT and sensor.ATTR_LAST_RESET not in state.attributes ): result.append(