diff --git a/homeassistant/components/esphome/__init__.py b/homeassistant/components/esphome/__init__.py index 8e7f2524a35..4e2e7c02aaa 100644 --- a/homeassistant/components/esphome/__init__.py +++ b/homeassistant/components/esphome/__init__.py @@ -13,7 +13,7 @@ from aioesphomeapi import ( APIIntEnum, APIVersion, DeviceInfo as EsphomeDeviceInfo, - EntityCategory, + EntityCategory as EsphomeEntityCategory, EntityInfo, EntityState, HomeassistantServiceCall, @@ -33,8 +33,6 @@ from homeassistant.const import ( CONF_MODE, CONF_PASSWORD, CONF_PORT, - ENTITY_CATEGORY_CONFIG, - ENTITY_CATEGORY_DIAGNOSTIC, EVENT_HOMEASSISTANT_STOP, ) from homeassistant.core import Event, HomeAssistant, ServiceCall, State, callback @@ -43,7 +41,7 @@ from homeassistant.helpers import template import homeassistant.helpers.config_validation as cv import homeassistant.helpers.device_registry as dr from homeassistant.helpers.dispatcher import async_dispatcher_connect -from homeassistant.helpers.entity import DeviceInfo, Entity +from homeassistant.helpers.entity import DeviceInfo, Entity, EntityCategory from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.event import async_track_state_change_event from homeassistant.helpers.json import JSONEncoder @@ -647,11 +645,13 @@ class EsphomeEnumMapper(Generic[_EnumT, _ValT]): ICON_SCHEMA = vol.Schema(cv.icon) -ENTITY_CATEGORIES: EsphomeEnumMapper[EntityCategory, str | None] = EsphomeEnumMapper( +ENTITY_CATEGORIES: EsphomeEnumMapper[ + EsphomeEntityCategory, EntityCategory | None +] = EsphomeEnumMapper( { - EntityCategory.NONE: None, - EntityCategory.CONFIG: ENTITY_CATEGORY_CONFIG, - EntityCategory.DIAGNOSTIC: ENTITY_CATEGORY_DIAGNOSTIC, + EsphomeEntityCategory.NONE: None, + EsphomeEntityCategory.CONFIG: EntityCategory.CONFIG, + EsphomeEntityCategory.DIAGNOSTIC: EntityCategory.DIAGNOSTIC, } ) @@ -799,7 +799,7 @@ class EsphomeEntity(Entity, Generic[_InfoT, _StateT]): return not self._static_info.disabled_by_default @property - def entity_category(self) -> str | None: + def entity_category(self) -> EntityCategory | None: """Return the category of the entity, if any.""" if not self._static_info.entity_category: return None diff --git a/homeassistant/components/esphome/sensor.py b/homeassistant/components/esphome/sensor.py index c3f9eff6060..45a73a5e5af 100644 --- a/homeassistant/components/esphome/sensor.py +++ b/homeassistant/components/esphome/sensor.py @@ -7,18 +7,17 @@ import math from aioesphomeapi import ( SensorInfo, SensorState, - SensorStateClass, + SensorStateClass as EsphomeSensorStateClass, TextSensorInfo, TextSensorState, ) from aioesphomeapi.model import LastResetType from homeassistant.components.sensor import ( - DEVICE_CLASS_TIMESTAMP, DEVICE_CLASSES, - STATE_CLASS_MEASUREMENT, - STATE_CLASS_TOTAL_INCREASING, + SensorDeviceClass, SensorEntity, + SensorStateClass, ) from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant @@ -61,11 +60,13 @@ async def async_setup_entry( # pylint: disable=invalid-overridden-method -_STATE_CLASSES: EsphomeEnumMapper[SensorStateClass, str | None] = EsphomeEnumMapper( +_STATE_CLASSES: EsphomeEnumMapper[ + EsphomeSensorStateClass, SensorStateClass | None +] = EsphomeEnumMapper( { - SensorStateClass.NONE: None, - SensorStateClass.MEASUREMENT: STATE_CLASS_MEASUREMENT, - SensorStateClass.TOTAL_INCREASING: STATE_CLASS_TOTAL_INCREASING, + EsphomeSensorStateClass.NONE: None, + EsphomeSensorStateClass.MEASUREMENT: SensorStateClass.MEASUREMENT, + EsphomeSensorStateClass.TOTAL_INCREASING: SensorStateClass.TOTAL_INCREASING, } ) @@ -85,7 +86,7 @@ class EsphomeSensor(EsphomeEntity[SensorInfo, SensorState], SensorEntity): return None if self._state.missing_state: return None - if self.device_class == DEVICE_CLASS_TIMESTAMP: + if self.device_class == SensorDeviceClass.TIMESTAMP: return dt.utc_from_timestamp(self._state.state) return f"{self._state.state:.{self._static_info.accuracy_decimals}f}" @@ -104,18 +105,18 @@ class EsphomeSensor(EsphomeEntity[SensorInfo, SensorState], SensorEntity): return self._static_info.device_class @property - def state_class(self) -> str | None: + def state_class(self) -> SensorStateClass | None: """Return the state class of this entity.""" if not self._static_info.state_class: return None state_class = self._static_info.state_class reset_type = self._static_info.last_reset_type if ( - state_class == SensorStateClass.MEASUREMENT + state_class == EsphomeSensorStateClass.MEASUREMENT and reset_type == LastResetType.AUTO ): # Legacy, last_reset_type auto was the equivalent to the TOTAL_INCREASING state class - return STATE_CLASS_TOTAL_INCREASING + return SensorStateClass.TOTAL_INCREASING return _STATE_CLASSES.from_esphome(self._static_info.state_class)