Merge coordinators in Airgradient

homeassistant/components/airgradient/__init__.py

@@ -2,18 +2,14 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
-
-from airgradient import AirGradientClient, get_model_name
+from airgradient import AirGradientClient
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
-from .const import DOMAIN
-from .coordinator import AirGradientConfigCoordinator, AirGradientMeasurementCoordinator
+from .coordinator import AirGradientCoordinator
 
 PLATFORMS: list[Platform] = [
     Platform.BUTTON,
@@ -25,15 +21,7 @@
 ]
 
 
-@dataclass
-class AirGradientData:
-    """AirGradient data class."""
-
-    measurement: AirGradientMeasurementCoordinator
-    config: AirGradientConfigCoordinator
-
-
-type AirGradientConfigEntry = ConfigEntry[AirGradientData]
+type AirGradientConfigEntry = ConfigEntry[AirGradientCoordinator]
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: AirGradientConfigEntry) -> bool:
@@ -43,27 +31,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: AirGradientConfigEntry)
         entry.data[CONF_HOST], session=async_get_clientsession(hass)
     )
 
-    measurement_coordinator = AirGradientMeasurementCoordinator(hass, client)
-    config_coordinator = AirGradientConfigCoordinator(hass, client)
-
-    await measurement_coordinator.async_config_entry_first_refresh()
-    await config_coordinator.async_config_entry_first_refresh()
-
-    device_registry = dr.async_get(hass)
-    device_registry.async_get_or_create(
-        config_entry_id=entry.entry_id,
-        identifiers={(DOMAIN, measurement_coordinator.serial_number)},
-        manufacturer="AirGradient",
-        model=get_model_name(measurement_coordinator.data.model),
-        model_id=measurement_coordinator.data.model,
-        serial_number=measurement_coordinator.data.serial_number,
-        sw_version=measurement_coordinator.data.firmware_version,
-    )
+    coordinator = AirGradientCoordinator(hass, client)
 
-    entry.runtime_data = AirGradientData(
-        measurement=measurement_coordinator,
-        config=config_coordinator,
-    )
+    await coordinator.async_config_entry_first_refresh()
+
+    entry.runtime_data = coordinator
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 

homeassistant/components/airgradient/button.py

@@ -15,8 +15,9 @@
 from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import DOMAIN, AirGradientConfigEntry
-from .coordinator import AirGradientConfigCoordinator
+from . import AirGradientConfigEntry
+from .const import DOMAIN
+from .coordinator import AirGradientCoordinator
 from .entity import AirGradientEntity
 
 
