Improve device class handling in ESPHome (#82923)

homeassistant/components/esphome/binary_sensor.py

@@ -1,9 +1,14 @@
 """Support for ESPHome binary sensors."""
 from __future__ import annotations
 
+from contextlib import suppress
+
 from aioesphomeapi import BinarySensorInfo, BinarySensorState
 
-from homeassistant.components.binary_sensor import DEVICE_CLASSES, BinarySensorEntity
+from homeassistant.components.binary_sensor import (
+    BinarySensorDeviceClass,
+    BinarySensorEntity,
+)
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
@@ -45,11 +50,11 @@ class EsphomeBinarySensor(
         return self._state.state
 
     @property
-    def device_class(self) -> str | None:
+    def device_class(self) -> BinarySensorDeviceClass | None:
         """Return the class of this device, from component DEVICE_CLASSES."""
-        if self._static_info.device_class not in DEVICE_CLASSES:
-            return None
-        return self._static_info.device_class
+        with suppress(ValueError):
+            return BinarySensorDeviceClass(self._static_info.device_class)
+        return None
 
     @property
     def available(self) -> bool:

homeassistant/components/esphome/sensor.py

@@ -1,6 +1,7 @@
 """Support for esphome sensors."""
 from __future__ import annotations
 
+from contextlib import suppress
 from datetime import datetime
 import math
 
@@ -14,7 +15,6 @@ from aioesphomeapi import (
 from aioesphomeapi.model import LastResetType
 
 from homeassistant.components.sensor import (
-    DEVICE_CLASSES,
     SensorDeviceClass,
     SensorEntity,
     SensorStateClass,
@@ -96,11 +96,11 @@ class EsphomeSensor(EsphomeEntity[SensorInfo, SensorState], SensorEntity):
         return self._static_info.unit_of_measurement
 
     @property
-    def device_class(self) -> str | None:
+    def device_class(self) -> SensorDeviceClass | None:
         """Return the class of this device, from component DEVICE_CLASSES."""
-        if self._static_info.device_class not in DEVICE_CLASSES:
-            return None
-        return self._static_info.device_class
+        with suppress(ValueError):
+            return SensorDeviceClass(self._static_info.device_class)
+        return None
 
     @property
     def state_class(self) -> SensorStateClass | None:

homeassistant/components/esphome/switch.py

@@ -1,11 +1,12 @@
 """Support for ESPHome switches."""
 from __future__ import annotations
 
+from contextlib import suppress
 from typing import Any
 
 from aioesphomeapi import SwitchInfo, SwitchState
 
-from homeassistant.components.switch import DEVICE_CLASSES, SwitchEntity
+from homeassistant.components.switch import SwitchDeviceClass, SwitchEntity
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
@@ -43,11 +44,11 @@ class EsphomeSwitch(EsphomeEntity[SwitchInfo, SwitchState], SwitchEntity):
         return self._state.state
 
     @property
-    def device_class(self) -> str | None:
+    def device_class(self) -> SwitchDeviceClass | None:
         """Return the class of this device."""
-        if self._static_info.device_class not in DEVICE_CLASSES:
-            return None
-        return self._static_info.device_class
+        with suppress(ValueError):
+            return SwitchDeviceClass(self._static_info.device_class)
+        return None
 
     async def async_turn_on(self, **kwargs: Any) -> None:
         """Turn the entity on."""