Better type entry data

iMicknl · Jan 7, 2024 · 0bb1f4c · 0bb1f4c
1 parent 50fcc36
commit 0bb1f4c
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 20 deletions.
diff --git a/custom_components/sagemcom_fast/__init__.py b/custom_components/sagemcom_fast/__init__.py
@@ -1,7 +1,7 @@
 """The Sagemcom F@st integration."""
 from __future__ import annotations
 
-import asyncio
+from dataclasses import dataclass
 from datetime import timedelta
 
 from aiohttp.client_exceptions import ClientError
@@ -26,6 +26,7 @@
     MaximumSessionCountException,
     UnauthorizedException,
 )
+from sagemcom_api.models import DeviceInfo as GatewayDeviceInfo
 
 from .const import (
     CONF_ENCRYPTION_METHOD,
@@ -39,8 +40,16 @@
 SERVICE_REBOOT = "reboot"
 
 
+@dataclass
+class HomeAssistantSagemcomFastData:
+    """SagemcomFast data stored in the Home Assistant data object."""
+
+    coordinator: SagemcomDataUpdateCoordinator
+    gateway: GatewayDeviceInfo
+
+
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
-    """Set up Sagemcom from a config entry."""
+    """Set up Sagemcom F@st from a config entry."""
 
     host = entry.data[CONF_HOST]
     username = entry.data[CONF_USERNAME]
@@ -94,10 +103,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
 
     await coordinator.async_refresh()
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = {
-        "coordinator": coordinator,
-        "update_listener": entry.add_update_listener(update_listener),
-    }
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = HomeAssistantSagemcomFastData(
+        coordinator=coordinator, gateway=gateway
+    )
 
     # Create gateway device in Home Assistant
     device_registry = hass.helpers.device_registry.async_get(hass)
@@ -125,22 +133,15 @@ async def async_command_reboot(call):
         hass, DOMAIN, SERVICE_REBOOT, async_command_reboot
     )
 
+    entry.async_on_unload(entry.add_update_listener(update_listener))
+
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry):
+async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
 
-    unload_ok = all(
-        await asyncio.gather(
-            *[
-                hass.config_entries.async_forward_entry_unload(entry, component)
-                for component in PLATFORMS
-            ]
-        )
-    )
-    if unload_ok:
-        hass.data[DOMAIN][entry.entry_id]["update_listener"]()
+    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
         hass.data[DOMAIN].pop(entry.entry_id)
 
     return unload_ok

diff --git a/custom_components/sagemcom_fast/device_tracker.py b/custom_components/sagemcom_fast/device_tracker.py
@@ -9,18 +9,19 @@
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 from sagemcom_api.models import Device
 
+from . import HomeAssistantSagemcomFastData
 from .const import DOMAIN
 from .coordinator import SagemcomDataUpdateCoordinator
 
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up from config entry."""
 
-    coordinator = hass.data[DOMAIN][config_entry.entry_id]["coordinator"]
+    data: HomeAssistantSagemcomFastData = hass.data[DOMAIN][config_entry.entry_id]
 
     async_add_entities(
-        SagemcomScannerEntity(coordinator, idx, config_entry.entry_id)
-        for idx, device in coordinator.data.items()
+        SagemcomScannerEntity(data.coordinator, idx, config_entry.entry_id)
+        for idx, device in data.coordinator.data.items()
     )