Add DeviceClass to NumberEntity

robinostlund · Mar 1, 2024 · 200672f · 200672f
1 parent 1cd6c1c
commit 200672f
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/custom_components/volkswagencarnet/number.py b/custom_components/volkswagencarnet/number.py
@@ -3,7 +3,10 @@
 import logging
 from typing import Union
 
-from homeassistant.components.number import NumberEntity
+from homeassistant.components.number import (
+    NumberEntity,
+    NumberDeviceClass,
+)
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity import EntityCategory
 
@@ -74,6 +77,14 @@ def native_value(self) -> float:
             return self.instrument.state
         return None
 
+    @property
+    def device_class(self) -> Union[NumberDeviceClass, None]:
+        """Return the device class."""
+        if self.instrument.device_class is None or self.instrument.device_class in NumberDeviceClass:
+            return self.instrument.device_class
+        _LOGGER.warning(f"Unknown device class {self.instrument.device_class}")
+        return None
+
     @property
     def entity_category(self) -> Union[EntityCategory, str, None]:
         """Return entity category."""