Add dataclass to Ondilo Ico (#116928)

homeassistant/components/ondilo_ico/coordinator.py

@@ -1,5 +1,6 @@
 """Define an object to coordinate fetching Ondilo ICO data."""
 
+from dataclasses import dataclass
 from datetime import timedelta
 import logging
 from typing import Any
@@ -15,7 +16,16 @@ from .api import OndiloClient
 _LOGGER = logging.getLogger(__name__)
 
 
-class OndiloIcoCoordinator(DataUpdateCoordinator[list[dict[str, Any]]]):
+@dataclass
+class OndiloIcoData:
+    """Class for storing the data."""
+
+    ico: dict[str, Any]
+    pool: dict[str, Any]
+    sensors: dict[str, Any]
+
+
+class OndiloIcoCoordinator(DataUpdateCoordinator[dict[str, OndiloIcoData]]):
     """Class to manage fetching Ondilo ICO data from API."""
 
     def __init__(self, hass: HomeAssistant, api: OndiloClient) -> None:
@@ -28,7 +38,7 @@ class OndiloIcoCoordinator(DataUpdateCoordinator[list[dict[str, Any]]]):
         )
         self.api = api
 
-    async def _async_update_data(self) -> list[dict[str, Any]]:
+    async def _async_update_data(self) -> dict[str, OndiloIcoData]:
         """Fetch data from API endpoint."""
         try:
             return await self.hass.async_add_executor_job(self._update_data)
@@ -37,9 +47,9 @@ class OndiloIcoCoordinator(DataUpdateCoordinator[list[dict[str, Any]]]):
             _LOGGER.exception("Error getting pools")
             raise UpdateFailed(f"Error communicating with API: {err}") from err
 
-    def _update_data(self) -> list[dict[str, Any]]:
+    def _update_data(self) -> dict[str, OndiloIcoData]:
         """Fetch data from API endpoint."""
-        res = []
+        res = {}
         pools = self.api.get_pools()
         _LOGGER.debug("Pools: %s", pools)
         for pool in pools:
@@ -54,12 +64,10 @@ class OndiloIcoCoordinator(DataUpdateCoordinator[list[dict[str, Any]]]):
             except OndiloError:
                 _LOGGER.exception("Error communicating with API for %s", pool["id"])
                 continue
-            res.append(
+            res[pool["id"]] = OndiloIcoData(
-                {
+                ico=ico,
-                    **pool,
+                pool=pool,
-                    "ICO": ico,
+                sensors={sensor["data_type"]: sensor["value"] for sensor in sensors},
-                    "sensors": sensors,
-                }
             )
         if not res:
             raise UpdateFailed("No data available")

homeassistant/components/ondilo_ico/sensor.py

@@ -18,10 +18,11 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.typing import StateType
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from .const import DOMAIN
-from .coordinator import OndiloIcoCoordinator
+from .coordinator import OndiloIcoCoordinator, OndiloIcoData
 
 SENSOR_TYPES: tuple[SensorEntityDescription, ...] = (
     SensorEntityDescription(
@@ -76,11 +77,10 @@ async def async_setup_entry(
     coordinator: OndiloIcoCoordinator = hass.data[DOMAIN][entry.entry_id]
 
     async_add_entities(
-        OndiloICO(coordinator, poolidx, description)
+        OndiloICO(coordinator, pool_id, description)
-        for poolidx, pool in enumerate(coordinator.data)
+        for pool_id, pool in coordinator.data.items()
-        for sensor in pool["sensors"]
         for description in SENSOR_TYPES
-        if description.key == sensor["data_type"]
+        if description.key in pool.sensors
     )
 
 
@@ -92,44 +92,31 @@ class OndiloICO(CoordinatorEntity[OndiloIcoCoordinator], SensorEntity):
     def __init__(
         self,
         coordinator: OndiloIcoCoordinator,
-        poolidx: int,
+        pool_id: str,
         description: SensorEntityDescription,
     ) -> None:
         """Initialize sensor entity with data from coordinator."""
         super().__init__(coordinator)
         self.entity_description = description
 
-        self._poolid = self.coordinator.data[poolidx]["id"]
+        self._pool_id = pool_id
 
-        pooldata = self._pooldata()
+        data = self.pool_data
-        self._attr_unique_id = f"{pooldata['ICO']['serial_number']}-{description.key}"
+        self._attr_unique_id = f"{data.ico['serial_number']}-{description.key}"
         self._attr_device_info = DeviceInfo(
-            identifiers={(DOMAIN, pooldata["ICO"]["serial_number"])},
+            identifiers={(DOMAIN, data.ico["serial_number"])},
             manufacturer="Ondilo",
             model="ICO",
-            name=pooldata["name"],
+            name=data.pool["name"],
-            sw_version=pooldata["ICO"]["sw_version"],
+            sw_version=data.ico["sw_version"],
-        )
-
-    def _pooldata(self):
-        """Get pool data dict."""
-        return next(
-            (pool for pool in self.coordinator.data if pool["id"] == self._poolid),
-            None,
-        )
-
-    def _devdata(self):
-        """Get device data dict."""
-        return next(
-            (
-                data_type
-                for data_type in self._pooldata()["sensors"]
-                if data_type["data_type"] == self.entity_description.key
-            ),
-            None,
         )
 
     @property
-    def native_value(self):
+    def pool_data(self) -> OndiloIcoData:
+        """Get pool data."""
+        return self.coordinator.data[self._pool_id]
+
+    @property
+    def native_value(self) -> StateType:
         """Last value of the sensor."""
-        return self._devdata()["value"]
+        return self.pool_data.sensors[self.entity_description.key]