From 1b4ba68e184aad0de7e1334b4dfc82eaf2c14016 Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Sat, 21 Sep 2024 13:15:42 +0200
Subject: [PATCH] Use HassKey in weather (#126329)

---
 homeassistant/components/weather/__init__.py      | 9 ++++-----
 homeassistant/components/weather/const.py         | 9 ++++++++-
 homeassistant/components/weather/websocket_api.py | 8 ++------
 3 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/homeassistant/components/weather/__init__.py b/homeassistant/components/weather/__init__.py
index 28f3e6b5c53..03b8addc1c9 100644
--- a/homeassistant/components/weather/__init__.py
+++ b/homeassistant/components/weather/__init__.py
@@ -63,6 +63,7 @@ from .const import (  # noqa: F401
     ATTR_WEATHER_WIND_SPEED,
     ATTR_WEATHER_WIND_SPEED_UNIT,
     DOMAIN,
+    DOMAIN_DATA,
     INTENT_GET_WEATHER,
     UNIT_CONVERSIONS,
     VALID_UNITS,
@@ -196,7 +197,7 @@ class Forecast(TypedDict, total=False):
 
 async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     """Set up the weather component."""
-    component = hass.data[DOMAIN] = EntityComponent[WeatherEntity](
+    component = hass.data[DOMAIN_DATA] = EntityComponent[WeatherEntity](
         _LOGGER, DOMAIN, hass, SCAN_INTERVAL
     )
     component.async_register_entity_service(
@@ -217,14 +218,12 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up a config entry."""
-    component: EntityComponent[WeatherEntity] = hass.data[DOMAIN]
-    return await component.async_setup_entry(entry)
+    return await hass.data[DOMAIN_DATA].async_setup_entry(entry)
 
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
-    component: EntityComponent[WeatherEntity] = hass.data[DOMAIN]
-    return await component.async_unload_entry(entry)
+    return await hass.data[DOMAIN_DATA].async_unload_entry(entry)
 
 
 class WeatherEntityDescription(EntityDescription, frozen_or_thawed=True):
diff --git a/homeassistant/components/weather/const.py b/homeassistant/components/weather/const.py
index 251bbd622fc..ef8eada2b3f 100644
--- a/homeassistant/components/weather/const.py
+++ b/homeassistant/components/weather/const.py
@@ -4,7 +4,7 @@ from __future__ import annotations
 
 from collections.abc import Callable
 from enum import IntFlag
-from typing import Final
+from typing import TYPE_CHECKING, Final
 
 from homeassistant.const import (
     UnitOfLength,
@@ -13,6 +13,7 @@ from homeassistant.const import (
     UnitOfSpeed,
     UnitOfTemperature,
 )
+from homeassistant.util.hass_dict import HassKey
 from homeassistant.util.unit_conversion import (
     DistanceConverter,
     PressureConverter,
@@ -20,6 +21,11 @@ from homeassistant.util.unit_conversion import (
     TemperatureConverter,
 )
 
+if TYPE_CHECKING:
+    from homeassistant.helpers.entity_component import EntityComponent
+
+    from . import WeatherEntity
+
 
 class WeatherEntityFeature(IntFlag):
     """Supported features of the update entity."""
@@ -48,6 +54,7 @@ ATTR_WEATHER_CLOUD_COVERAGE = "cloud_coverage"
 ATTR_WEATHER_UV_INDEX = "uv_index"
 
 DOMAIN: Final = "weather"
+DOMAIN_DATA: HassKey[EntityComponent[WeatherEntity]] = HassKey(DOMAIN)
 
 INTENT_GET_WEATHER = "HassGetWeather"
 
diff --git a/homeassistant/components/weather/websocket_api.py b/homeassistant/components/weather/websocket_api.py
index 98adbd1bd02..fb9759c9bdf 100644
--- a/homeassistant/components/weather/websocket_api.py
+++ b/homeassistant/components/weather/websocket_api.py
@@ -9,10 +9,9 @@ import voluptuous as vol
 from homeassistant.components import websocket_api
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import config_validation as cv
-from homeassistant.helpers.entity_component import EntityComponent
 from homeassistant.util.json import JsonValueType
 
-from .const import DOMAIN, VALID_UNITS, WeatherEntityFeature
+from .const import DOMAIN, DOMAIN_DATA, VALID_UNITS, WeatherEntityFeature
 
 FORECAST_TYPE_TO_FLAG = {
     "daily": WeatherEntityFeature.FORECAST_DAILY,
@@ -56,13 +55,10 @@ async def ws_subscribe_forecast(
     hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict[str, Any]
 ) -> None:
     """Subscribe to weather forecasts."""
-    from . import WeatherEntity  # pylint: disable=import-outside-toplevel
-
-    component: EntityComponent[WeatherEntity] = hass.data[DOMAIN]
     entity_id: str = msg["entity_id"]
     forecast_type: Literal["daily", "hourly", "twice_daily"] = msg["forecast_type"]
 
-    if not (entity := component.get_entity(msg["entity_id"])):
+    if not (entity := hass.data[DOMAIN_DATA].get_entity(msg["entity_id"])):
         connection.send_error(
             msg["id"],
             "invalid_entity_id",