Use new enums in esphome (#61391)

Co-authored-by: epenet <epenet@users.noreply.github.com>

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

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)