@@ -47,8 +48,8 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up AirGradient button entities based on a config entry."""
-    model = entry.runtime_data.measurement.data.model
-    coordinator = entry.runtime_data.config
+    coordinator = entry.runtime_data
+    model = coordinator.data.measures.model
 
     added_entities = False
 
@@ -57,7 +58,7 @@ def _check_entities() -> None:
         nonlocal added_entities
 
         if (
-            coordinator.data.configuration_control is ConfigurationControl.LOCAL
+            coordinator.data.config.configuration_control is ConfigurationControl.LOCAL
             and not added_entities
         ):
             entities = [AirGradientButton(coordinator, CO2_CALIBRATION)]
@@ -67,7 +68,8 @@ def _check_entities() -> None:
             async_add_entities(entities)
             added_entities = True
         elif (
-            coordinator.data.configuration_control is not ConfigurationControl.LOCAL
+            coordinator.data.config.configuration_control
+            is not ConfigurationControl.LOCAL
             and added_entities
         ):
             entity_registry = er.async_get(hass)
@@ -87,11 +89,10 @@ class AirGradientButton(AirGradientEntity, ButtonEntity):
     """Defines an AirGradient button."""
 
     entity_description: AirGradientButtonEntityDescription
-    coordinator: AirGradientConfigCoordinator
 
     def __init__(
         self,
-        coordinator: AirGradientConfigCoordinator,
+        coordinator: AirGradientCoordinator,
         description: AirGradientButtonEntityDescription,
     ) -> None:
         """Initialize airgradient button."""

homeassistant/components/airgradient/coordinator.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from dataclasses import dataclass
 from datetime import timedelta
 from typing import TYPE_CHECKING
 
@@ -16,7 +17,15 @@
     from . import AirGradientConfigEntry
 
 
-class AirGradientCoordinator[_DataT](DataUpdateCoordinator[_DataT]):
+@dataclass
+class AirGradientData:
+    """Class for AirGradient data."""
+
+    measures: Measures
+    config: Config
+
+
+class AirGradientCoordinator(DataUpdateCoordinator[AirGradientData]):
     """Class to manage fetching AirGradient data."""
 
     config_entry: AirGradientConfigEntry
@@ -33,25 +42,11 @@ def __init__(self, hass: HomeAssistant, client: AirGradientClient) -> None:
         assert self.config_entry.unique_id
         self.serial_number = self.config_entry.unique_id
 
-    async def _async_update_data(self) -> _DataT:
+    async def _async_update_data(self) -> AirGradientData:
         try:
-            return await self._update_data()
+            measures = await self.client.get_current_measures()
+            config = await self.client.get_config()
         except AirGradientError as error:
             raise UpdateFailed(error) from error
-
-    async def _update_data(self) -> _DataT:
-        raise NotImplementedError
-
-
-class AirGradientMeasurementCoordinator(AirGradientCoordinator[Measures]):
-    """Class to manage fetching AirGradient data."""
-
-    async def _update_data(self) -> Measures:
-        return await self.client.get_current_measures()
-
-
-class AirGradientConfigCoordinator(AirGradientCoordinator[Config]):
-    """Class to manage fetching AirGradient data."""
-
-    async def _update_data(self) -> Config:
-        return await self.client.get_config()
+        else:
+            return AirGradientData(measures, config)

homeassistant/components/airgradient/entity.py

@@ -1,5 +1,7 @@
 """Base class for AirGradient entities."""
 
+from airgradient import get_model_name
+
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
@@ -15,6 +17,12 @@ class AirGradientEntity(CoordinatorEntity[AirGradientCoordinator]):
     def __init__(self, coordinator: AirGradientCoordinator) -> None:
         """Initialize airgradient entity."""
         super().__init__(coordinator)
+        measures = coordinator.data.measures
         self._attr_device_info = DeviceInfo(
             identifiers={(DOMAIN, coordinator.serial_number)},
+            manufacturer="AirGradient",
+            model=get_model_name(measures.model),
+            model_id=measures.model,
+            serial_number=coordinator.serial_number,
+            sw_version=measures.firmware_version,
         )

homeassistant/components/airgradient/number.py

@@ -18,7 +18,7 @@
 
 from . import AirGradientConfigEntry
 from .const import DOMAIN
-from .coordinator import AirGradientConfigCoordinator
+from .coordinator import AirGradientCoordinator
 from .entity import AirGradientEntity
 
 
@@ -62,8 +62,8 @@ async def async_setup_entry(
 ) -> None:
     """Set up AirGradient number entities based on a config entry."""
 
-    model = entry.runtime_data.measurement.data.model
-    coordinator = entry.runtime_data.config
+    coordinator = entry.runtime_data
+    model = coordinator.data.measures.model
 
     added_entities = False
 
@@ -72,7 +72,7 @@ def _async_check_entities() -> None:
         nonlocal added_entities
 
         if (
-            coordinator.data.configuration_control is ConfigurationControl.LOCAL
+            coordinator.data.config.configuration_control is ConfigurationControl.LOCAL
             and not added_entities
         ):
             entities = []
@@ -84,7 +84,8 @@ def _async_check_entities() -> None:
             async_add_entities(entities)
             added_entities = True
         elif (
-            coordinator.data.configuration_control is not ConfigurationControl.LOCAL
+            coordinator.data.config.configuration_control
+            is not ConfigurationControl.LOCAL
             and added_entities
         ):
             entity_registry = er.async_get(hass)
@@ -104,11 +105,10 @@ class AirGradientNumber(AirGradientEntity, NumberEntity):
     """Defines an AirGradient number entity."""
 
     entity_description: AirGradientNumberEntityDescription
-    coordinator: AirGradientConfigCoordinator
 
     def __init__(
         self,
-        coordinator: AirGradientConfigCoordinator,
+        coordinator: AirGradientCoordinator,
         description: AirGradientNumberEntityDescription,
     ) -> None:
         """Initialize AirGradient number."""
@@ -119,7 +119,7 @@ def __init__(
     @property
     def native_value(self) -> int | None:
         """Return the state of the number."""
-        return self.entity_description.value_fn(self.coordinator.data)
+        return self.entity_description.value_fn(self.coordinator.data.config)
 
     async def async_set_native_value(self, value: float) -> None:
         """Set the selected value."""

homeassistant/components/airgradient/select.py

@@ -18,7 +18,7 @@
 
 from . import AirGradientConfigEntry
 from .const import DOMAIN, PM_STANDARD, PM_STANDARD_REVERSE
-from .coordinator import AirGradientConfigCoordinator
+from .coordinator import AirGradientCoordinator
 from .entity import AirGradientEntity
 
 
@@ -144,21 +144,19 @@ async def async_setup_entry(
 ) -> None:
     """Set up AirGradient select entities based on a config entry."""
 
-    coordinator = entry.runtime_data.config
-    measurement_coordinator = entry.runtime_data.measurement
+    coordinator = entry.runtime_data
+    model = coordinator.data.measures.model
 
     async_add_entities([AirGradientSelect(coordinator, CONFIG_CONTROL_ENTITY)])
 
-    model = measurement_coordinator.data.model
-
     added_entities = False
 
     @callback
     def _async_check_entities() -> None:
         nonlocal added_entities
 
         if (
-            coordinator.data.configuration_control is ConfigurationControl.LOCAL
+            coordinator.data.config.configuration_control is ConfigurationControl.LOCAL
             and not added_entities
         ):
             entities: list[AirGradientSelect] = [
@@ -179,7 +177,8 @@ def _async_check_entities() -> None:
             async_add_entities(entities)
             added_entities = True
         elif (
-            coordinator.data.configuration_control is not ConfigurationControl.LOCAL
+            coordinator.data.config.configuration_control
+            is not ConfigurationControl.LOCAL
             and added_entities
         ):
             entity_registry = er.async_get(hass)
@@ -201,11 +200,10 @@ class AirGradientSelect(AirGradientEntity, SelectEntity):
     """Defines an AirGradient select entity."""
 
     entity_description: AirGradientSelectEntityDescription
-    coordinator: AirGradientConfigCoordinator
 
     def __init__(
         self,
-        coordinator: AirGradientConfigCoordinator,
+        coordinator: AirGradientCoordinator,
         description: AirGradientSelectEntityDescription,
     ) -> None:
         """Initialize AirGradient select."""
@@ -216,7 +214,7 @@ def __init__(
     @property
     def current_option(self) -> str | None:
         """Return the state of the select."""
-        return self.entity_description.value_fn(self.coordinator.data)
+        return self.entity_description.value_fn(self.coordinator.data.config)
 
     async def async_select_option(self, option: str) -> None:
         """Change the selected option."""