Migrate Nuki to use dataclass for entry data (#101785)

Co-authored-by: J. Nick Koston <nick@koston.org>
Co-authored-by: Franck Nijhof <frenck@frenck.nl>

homeassistant/components/nuki/__init__.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 
 import asyncio
 from collections import defaultdict
+from dataclasses import dataclass
 from datetime import timedelta
 from http import HTTPStatus
 import logging
@@ -38,15 +39,7 @@ from homeassistant.helpers.update_coordinator import (
     UpdateFailed,
 )
 
-from .const import (
+from .const import DEFAULT_TIMEOUT, DOMAIN, ERROR_STATES
-    DATA_BRIDGE,
-    DATA_COORDINATOR,
-    DATA_LOCKS,
-    DATA_OPENERS,
-    DEFAULT_TIMEOUT,
-    DOMAIN,
-    ERROR_STATES,
-)
 from .helpers import NukiWebhookException, parse_id
 
 _NukiDeviceT = TypeVar("_NukiDeviceT", bound=NukiDevice)
@@ -57,6 +50,16 @@ PLATFORMS = [Platform.BINARY_SENSOR, Platform.LOCK, Platform.SENSOR]
 UPDATE_INTERVAL = timedelta(seconds=30)
 
 
+@dataclass(slots=True)
+class NukiEntryData:
+    """Class to hold Nuki data."""
+
+    coordinator: NukiCoordinator
+    bridge: NukiBridge
+    locks: list[NukiLock]
+    openers: list[NukiOpener]
+
+
 def _get_bridge_devices(bridge: NukiBridge) -> tuple[list[NukiLock], list[NukiOpener]]:
     return bridge.locks, bridge.openers
 
@@ -74,14 +77,15 @@ async def _create_webhook(
         except ValueError:
             return web.Response(status=HTTPStatus.BAD_REQUEST)
 
-        locks = hass.data[DOMAIN][entry.entry_id][DATA_LOCKS]
+        entry_data: NukiEntryData = hass.data[DOMAIN][entry.entry_id]
-        openers = hass.data[DOMAIN][entry.entry_id][DATA_OPENERS]
+        locks = entry_data.locks
+        openers = entry_data.openers
 
         devices = [x for x in locks + openers if x.nuki_id == data["nukiId"]]
         if len(devices) == 1:
             devices[0].update_from_callback(data)
 
-        coordinator = hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR]
+        coordinator = entry_data.coordinator
         coordinator.async_set_updated_data(None)
 
         return web.Response(status=HTTPStatus.OK)
@@ -232,13 +236,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     )
 
     coordinator = NukiCoordinator(hass, bridge, locks, openers)
-
+    hass.data[DOMAIN][entry.entry_id] = NukiEntryData(
-    hass.data[DOMAIN][entry.entry_id] = {
+        coordinator=coordinator,
-        DATA_COORDINATOR: coordinator,
+        bridge=bridge,
-        DATA_BRIDGE: bridge,
+        locks=locks,
-        DATA_LOCKS: locks,
+        openers=openers,
-        DATA_OPENERS: openers,
+    )
-    }
 
     # Fetch initial data so we have data when entities subscribe
     await coordinator.async_refresh()
@@ -251,11 +254,13 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload the Nuki entry."""
     webhook.async_unregister(hass, entry.entry_id)
+    entry_data: NukiEntryData = hass.data[DOMAIN][entry.entry_id]
+
     try:
         async with asyncio.timeout(10):
             await hass.async_add_executor_job(
                 _remove_webhook,
-                hass.data[DOMAIN][entry.entry_id][DATA_BRIDGE],
+                entry_data.bridge,
                 entry.entry_id,
             )
     except InvalidCredentialsException as err:

homeassistant/components/nuki/binary_sensor.py

@@ -12,22 +12,21 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import NukiCoordinator, NukiEntity
+from . import NukiEntity, NukiEntryData
-from .const import ATTR_NUKI_ID, DATA_COORDINATOR, DATA_LOCKS, DOMAIN as NUKI_DOMAIN
+from .const import ATTR_NUKI_ID, DOMAIN as NUKI_DOMAIN
 
 
 async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up the Nuki lock binary sensor."""
