Improve typing of nws (#98485)

* Improve typing of nws

* Address review comments

homeassistant/components/nws/__init__.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
 import datetime
 import logging
 from typing import TYPE_CHECKING
@@ -18,15 +19,7 @@ from homeassistant.helpers.event import async_track_point_in_utc_time
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 from homeassistant.util.dt import utcnow
 
-from .const import (
-    CONF_STATION,
-    COORDINATOR_FORECAST,
-    COORDINATOR_FORECAST_HOURLY,
-    COORDINATOR_OBSERVATION,
-    DOMAIN,
-    NWS_DATA,
-    UPDATE_TIME_PERIOD,
-)
+from .const import CONF_STATION, DOMAIN, UPDATE_TIME_PERIOD
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -42,6 +35,16 @@ def base_unique_id(latitude: float, longitude: float) -> str:
     return f"{latitude}_{longitude}"
 
 
+@dataclass
+class NWSData:
+    """Data for the National Weather Service integration."""
+
+    api: SimpleNWS
+    coordinator_observation: NwsDataUpdateCoordinator
+    coordinator_forecast: NwsDataUpdateCoordinator
+    coordinator_forecast_hourly: NwsDataUpdateCoordinator
+
+
 class NwsDataUpdateCoordinator(DataUpdateCoordinator[None]):
     """NWS data update coordinator.
 
@@ -150,12 +153,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         ),
     )
     nws_hass_data = hass.data.setdefault(DOMAIN, {})
-    nws_hass_data[entry.entry_id] = {
-        NWS_DATA: nws_data,
-        COORDINATOR_OBSERVATION: coordinator_observation,
-        COORDINATOR_FORECAST: coordinator_forecast,
-        COORDINATOR_FORECAST_HOURLY: coordinator_forecast_hourly,
-    }
+    nws_hass_data[entry.entry_id] = NWSData(
+        nws_data,
+        coordinator_observation,
+        coordinator_forecast,
+        coordinator_forecast_hourly,
+    )
 
     # Fetch initial data so we have data when entities subscribe
     await coordinator_observation.async_refresh()

homeassistant/components/nws/const.py

@@ -74,11 +74,6 @@ CONDITION_CLASSES: dict[str, list[str]] = {
 DAYNIGHT = "daynight"
 HOURLY = "hourly"
 
-NWS_DATA = "nws data"
-COORDINATOR_OBSERVATION = "coordinator_observation"
-COORDINATOR_FORECAST = "coordinator_forecast"
-COORDINATOR_FORECAST_HOURLY = "coordinator_forecast_hourly"
-
 OBSERVATION_VALID_TIME = timedelta(minutes=20)
 FORECAST_VALID_TIME = timedelta(minutes=45)
 # A lot of stations update once hourly plus some wiggle room

homeassistant/components/nws/sensor.py

@@ -5,8 +5,6 @@ from dataclasses import dataclass
 from types import MappingProxyType
 from typing import Any
 
-from pynws import SimpleNWS
-
 from homeassistant.components.sensor import (
     SensorDeviceClass,
     SensorEntity,
@@ -36,15 +34,8 @@ from homeassistant.util.unit_conversion import (
 )
 from homeassistant.util.unit_system import US_CUSTOMARY_SYSTEM
 
-from . import NwsDataUpdateCoordinator, base_unique_id, device_info
-from .const import (
-    ATTRIBUTION,
-    CONF_STATION,
-    COORDINATOR_OBSERVATION,
-    DOMAIN,
-    NWS_DATA,
-    OBSERVATION_VALID_TIME,
-)
+from . import NWSData, NwsDataUpdateCoordinator, base_unique_id, device_info
+from .const import ATTRIBUTION, CONF_STATION, DOMAIN, OBSERVATION_VALID_TIME
 
 PARALLEL_UPDATES = 0
 
@@ -152,14 +143,14 @@ async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up the NWS weather platform."""
-    hass_data = hass.data[DOMAIN][entry.entry_id]
+    nws_data: NWSData = hass.data[DOMAIN][entry.entry_id]
     station = entry.data[CONF_STATION]
 
     async_add_entities(
         NWSSensor(
             hass=hass,
             entry_data=entry.data,
-            hass_data=hass_data,
+            nws_data=nws_data,
             description=description,
             station=station,
         )
@@ -177,13 +168,13 @@ class NWSSensor(CoordinatorEntity[NwsDataUpdateCoordinator], SensorEntity):
         self,
         hass: HomeAssistant,
         entry_data: MappingProxyType[str, Any],
-        hass_data: dict[str, Any],
+        nws_data: NWSData,
         description: NWSSensorEntityDescription,
         station: str,
     ) -> None:
         """Initialise the platform with a data instance."""
-        super().__init__(hass_data[COORDINATOR_OBSERVATION])
-        self._nws: SimpleNWS = hass_data[NWS_DATA]
+        super().__init__(nws_data.coordinator_observation)
+        self._nws = nws_data.api
         self._latitude = entry_data[CONF_LATITUDE]
         self._longitude = entry_data[CONF_LONGITUDE]
         self.entity_description = description

homeassistant/components/nws/weather.py

@@ -35,19 +35,15 @@ from homeassistant.util.dt import utcnow
 from homeassistant.util.unit_conversion import SpeedConverter, TemperatureConverter
 from homeassistant.util.unit_system import UnitSystem
 
-from . import base_unique_id, device_info
+from . import NWSData, base_unique_id, device_info
 from .const import (
     ATTR_FORECAST_DETAILED_DESCRIPTION,
     ATTRIBUTION,
     CONDITION_CLASSES,
-    COORDINATOR_FORECAST,
-    COORDINATOR_FORECAST_HOURLY,
-    COORDINATOR_OBSERVATION,
     DAYNIGHT,
     DOMAIN,
     FORECAST_VALID_TIME,
     HOURLY,
-    NWS_DATA,
     OBSERVATION_VALID_TIME,
 )
 
@@ -84,12 +80,12 @@ async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up the NWS weather platform."""
-    hass_data = hass.data[DOMAIN][entry.entry_id]
+    nws_data: NWSData = hass.data[DOMAIN][entry.entry_id]
 
     async_add_entities(
         [
-            NWSWeather(entry.data, hass_data, DAYNIGHT, hass.config.units),
-            NWSWeather(entry.data, hass_data, HOURLY, hass.config.units),
+            NWSWeather(entry.data, nws_data, DAYNIGHT, hass.config.units),
+            NWSWeather(entry.data, nws_data, HOURLY, hass.config.units),
         ],
         False,
     )
@@ -112,19 +108,19 @@ class NWSWeather(WeatherEntity):
     def __init__(
         self,
         entry_data: MappingProxyType[str, Any],
-        hass_data: dict[str, Any],
+        nws_data: NWSData,
         mode: str,
         units: UnitSystem,
     ) -> None:
         """Initialise the platform with a data instance and station name."""
-        self.nws = hass_data[NWS_DATA]
+        self.nws = nws_data.api
         self.latitude = entry_data[CONF_LATITUDE]
         self.longitude = entry_data[CONF_LONGITUDE]
-        self.coordinator_observation = hass_data[COORDINATOR_OBSERVATION]
+        self.coordinator_observation = nws_data.coordinator_observation
         if mode == DAYNIGHT:
-            self.coordinator_forecast = hass_data[COORDINATOR_FORECAST]
+            self.coordinator_forecast = nws_data.coordinator_forecast
         else:
-            self.coordinator_forecast = hass_data[COORDINATOR_FORECAST_HOURLY]
+            self.coordinator_forecast = nws_data.coordinator_forecast_hourly
         self.station = self.nws.station
 
         self.mode = mode