Use dataclass for Withings domain data (#102547)

Co-authored-by: J. Nick Koston <nick@koston.org>

homeassistant/components/withings/__init__.py

@@ -7,6 +7,7 @@ from __future__ import annotations
 import asyncio
 from collections.abc import Awaitable, Callable
 import contextlib
+from dataclasses import dataclass, field
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any
 
@@ -50,17 +51,7 @@ from homeassistant.helpers.event import async_call_later
 from homeassistant.helpers.issue_registry import IssueSeverity, async_create_issue
 from homeassistant.helpers.typing import ConfigType
 
-from .const import (
+from .const import CONF_PROFILES, CONF_USE_WEBHOOK, DEFAULT_TITLE, DOMAIN, LOGGER
-    BED_PRESENCE_COORDINATOR,
-    CONF_PROFILES,
-    CONF_USE_WEBHOOK,
-    DEFAULT_TITLE,
-    DOMAIN,
-    GOALS_COORDINATOR,
-    LOGGER,
-    MEASUREMENT_COORDINATOR,
-    SLEEP_COORDINATOR,
-)
 from .coordinator import (
     WithingsBedPresenceDataUpdateCoordinator,
     WithingsDataUpdateCoordinator,
@@ -132,6 +123,26 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     return True
 
 
+@dataclass(slots=True)
+class WithingsData:
+    """Dataclass to hold withings domain data."""
+
+    measurement_coordinator: WithingsMeasurementDataUpdateCoordinator
+    sleep_coordinator: WithingsSleepDataUpdateCoordinator
+    bed_presence_coordinator: WithingsBedPresenceDataUpdateCoordinator
+    goals_coordinator: WithingsGoalsDataUpdateCoordinator
+    coordinators: set[WithingsDataUpdateCoordinator] = field(default_factory=set)
+
+    def __post_init__(self) -> None:
+        """Collect all coordinators in a set."""
+        self.coordinators = {
+            self.measurement_coordinator,
+            self.sleep_coordinator,
+            self.bed_presence_coordinator,
+            self.goals_coordinator,
+        }
+
+
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Withings from a config entry."""
     if CONF_WEBHOOK_ID not in entry.data or entry.unique_id is None:
@@ -156,19 +167,17 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         return token
 
     client.refresh_token_function = _refresh_token
-    coordinators: dict[str, WithingsDataUpdateCoordinator] = {
+    withings_data = WithingsData(
-        MEASUREMENT_COORDINATOR: WithingsMeasurementDataUpdateCoordinator(hass, client),
+        measurement_coordinator=WithingsMeasurementDataUpdateCoordinator(hass, client),
-        SLEEP_COORDINATOR: WithingsSleepDataUpdateCoordinator(hass, client),
+        sleep_coordinator=WithingsSleepDataUpdateCoordinator(hass, client),
-        BED_PRESENCE_COORDINATOR: WithingsBedPresenceDataUpdateCoordinator(
+        bed_presence_coordinator=WithingsBedPresenceDataUpdateCoordinator(hass, client),
-            hass, client
+        goals_coordinator=WithingsGoalsDataUpdateCoordinator(hass, client),
-        ),
+    )
-        GOALS_COORDINATOR: WithingsGoalsDataUpdateCoordinator(hass, client),
-    }
 
-    for coordinator in coordinators.values():
+    for coordinator in withings_data.coordinators:
         await coordinator.async_config_entry_first_refresh()
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinators
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = withings_data
 
     async def unregister_webhook(
         _: Any,
@@ -176,7 +185,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         LOGGER.debug("Unregister Withings webhook (%s)", entry.data[CONF_WEBHOOK_ID])
         webhook_unregister(hass, entry.data[CONF_WEBHOOK_ID])
         await async_unsubscribe_webhooks(client)
-        for coordinator in coordinators.values():
+        for coordinator in withings_data.coordinators:
             coordinator.webhook_subscription_listener(False)
 
     async def register_webhook(
@@ -203,12 +212,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             DOMAIN,
             webhook_name,
             entry.data[CONF_WEBHOOK_ID],
-            get_webhook_handler(coordinators),
+            get_webhook_handler(withings_data),
             allowed_methods=[METH_POST],
         )
 
         await async_subscribe_webhooks(client, webhook_url)
-        for coordinator in coordinators.values():
+        for coordinator in withings_data.coordinators:
             coordinator.webhook_subscription_listener(True)
         LOGGER.debug("Register Withings webhook: %s", webhook_url)
         entry.async_on_unload(
@@ -325,7 +334,7 @@ def json_message_response(message: str, message_code: int) -> Response:
 
 
 def get_webhook_handler(
-    coordinators: dict[str, WithingsDataUpdateCoordinator],
+    withings_data: WithingsData,
 ) -> Callable[[HomeAssistant, str, Request], Awaitable[Response | None]]:
     """Return webhook handler."""
 
@@ -349,7 +358,7 @@ def get_webhook_handler(
             NotificationCategory.UNKNOWN,
         )
 
-        for coordinator in coordinators.values():
+        for coordinator in withings_data.coordinators:
             if notification_category in coordinator.notification_categories:
                 await coordinator.async_webhook_data_updated(notification_category)
 

homeassistant/components/withings/binary_sensor.py

@@ -9,7 +9,7 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import BED_PRESENCE_COORDINATOR, DOMAIN
+from .const import DOMAIN
 from .coordinator import WithingsBedPresenceDataUpdateCoordinator
 from .entity import WithingsEntity
 
@@ -20,9 +20,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the sensor config entry."""
-    coordinator: WithingsBedPresenceDataUpdateCoordinator = hass.data[DOMAIN][
+    coordinator = hass.data[DOMAIN][entry.entry_id].bed_presence_coordinator
-        entry.entry_id
-    ][BED_PRESENCE_COORDINATOR]
 
     entities = [WithingsBinarySensor(coordinator)]
 

homeassistant/components/withings/const.py

@@ -1,25 +1,13 @@
 """Constants used by the Withings component."""
 import logging
 
+LOGGER = logging.getLogger(__package__)
+
 DEFAULT_TITLE = "Withings"
 CONF_PROFILES = "profiles"
 CONF_USE_WEBHOOK = "use_webhook"
 
-DATA_MANAGER = "data_manager"
-
-CONFIG = "config"
 DOMAIN = "withings"
-LOG_NAMESPACE = "homeassistant.components.withings"
-PROFILE = "profile"
-PUSH_HANDLER = "push_handler"
-
-MEASUREMENT_COORDINATOR = "measurement_coordinator"
-SLEEP_COORDINATOR = "sleep_coordinator"
-BED_PRESENCE_COORDINATOR = "bed_presence_coordinator"
-GOALS_COORDINATOR = "goals_coordinator"
-
-LOGGER = logging.getLogger(__package__)
-
 
 SCORE_POINTS = "points"
 UOM_BEATS_PER_MINUTE = "bpm"

homeassistant/components/withings/diagnostics.py

@@ -10,12 +10,8 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_WEBHOOK_ID
 from homeassistant.core import HomeAssistant
 
-from . import (
+from . import CONF_CLOUDHOOK_URL, WithingsData
-    CONF_CLOUDHOOK_URL,
+from .const import DOMAIN
-    WithingsMeasurementDataUpdateCoordinator,
-    WithingsSleepDataUpdateCoordinator,
-)
-from .const import DOMAIN, MEASUREMENT_COORDINATOR, SLEEP_COORDINATOR
 
 
 async def async_get_config_entry_diagnostics(
@@ -29,17 +25,12 @@ async def async_get_config_entry_diagnostics(
 
     has_cloudhooks = CONF_CLOUDHOOK_URL in entry.data
 
-    measurement_coordinator: WithingsMeasurementDataUpdateCoordinator = hass.data[
+    withings_data: WithingsData = hass.data[DOMAIN][entry.entry_id]
-        DOMAIN
-    ][entry.entry_id][MEASUREMENT_COORDINATOR]
-    sleep_coordinator: WithingsSleepDataUpdateCoordinator = hass.data[DOMAIN][
-        entry.entry_id
-    ][SLEEP_COORDINATOR]
 
     return {
         "has_valid_external_webhook_url": has_valid_external_webhook_url,
         "has_cloudhooks": has_cloudhooks,
-        "webhooks_connected": measurement_coordinator.webhooks_connected,
+        "webhooks_connected": withings_data.measurement_coordinator.webhooks_connected,
-        "received_measurements": list(measurement_coordinator.data),
+        "received_measurements": list(withings_data.measurement_coordinator.data),
-        "received_sleep_data": sleep_coordinator.data is not None,
+        "received_sleep_data": withings_data.sleep_coordinator.data is not None,
     }

homeassistant/components/withings/sensor.py

@@ -25,12 +25,10 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
 
+from . import WithingsData
 from .const import (
     DOMAIN,
-    GOALS_COORDINATOR,
-    MEASUREMENT_COORDINATOR,
     SCORE_POINTS,
-    SLEEP_COORDINATOR,
     UOM_BEATS_PER_MINUTE,
     UOM_BREATHS_PER_MINUTE,
     UOM_FREQUENCY,
@@ -462,9 +460,9 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the sensor config entry."""
-    measurement_coordinator: WithingsMeasurementDataUpdateCoordinator = hass.data[
+    withings_data: WithingsData = hass.data[DOMAIN][entry.entry_id]
-        DOMAIN
+
-    ][entry.entry_id][MEASUREMENT_COORDINATOR]
+    measurement_coordinator = withings_data.measurement_coordinator
 
     entities: list[SensorEntity] = []
     entities.extend(
@@ -492,9 +490,7 @@ async def async_setup_entry(
 
     measurement_coordinator.async_add_listener(_async_measurement_listener)
 
-    goals_coordinator: WithingsGoalsDataUpdateCoordinator = hass.data[DOMAIN][
+    goals_coordinator = withings_data.goals_coordinator
-        entry.entry_id
-    ][GOALS_COORDINATOR]
 
     current_goals = get_current_goals(goals_coordinator.data)
 
@@ -516,9 +512,7 @@ async def async_setup_entry(
 
     goals_coordinator.async_add_listener(_async_goals_listener)
 
-    sleep_coordinator: WithingsSleepDataUpdateCoordinator = hass.data[DOMAIN][
+    sleep_coordinator = withings_data.sleep_coordinator
-        entry.entry_id
-    ][SLEEP_COORDINATOR]
 
     entities.extend(
         WithingsSleepSensor(sleep_coordinator, attribute) for attribute in SLEEP_SENSORS