Improve type annotations for Airly integration (#49898)

.strict-typing

@@ -3,6 +3,7 @@
 # to enable strict mypy checks.
 
 homeassistant.components
+homeassistant.components.airly.*
 homeassistant.components.automation.*
 homeassistant.components.binary_sensor.*
 homeassistant.components.bond.*

homeassistant/components/airly/__init__.py

@@ -1,15 +1,21 @@
 """The Airly integration."""
+from __future__ import annotations
+
 from datetime import timedelta
 import logging
 from math import ceil
 
+from aiohttp import ClientSession
 from aiohttp.client_exceptions import ClientConnectorError
 from airly import Airly
 from airly.exceptions import AirlyError
 import async_timeout
 
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_API_KEY, CONF_LATITUDE, CONF_LONGITUDE
+from homeassistant.core import HomeAssistant
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
+from homeassistant.helpers.device_registry import async_get_registry
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 from homeassistant.util import dt as dt_util
 
@@ -30,7 +36,7 @@
 _LOGGER = logging.getLogger(__name__)
 
 
-def set_update_interval(instances, requests_remaining):
+def set_update_interval(instances_count: int, requests_remaining: int) -> timedelta:
     """
     Return data update interval.
 
@@ -46,7 +52,7 @@ def set_update_interval(instances, requests_remaining):
     interval = timedelta(
         minutes=min(
             max(
-                ceil(minutes_to_midnight / requests_remaining * instances),
+                ceil(minutes_to_midnight / requests_remaining * instances_count),
                 MIN_UPDATE_INTERVAL,
             ),
             MAX_UPDATE_INTERVAL,
@@ -58,19 +64,28 @@ def set_update_interval(instances, requests_remaining):
     return interval
 
 
-async def async_setup_entry(hass, config_entry):
+async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Airly as config entry."""
-    api_key = config_entry.data[CONF_API_KEY]
-    latitude = config_entry.data[CONF_LATITUDE]
-    longitude = config_entry.data[CONF_LONGITUDE]
-    use_nearest = config_entry.data.get(CONF_USE_NEAREST, False)
+    api_key = entry.data[CONF_API_KEY]
+    latitude = entry.data[CONF_LATITUDE]
+    longitude = entry.data[CONF_LONGITUDE]
+    use_nearest = entry.data.get(CONF_USE_NEAREST, False)
 
     # For backwards compat, set unique ID
-    if config_entry.unique_id is None:
+    if entry.unique_id is None:
         hass.config_entries.async_update_entry(
-            config_entry, unique_id=f"{latitude}-{longitude}"
+            entry, unique_id=f"{latitude}-{longitude}"
         )
 
+    # identifiers in device_info should use Tuple[str, str, str] type, but latitude and
+    # longitude are float, so we convert old device entries to use correct types
+    device_registry = await async_get_registry(hass)
+    old_ids = (DOMAIN, latitude, longitude)
+    device_entry = device_registry.async_get_device({old_ids})
+    if device_entry and entry.entry_id in device_entry.config_entries:
+        new_ids = (DOMAIN, str(latitude), str(longitude))
+        device_registry.async_update_device(device_entry.id, new_identifiers={new_ids})
+
     websession = async_get_clientsession(hass)
 
     update_interval = timedelta(minutes=MIN_UPDATE_INTERVAL)
@@ -81,21 +96,19 @@ async def async_setup_entry(hass, config_entry):
     await coordinator.async_config_entry_first_refresh()
 
     hass.data.setdefault(DOMAIN, {})
-    hass.data[DOMAIN][config_entry.entry_id] = coordinator
+    hass.data[DOMAIN][entry.entry_id] = coordinator
 
-    hass.config_entries.async_setup_platforms(config_entry, PLATFORMS)
+    hass.config_entries.async_setup_platforms(entry, PLATFORMS)
 
     return True
 
 
-async def async_unload_entry(hass, config_entry):
+async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(
-        config_entry, PLATFORMS
-    )
+    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
 
     if unload_ok:
-        hass.data[DOMAIN].pop(config_entry.entry_id)
+        hass.data[DOMAIN].pop(entry.entry_id)
 
     return unload_ok
 
@@ -105,13 +118,13 @@ class AirlyDataUpdateCoordinator(DataUpdateCoordinator):
 
     def __init__(
         self,
-        hass,
-        session,
-        api_key,
-        latitude,
-        longitude,
-        update_interval,
-        use_nearest,
+        hass: HomeAssistant,
+        session: ClientSession,
+        api_key: str,
+        latitude: float,
+        longitude: float,
+        update_interval: timedelta,
+        use_nearest: bool,
     ):
         """Initialize."""
         self.latitude = latitude
@@ -121,9 +134,9 @@ def __init__(
 
         super().__init__(hass, _LOGGER, name=DOMAIN, update_interval=update_interval)
 
-    async def _async_update_data(self):
+    async def _async_update_data(self) -> dict[str, str | float | int]:
         """Update data via library."""
-        data = {}
+        data: dict[str, str | float | int] = {}
         if self.use_nearest:
             measurements = self.airly.create_measurements_session_nearest(
                 self.latitude, self.longitude, max_distance_km=5

homeassistant/components/airly/air_quality.py

@@ -1,13 +1,22 @@
 """Support for the Airly air_quality service."""
+from __future__ import annotations
+
+from typing import Any
+
 from homeassistant.components.air_quality import (
     ATTR_AQI,
     ATTR_PM_2_5,
     ATTR_PM_10,
     AirQualityEntity,
 )
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_NAME
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity import DeviceInfo
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
+from . import AirlyDataUpdateCoordinator
 from .const import (
     ATTR_API_ADVICE,
     ATTR_API_CAQI,
@@ -36,88 +45,81 @@
 PARALLEL_UPDATES = 1
 
 
-async def async_setup_entry(hass, config_entry, async_add_entities):
+async def async_setup_entry(
+    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+) -> None:
     """Set up Airly air_quality entity based on a config entry."""
-    name = config_entry.data[CONF_NAME]
+    name = entry.data[CONF_NAME]
 
-    coordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator = hass.data[DOMAIN][entry.entry_id]
 
     async_add_entities([AirlyAirQuality(coordinator, name)], False)
 
 
-def round_state(func):
-    """Round state."""
-
-    def _decorator(self):
-        res = func(self)
-        if isinstance(res, float):
-            return round(res)
-        return res
-
-    return _decorator
-
-
 class AirlyAirQuality(CoordinatorEntity, AirQualityEntity):
     """Define an Airly air quality."""
 
-    def __init__(self, coordinator, name):
+    coordinator: AirlyDataUpdateCoordinator
+
+    def __init__(self, coordinator: AirlyDataUpdateCoordinator, name: str) -> None:
         """Initialize."""
         super().__init__(coordinator)
         self._name = name
         self._icon = "mdi:blur"
 
     @property
-    def name(self):
+    def name(self) -> str:
         """Return the name."""
         return self._name
 
     @property
-    def icon(self):
+    def icon(self) -> str:
         """Return the icon."""
         return self._icon
 
     @property
-    @round_state
-    def air_quality_index(self):
+    def air_quality_index(self) -> float | None:
         """Return the air quality index."""
-        return self.coordinator.data[ATTR_API_CAQI]
+        return round_state(self.coordinator.data[ATTR_API_CAQI])
 
     @property
-    @round_state
-    def particulate_matter_2_5(self):
+    def particulate_matter_2_5(self) -> float | None:
         """Return the particulate matter 2.5 level."""
-        return self.coordinator.data.get(ATTR_API_PM25)
+        return round_state(self.coordinator.data.get(ATTR_API_PM25))
 
     @property
-    @round_state
-    def particulate_matter_10(self):
+    def particulate_matter_10(self) -> float | None:
         """Return the particulate matter 10 level."""
-        return self.coordinator.data.get(ATTR_API_PM10)
+        return round_state(self.coordinator.data.get(ATTR_API_PM10))
 
     @property
-    def attribution(self):
+    def attribution(self) -> str:
         """Return the attribution."""
         return ATTRIBUTION
 
     @property
-    def unique_id(self):
+    def unique_id(self) -> str:
         """Return a unique_id for this entity."""
         return f"{self.coordinator.latitude}-{self.coordinator.longitude}"
 
     @property
-    def device_info(self):
+    def device_info(self) -> DeviceInfo:
         """Return the device info."""
         return {
             "identifiers": {
-                (DOMAIN, self.coordinator.latitude, self.coordinator.longitude)
+                (
+                    DOMAIN,
+                    str(self.coordinator.latitude),
+                    str(self.coordinator.longitude),
+                )
             },
             "name": DEFAULT_NAME,
             "manufacturer": MANUFACTURER,
             "entry_type": "service",
         }
 
     @property
-    def extra_state_attributes(self):
+    def extra_state_attributes(self) -> dict[str, Any]:
         """Return the state attributes."""
         attrs = {
             LABEL_AQI_DESCRIPTION: self.coordinator.data[ATTR_API_CAQI_DESCRIPTION],
@@ -135,3 +137,8 @@ def extra_state_attributes(self):
                 self.coordinator.data[ATTR_API_PM10_PERCENT]
             )
         return attrs
+
+
+def round_state(state: float | None) -> float | None:
+    """Round state."""
+    return round(state) if state else state

homeassistant/components/airly/config_flow.py

@@ -1,4 +1,9 @@
 """Adds config flow for Airly."""
+from __future__ import annotations
+
+from typing import Any
+
+from aiohttp import ClientSession
 from airly import Airly
 from airly.exceptions import AirlyError
 import async_timeout
@@ -13,6 +18,7 @@
     HTTP_NOT_FOUND,
     HTTP_UNAUTHORIZED,
 )
+from homeassistant.data_entry_flow import FlowResult
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 import homeassistant.helpers.config_validation as cv
 
@@ -24,7 +30,9 @@ class AirlyFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
 
     VERSION = 1
 
-    async def async_step_user(self, user_input=None):
+    async def async_step_user(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
         """Handle a flow initialized by the user."""
         errors = {}
         use_nearest = False
@@ -84,7 +92,13 @@ async def async_step_user(self, user_input=None):
         )
 
 
-async def test_location(client, api_key, latitude, longitude, use_nearest=False):
+async def test_location(
+    client: ClientSession,
+    api_key: str,
+    latitude: float,
+    longitude: float,
+    use_nearest: bool = False,
+) -> bool:
     """Return true if location is valid."""
     airly = Airly(api_key, client)
     if use_nearest: