✨ Expose stable/numerical alert level entity

For people building automations
kamaradclimber · Nov 23, 2024 · 770813f · 770813f
1 parent d0597d5
commit 770813f
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 11 deletions.
diff --git a/custom_components/vigieau/__init__.py b/custom_components/vigieau/__init__.py
@@ -216,10 +216,14 @@ def __init__(
         coordinator: VigieauAPICoordinator,
         hass: HomeAssistant,
         config_entry: ConfigEntry,
+        numeric_state: bool,
     ):
         super().__init__(coordinator)
+        self._numeric_state = numeric_state
         self.hass = hass
         self._attr_name = f"Alert level in {config_entry.data.get(CONF_CITY)}"
+        if self._numeric_state:
+            self._attr_name += " (numeric)"
         self._attr_native_value = None
         self._attr_state_attributes = None
         if MIGRATED_FROM_VERSION_1 in config_entry.data:
@@ -228,6 +232,9 @@ def __init__(
             self._attr_unique_id = f"sensor-vigieau-{self._attr_name}-{config_entry.data.get(CONF_INSEE_CODE)}"
         else:
             self._attr_unique_id = f"sensor-vigieau-{self._attr_name}-{config_entry.data.get(CONF_INSEE_CODE)}-{config_entry.data.get(CONF_ZONE_TYPE)}"
+        if self._numeric_state:
+            self._attr_unique_id += "-numeric"
+            self._attr_entity_category = EntityCategory.DIAGNOSTIC
 
         self._attr_device_info = DeviceInfo(
             name=f"{NAME} {config_entry.data.get(CONF_CITY)} {zone_type_to_str(config_entry.data.get(CONF_ZONE_TYPE))}",
@@ -254,16 +261,21 @@ def _handle_coordinator_update(self) -> None:
         if not self.coordinator.last_update_success:
             _LOGGER.debug("Last coordinator failed, assuming state has not changed")
             return
-        self._attr_native_value = self.coordinator.data["niveauGravite"]
+        self.numeric_state_value = self.coordinator.data["_numeric_state_value"]
+
+        if self._numeric_state:
+            self._attr_native_value = self.numeric_state_value
+        else:
+            self._attr_native_value = self.coordinator.data["niveauGravite"]
+            if self.numeric_state_value == 0:
+                self._attr_native_value = "vigilance (pas de restriction)"
 
         self._attr_icon = {
-            "vigilance": "mdi:water-check",
-            "vigilance_(pas_de_restriction)": "mdi:water-check",
-            "alerte": "mdi:water-alert",
-            "alerte_renforcée": "mdi:water-remove",
-            "alerte_renforcee": "mdi:water-remove",
-            "crise": "mdi:water-off",
-        }[self._attr_native_value.lower().replace(" ", "_")]
+            0: "mdi:water-check",
+            1: "mdi:water-alert",
+            2: "mdi:water-remove",
+            3: "mdi:water-off",
+        }[self.numeric_state_value]
 
         self.enrich_attributes(self.coordinator.data, "cheminFichier", "source")
         self.enrich_attributes(

diff --git a/custom_components/vigieau/api.py b/custom_components/vigieau/api.py
@@ -123,13 +123,21 @@ async def get_data(
             and re.match("Aucune zone.+en vigueur", (await resp.json())["message"])
         ):
             _LOGGER.debug(f"Vigieau replied with no restriction, faking data")
-            data = {"niveauGravite": "vigilance (pas de restriction)", "usages": [], "arrete": {}}
+            data = {"niveauGravite": "vigilance", "usages": [], "arrete": {}}
         elif resp.status == 200 and (await resp.text()) == "":
             _LOGGER.debug(f"Vigieau replied with no data at all, faking data")
-            data = {"niveauGravite": "vigilance (pas de restriction)", "usages": [], "arrete": {}}
+            data = {"niveauGravite": "vigilance", "usages": [], "arrete": {}}
         elif resp.status in range(200, 300):
             data = await resp.json()
         else:
             raise VigieauAPIError(f"Failed fetching vigieau data", resp.text)
         _LOGGER.debug(f"Data fetched from vigieau: {data}")
+        # enriching with numeric state value
+        data["_numeric_state_value"] = {
+            "vigilance": 0,
+            "alerte": 1,
+            "alerte_renforcée": 2,
+            "alerte_renforcee": 2,
+            "crise": 3,
+        }[data["niveauGravite"]]
         return data
diff --git a/custom_components/vigieau/sensor.py b/custom_components/vigieau/sensor.py
@@ -23,7 +23,8 @@ async def async_setup_entry(
         )
         for sensor_description in SENSOR_DEFINITIONS
     ]
-    sensors.append(AlertLevelEntity(vigieau_coordinator, hass, entry))
+    sensors.append(AlertLevelEntity(vigieau_coordinator, hass, entry, numeric_state=False))
+    sensors.append(AlertLevelEntity(vigieau_coordinator, hass, entry, numeric_state=True))
 
     async_add_entities(sensors)
     await vigieau_coordinator.async_config_entry_first_refresh()