From 5c6525b13a4b61ba81b737aef08b02e2d5a20c79 Mon Sep 17 00:00:00 2001 From: epenet <6771947+epenet@users.noreply.github.com> Date: Tue, 7 Jul 2026 14:31:58 +0200 Subject: [PATCH] Use state attribute enums in derivative (#175848) Co-authored-by: Claude Opus 4.8 (1M context) --- .../components/derivative/config_flow.py | 14 ++++++-------- homeassistant/components/derivative/sensor.py | 19 ++++++++++++------- 2 files changed, 18 insertions(+), 15 deletions(-) diff --git a/homeassistant/components/derivative/config_flow.py b/homeassistant/components/derivative/config_flow.py index 53e4e5dfa727..437eb92e336c 100644 --- a/homeassistant/components/derivative/config_flow.py +++ b/homeassistant/components/derivative/config_flow.py @@ -8,12 +8,7 @@ import voluptuous as vol from homeassistant.components.counter import DOMAIN as COUNTER_DOMAIN from homeassistant.components.input_number import DOMAIN as INPUT_NUMBER_DOMAIN from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN -from homeassistant.const import ( - ATTR_UNIT_OF_MEASUREMENT, - CONF_NAME, - CONF_SOURCE, - UnitOfTime, -) +from homeassistant.const import CONF_NAME, CONF_SOURCE, EntityStateAttribute, UnitOfTime from homeassistant.core import callback from homeassistant.helpers import selector from homeassistant.helpers.schema_config_entry_flow import ( @@ -59,13 +54,16 @@ def entity_selector_compatible( """Return an entity selector which compatible entities.""" current = handler.hass.states.get(handler.options[CONF_SOURCE]) unit_of_measurement = ( - current.attributes.get(ATTR_UNIT_OF_MEASUREMENT) if current else None + current.attributes.get(EntityStateAttribute.UNIT_OF_MEASUREMENT) + if current + else None ) entities = [ ent.entity_id for ent in handler.hass.states.async_all(ALLOWED_DOMAINS) - if ent.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == unit_of_measurement + if ent.attributes.get(EntityStateAttribute.UNIT_OF_MEASUREMENT) + == unit_of_measurement and ent.domain in ALLOWED_DOMAINS ] diff --git a/homeassistant/components/derivative/sensor.py b/homeassistant/components/derivative/sensor.py index ab7431fbef2c..0e25e72ea1b4 100644 --- a/homeassistant/components/derivative/sensor.py +++ b/homeassistant/components/derivative/sensor.py @@ -8,23 +8,22 @@ from typing import override import voluptuous as vol from homeassistant.components.sensor import ( - ATTR_STATE_CLASS, DEVICE_CLASS_UNITS, PLATFORM_SCHEMA as SENSOR_PLATFORM_SCHEMA, RestoreSensor, SensorDeviceClass, SensorEntity, + SensorEntityCapabilityAttribute, SensorStateClass, ) from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( - ATTR_DEVICE_CLASS, - ATTR_UNIT_OF_MEASUREMENT, CONF_NAME, CONF_SOURCE, CONF_UNIQUE_ID, STATE_UNAVAILABLE, STATE_UNKNOWN, + EntityStateAttribute, Platform, UnitOfTime, ) @@ -243,7 +242,9 @@ class DerivativeSensor(RestoreSensor, SensorEntity): if not source_state: return - source_class_raw = source_state.attributes.get(ATTR_DEVICE_CLASS) + source_class_raw = source_state.attributes.get( + EntityStateAttribute.DEVICE_CLASS + ) source_class: SensorDeviceClass | None = None if isinstance(source_class_raw, str): try: @@ -252,7 +253,9 @@ class DerivativeSensor(RestoreSensor, SensorEntity): source_class = None if self._string_unit_prefix is not None and self._string_unit_time is not None: original_unit = self._attr_native_unit_of_measurement - source_unit = source_state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) + source_unit = source_state.attributes.get( + EntityStateAttribute.UNIT_OF_MEASUREMENT + ) if ( ( source_class @@ -366,7 +369,9 @@ class DerivativeSensor(RestoreSensor, SensorEntity): last_state = await self.async_get_last_state() if last_state: - self._attr_device_class = last_state.attributes.get(ATTR_DEVICE_CLASS) + self._attr_device_class = last_state.attributes.get( + EntityStateAttribute.DEVICE_CLASS + ) @override async def async_added_to_hass(self) -> None: @@ -540,7 +545,7 @@ class DerivativeSensor(RestoreSensor, SensorEntity): # A negative derivative for a total increasing sensor likely indicates the # sensor has been reset. To prevent inaccurate data, discard this sample. if ( - new_state.attributes.get(ATTR_STATE_CLASS) + new_state.attributes.get(SensorEntityCapabilityAttribute.STATE_CLASS) == SensorStateClass.TOTAL_INCREASING and new_derivative < 0 ):