Store runtime data in entry in Withings (#116439)

* Add entry runtime data to Withings

* Store runtime data in entry in Withings

* Fix

* Fix

* Update homeassistant/components/withings/coordinator.py

Co-authored-by: Marc Mueller <30130371+cdce8p@users.noreply.github.com>

---------

Co-authored-by: Marc Mueller <30130371+cdce8p@users.noreply.github.com>

homeassistant/components/withings/__init__.py

@@ -10,7 +10,7 @@ from collections.abc import Awaitable, Callable
 import contextlib
 from dataclasses import dataclass, field
 from datetime import timedelta
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any
 
 from aiohttp import ClientError
 from aiohttp.hdrs import METH_POST
@@ -59,6 +59,7 @@ PLATFORMS = [Platform.BINARY_SENSOR, Platform.CALENDAR, Platform.SENSOR]
 SUBSCRIBE_DELAY = timedelta(seconds=5)
 UNSUBSCRIBE_DELAY = timedelta(seconds=1)
 CONF_CLOUDHOOK_URL = "cloudhook_url"
+WithingsConfigEntry = ConfigEntry["WithingsData"]
 
 
 @dataclass(slots=True)
@@ -86,7 +87,7 @@ class WithingsData:
         }
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: WithingsConfigEntry) -> bool:
     """Set up Withings from a config entry."""
     if CONF_WEBHOOK_ID not in entry.data or entry.unique_id is None:
         new_data = entry.data.copy()
@@ -126,7 +127,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     for coordinator in withings_data.coordinators:
         await coordinator.async_config_entry_first_refresh()
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = withings_data
+    entry.runtime_data = withings_data
 
     webhook_manager = WithingsWebhookManager(hass, entry)
 
@@ -159,13 +160,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: WithingsConfigEntry) -> bool:
     """Unload Withings config entry."""
     webhook_unregister(hass, entry.data[CONF_WEBHOOK_ID])
 
-    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-        hass.data[DOMAIN].pop(entry.entry_id)
-    return unload_ok
 
 
 async def async_subscribe_webhooks(client: WithingsClient, webhook_url: str) -> None:
@@ -200,7 +199,7 @@ class WithingsWebhookManager:
     _webhooks_registered = False
     _register_lock = asyncio.Lock()
 
-    def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None:
+    def __init__(self, hass: HomeAssistant, entry: WithingsConfigEntry) -> None:
         """Initialize webhook manager."""
         self.hass = hass
         self.entry = entry
@@ -208,7 +207,7 @@ class WithingsWebhookManager:
     @property
     def withings_data(self) -> WithingsData:
         """Return Withings data."""
-        return cast(WithingsData, self.hass.data[DOMAIN][self.entry.entry_id])
+        return self.entry.runtime_data
 
     async def unregister_webhook(
         self,
@@ -297,7 +296,9 @@ async def async_unsubscribe_webhooks(client: WithingsClient) -> None:
         )
 
 
-async def _async_cloudhook_generate_url(hass: HomeAssistant, entry: ConfigEntry) -> str:
+async def _async_cloudhook_generate_url(
+    hass: HomeAssistant, entry: WithingsConfigEntry
+) -> str:
     """Generate the full URL for a webhook_id."""
     if CONF_CLOUDHOOK_URL not in entry.data:
         webhook_id = entry.data[CONF_WEBHOOK_ID]
@@ -312,7 +313,7 @@ async def _async_cloudhook_generate_url(hass: HomeAssistant, entry: ConfigEntry)
     return str(entry.data[CONF_CLOUDHOOK_URL])
 
 
-async def async_remove_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:
+async def async_remove_entry(hass: HomeAssistant, entry: WithingsConfigEntry) -> None:
     """Cleanup when entry is removed."""
     if cloud.async_active_subscription(hass):
         try:

homeassistant/components/withings/binary_sensor.py