-    data = hass.data[NUKI_DOMAIN][entry.entry_id]
+    entry_data: NukiEntryData = hass.data[NUKI_DOMAIN][entry.entry_id]
-    coordinator: NukiCoordinator = data[DATA_COORDINATOR]
 
     entities = []
 
-    for lock in data[DATA_LOCKS]:
+    for lock in entry_data.locks:
         if lock.is_door_sensor_activated:
-            entities.extend([NukiDoorsensorEntity(coordinator, lock)])
+            entities.extend([NukiDoorsensorEntity(entry_data.coordinator, lock)])
 
     async_add_entities(entities)
 

homeassistant/components/nuki/const.py

@@ -7,12 +7,6 @@ ATTR_NUKI_ID = "nuki_id"
 ATTR_ENABLE = "enable"
 ATTR_UNLATCH = "unlatch"
 
-# Data
-DATA_BRIDGE = "nuki_bridge_data"
-DATA_LOCKS = "nuki_locks_data"
-DATA_OPENERS = "nuki_openers_data"
-DATA_COORDINATOR = "nuki_coordinator"
-
 # Defaults
 DEFAULT_PORT = 8080
 DEFAULT_TIMEOUT = 20

homeassistant/components/nuki/lock.py

@@ -16,15 +16,12 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers import config_validation as cv, entity_platform
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import NukiCoordinator, NukiEntity
+from . import NukiEntity, NukiEntryData
 from .const import (
     ATTR_BATTERY_CRITICAL,
     ATTR_ENABLE,
     ATTR_NUKI_ID,
     ATTR_UNLATCH,
-    DATA_COORDINATOR,
-    DATA_LOCKS,
-    DATA_OPENERS,
     DOMAIN as NUKI_DOMAIN,
     ERROR_STATES,
 )
@@ -37,14 +34,14 @@ async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up the Nuki lock platform."""
-    data = hass.data[NUKI_DOMAIN][entry.entry_id]
+    entry_data: NukiEntryData = hass.data[NUKI_DOMAIN][entry.entry_id]
-    coordinator: NukiCoordinator = data[DATA_COORDINATOR]
+    coordinator = entry_data.coordinator
 
     entities: list[NukiDeviceEntity] = [
-        NukiLockEntity(coordinator, lock) for lock in data[DATA_LOCKS]
+        NukiLockEntity(coordinator, lock) for lock in entry_data.locks
     ]
     entities.extend(
-        [NukiOpenerEntity(coordinator, opener) for opener in data[DATA_OPENERS]]
+        [NukiOpenerEntity(coordinator, opener) for opener in entry_data.openers]
     )
     async_add_entities(entities)
 

homeassistant/components/nuki/sensor.py

@@ -9,19 +9,18 @@ from homeassistant.const import PERCENTAGE, EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import NukiEntity
+from . import NukiEntity, NukiEntryData
-from .const import ATTR_NUKI_ID, DATA_COORDINATOR, DATA_LOCKS, DOMAIN as NUKI_DOMAIN
+from .const import ATTR_NUKI_ID, DOMAIN as NUKI_DOMAIN
 
 
 async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up the Nuki lock sensor."""
-    data = hass.data[NUKI_DOMAIN][entry.entry_id]
+    entry_data: NukiEntryData = hass.data[NUKI_DOMAIN][entry.entry_id]
-    coordinator = data[DATA_COORDINATOR]
 
     async_add_entities(
-        NukiBatterySensor(coordinator, lock) for lock in data[DATA_LOCKS]
+        NukiBatterySensor(entry_data.coordinator, lock) for lock in entry_data.locks
     )