From 73f636c40d7078badccfd1509ea8cc84c99106be Mon Sep 17 00:00:00 2001 From: Marc Mueller <30130371+cdce8p@users.noreply.github.com> Date: Tue, 22 Apr 2025 12:42:24 +0200 Subject: [PATCH] Use HassKey for wemo data (#143322) --- homeassistant/components/wemo/__init__.py | 8 ++++---- homeassistant/components/wemo/coordinator.py | 6 +++--- homeassistant/components/wemo/models.py | 12 ++++-------- 3 files changed, 11 insertions(+), 15 deletions(-) diff --git a/homeassistant/components/wemo/__init__.py b/homeassistant/components/wemo/__init__.py index 3ef7ac92f98..96e61dfded6 100644 --- a/homeassistant/components/wemo/__init__.py +++ b/homeassistant/components/wemo/__init__.py @@ -21,7 +21,7 @@ from homeassistant.util.async_ import gather_with_limited_concurrency from .const import DOMAIN from .coordinator import DeviceCoordinator, async_register_device -from .models import WemoConfigEntryData, WemoData, async_wemo_data +from .models import DATA_WEMO, WemoConfigEntryData, WemoData # Max number of devices to initialize at once. This limit is in place to # avoid tying up too many executor threads with WeMo device setup. @@ -117,7 +117,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up a wemo config entry.""" - wemo_data = async_wemo_data(hass) + wemo_data = hass.data[DATA_WEMO] dispatcher = WemoDispatcher(entry) discovery = WemoDiscovery(hass, dispatcher, wemo_data.static_config, entry) wemo_data.config_entry_data = WemoConfigEntryData( @@ -138,7 +138,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Unload a wemo config entry.""" _LOGGER.debug("Unloading WeMo") - wemo_data = async_wemo_data(hass) + wemo_data = hass.data[DATA_WEMO] wemo_data.config_entry_data.discovery.async_stop_discovery() @@ -161,7 +161,7 @@ async def async_wemo_dispatcher_connect( module = dispatch.__module__ # Example: "homeassistant.components.wemo.switch" platform = Platform(module.rsplit(".", 1)[1]) - dispatcher = async_wemo_data(hass).config_entry_data.dispatcher + dispatcher = hass.data[DATA_WEMO].config_entry_data.dispatcher await dispatcher.async_connect_platform(platform, dispatch) diff --git a/homeassistant/components/wemo/coordinator.py b/homeassistant/components/wemo/coordinator.py index 0aaedf598d2..6cda83f6419 100644 --- a/homeassistant/components/wemo/coordinator.py +++ b/homeassistant/components/wemo/coordinator.py @@ -29,7 +29,7 @@ from homeassistant.helpers.device_registry import CONNECTION_UPNP, DeviceInfo from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed from .const import DOMAIN, WEMO_SUBSCRIPTION_EVENT -from .models import async_wemo_data +from .models import DATA_WEMO _LOGGER = logging.getLogger(__name__) @@ -316,9 +316,9 @@ def async_get_coordinator(hass: HomeAssistant, device_id: str) -> DeviceCoordina @callback def _async_coordinators(hass: HomeAssistant) -> dict[str, DeviceCoordinator]: - return async_wemo_data(hass).config_entry_data.device_coordinators + return hass.data[DATA_WEMO].config_entry_data.device_coordinators @callback def _async_registry(hass: HomeAssistant) -> SubscriptionRegistry: - return async_wemo_data(hass).registry + return hass.data[DATA_WEMO].registry diff --git a/homeassistant/components/wemo/models.py b/homeassistant/components/wemo/models.py index 80213c9ba33..b96cd502cd4 100644 --- a/homeassistant/components/wemo/models.py +++ b/homeassistant/components/wemo/models.py @@ -4,11 +4,11 @@ from __future__ import annotations from collections.abc import Sequence from dataclasses import dataclass -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING import pywemo -from homeassistant.core import HomeAssistant, callback +from homeassistant.util.hass_dict import HassKey from .const import DOMAIN @@ -16,6 +16,8 @@ if TYPE_CHECKING: # Avoid circular dependencies. from . import HostPortTuple, WemoDiscovery, WemoDispatcher from .coordinator import DeviceCoordinator +DATA_WEMO: HassKey[WemoData] = HassKey(DOMAIN) + @dataclass class WemoConfigEntryData: @@ -37,9 +39,3 @@ class WemoData: # unloaded. It's a programmer error if config_entry_data is accessed when the # config entry is not loaded config_entry_data: WemoConfigEntryData = None # type: ignore[assignment] - - -@callback -def async_wemo_data(hass: HomeAssistant) -> WemoData: - """Fetch WemoData with proper typing.""" - return cast(WemoData, hass.data[DOMAIN])