@@ -8,12 +8,12 @@ from homeassistant.components.binary_sensor import (
     BinarySensorDeviceClass,
     BinarySensorEntity,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 import homeassistant.helpers.entity_registry as er
 
+from . import WithingsConfigEntry
 from .const import DOMAIN
 from .coordinator import WithingsBedPresenceDataUpdateCoordinator
 from .entity import WithingsEntity
@@ -21,11 +21,11 @@ from .entity import WithingsEntity
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: WithingsConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the sensor config entry."""
-    coordinator = hass.data[DOMAIN][entry.entry_id].bed_presence_coordinator
+    coordinator = entry.runtime_data.bed_presence_coordinator
 
     ent_reg = er.async_get(hass)
 

homeassistant/components/withings/calendar.py

@@ -8,25 +8,24 @@ from datetime import datetime
 from aiowithings import WithingsClient, WorkoutCategory
 
 from homeassistant.components.calendar import CalendarEntity, CalendarEvent
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 import homeassistant.helpers.entity_registry as er
 
-from . import DOMAIN, WithingsData
+from . import DOMAIN, WithingsConfigEntry
 from .coordinator import WithingsWorkoutDataUpdateCoordinator
 from .entity import WithingsEntity
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: WithingsConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the calendar platform for entity."""
     ent_reg = er.async_get(hass)
-    withings_data: WithingsData = hass.data[DOMAIN][entry.entry_id]
+    withings_data = entry.runtime_data
 
     workout_coordinator = withings_data.workout_coordinator
 

homeassistant/components/withings/coordinator.py

@@ -1,8 +1,10 @@
 """Withings coordinator."""
 
+from __future__ import annotations
+
 from abc import abstractmethod
 from datetime import date, datetime, timedelta
-from typing import TypeVar
+from typing import TYPE_CHECKING, TypeVar
 
 from aiowithings import (
     Activity,
@@ -18,7 +20,6 @@ from aiowithings import (
     aggregate_measurements,
 )
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
@@ -26,6 +27,9 @@ from homeassistant.util import dt as dt_util
 
 from .const import LOGGER
 
+if TYPE_CHECKING:
+    from . import WithingsConfigEntry
+
 _T = TypeVar("_T")
 
 UPDATE_INTERVAL = timedelta(minutes=10)
@@ -34,7 +38,7 @@ UPDATE_INTERVAL = timedelta(minutes=10)
 class WithingsDataUpdateCoordinator(DataUpdateCoordinator[_T]):
     """Base coordinator."""
 
-    config_entry: ConfigEntry
+    config_entry: WithingsConfigEntry
     _default_update_interval: timedelta | None = UPDATE_INTERVAL
     _last_valid_update: datetime | None = None
     webhooks_connected: bool = False

homeassistant/components/withings/diagnostics.py

@@ -7,16 +7,14 @@ from typing import Any
 from yarl import URL
 
 from homeassistant.components.webhook import async_generate_url as webhook_generate_url
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_WEBHOOK_ID
 from homeassistant.core import HomeAssistant
 
-from . import CONF_CLOUDHOOK_URL, WithingsData
+from . import CONF_CLOUDHOOK_URL, WithingsConfigEntry
-from .const import DOMAIN
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: WithingsConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
 
@@ -26,7 +24,7 @@ async def async_get_config_entry_diagnostics(
 
     has_cloudhooks = CONF_CLOUDHOOK_URL in entry.data
 
-    withings_data: WithingsData = hass.data[DOMAIN][entry.entry_id]
+    withings_data = entry.runtime_data
 
     return {
         "has_valid_external_webhook_url": has_valid_external_webhook_url,

homeassistant/components/withings/sensor.py

@@ -22,7 +22,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     PERCENTAGE,
     Platform,
@@ -38,7 +37,7 @@ import homeassistant.helpers.entity_registry as er
 from homeassistant.helpers.typing import StateType
 from homeassistant.util import dt as dt_util
 
-from . import WithingsData
+from . import WithingsConfigEntry
 from .const import (
     DOMAIN,
     LOGGER,
@@ -619,13 +618,13 @@ def get_current_goals(goals: Goals) -> set[str]:
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: WithingsConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the sensor config entry."""
     ent_reg = er.async_get(hass)
 
-    withings_data: WithingsData = hass.data[DOMAIN][entry.entry_id]
+    withings_data = entry.runtime_data
 
     measurement_coordinator = withings_data.measurement_coordinator