Merge pull request #544 from stickpin/Add_EntityCategory

Reorganize entities by categories
robinostlund · Jan 17, 2024 · 7daa4fe · 7daa4fe
2 parents ff62871 + 4efaee4
commit 7daa4fe
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 0 deletions.
diff --git a/custom_components/volkswagencarnet/binary_sensor.py b/custom_components/volkswagencarnet/binary_sensor.py
@@ -5,6 +5,7 @@
 from homeassistant.components.binary_sensor import DEVICE_CLASSES, BinarySensorEntity, BinarySensorDeviceClass
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity import EntityCategory
 
 from . import VolkswagenEntity
 from .const import DATA_KEY, DATA, DOMAIN, UPDATE_CALLBACK
@@ -54,3 +55,11 @@ def device_class(self) -> Union[BinarySensorDeviceClass, str, None]:
             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."""
+        if self.instrument.entity_type == "diag":
+            return EntityCategory.DIAGNOSTIC
+        if self.instrument.entity_type == "config":
+            return EntityCategory.CONFIG
diff --git a/custom_components/volkswagencarnet/sensor.py b/custom_components/volkswagencarnet/sensor.py
@@ -11,6 +11,7 @@
 )
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity import EntityCategory
 
 from . import VolkswagenEntity
 from .const import DATA_KEY, DATA, DOMAIN, UPDATE_CALLBACK
@@ -79,3 +80,11 @@ def state_class(self) -> Union[SensorStateClass, None]:
             return self.instrument.state_class
         _LOGGER.warning(f"Unknown state class {self.instrument.state_class}")
         return None
+
+    @property
+    def entity_category(self) -> Union[EntityCategory, str, None]:
+        """Return entity category."""
+        if self.instrument.entity_type == "diag":
+            return EntityCategory.DIAGNOSTIC
+        if self.instrument.entity_type == "config":
+            return EntityCategory.CONFIG
diff --git a/custom_components/volkswagencarnet/switch.py b/custom_components/volkswagencarnet/switch.py
@@ -92,6 +92,14 @@ async def async_turn_off(self, **kwargs):
         await self.instrument.turn_off()
         self.notify_updated()
 
+    @property
+    def entity_category(self) -> Union[EntityCategory, str, None]:
+        """Return entity category."""
+        if self.instrument.entity_type == "diag":
+            return EntityCategory.DIAGNOSTIC
+        if self.instrument.entity_type == "config":
+            return EntityCategory.CONFIG
+
     @property
     def assumed_state(self):
         """Return state assumption."""