Use HassKey for wemo data (#143322)

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)
 
 

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

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])