Improve type annotations for GIOS integration (#50454)

.strict-typing

@@ -15,6 +15,7 @@ homeassistant.components.device_automation.*
 homeassistant.components.elgato.*
 homeassistant.components.frontend.*
 homeassistant.components.geo_location.*
+homeassistant.components.gios.*
 homeassistant.components.group.*
 homeassistant.components.history.*
 homeassistant.components.http.*

homeassistant/components/gios/__init__.py

@@ -1,11 +1,18 @@
 """The GIOS component."""
-import logging
+from __future__ import annotations
 
+import logging
+from typing import Any, Dict, cast
+
+from aiohttp import ClientSession
 from aiohttp.client_exceptions import ClientConnectorError
 from async_timeout import timeout
 from gios import ApiError, Gios, InvalidSensorsData, NoStationError
 
+from homeassistant.config_entries import ConfigEntry
+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 .const import API_TIMEOUT, CONF_STATION_ID, DOMAIN, SCAN_INTERVAL
@@ -15,10 +22,22 @@ _LOGGER = logging.getLogger(__name__)
 PLATFORMS = ["air_quality"]
 
 
-async def async_setup_entry(hass, entry):
+async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up GIOS as config entry."""
-    station_id = entry.data[CONF_STATION_ID]
-    _LOGGER.debug("Using station_id: %s", station_id)
+    station_id: int = entry.data[CONF_STATION_ID]
+    _LOGGER.debug("Using station_id: %d", station_id)
+
+    # We used to use int as config_entry unique_id, convert this to str.
+    if isinstance(entry.unique_id, int):  # type: ignore[unreachable]
+        hass.config_entries.async_update_entry(entry, unique_id=str(station_id))  # type: ignore[unreachable]
+
+    # We used to use int in device_entry identifiers, convert this to str.
+    device_registry = await async_get_registry(hass)
+    old_ids = (DOMAIN, station_id)
+    device_entry = device_registry.async_get_device({old_ids})  # type: ignore[arg-type]
+    if device_entry and entry.entry_id in device_entry.config_entries:
+        new_ids = (DOMAIN, str(station_id))
+        device_registry.async_update_device(device_entry.id, new_identifiers={new_ids})
 
     websession = async_get_clientsession(hass)
 
@@ -33,26 +52,32 @@ async def async_setup_entry(hass, entry):
     return True
 
 
-async def async_unload_entry(hass, entry):
+async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
-    hass.data[DOMAIN].pop(entry.entry_id)
-    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
+    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
+
+    if unload_ok:
+        hass.data[DOMAIN].pop(entry.entry_id)
+
+    return unload_ok
 
 
 class GiosDataUpdateCoordinator(DataUpdateCoordinator):
     """Define an object to hold GIOS data."""
 
-    def __init__(self, hass, session, station_id):
+    def __init__(
+        self, hass: HomeAssistant, session: ClientSession, station_id: int
+    ) -> None:
         """Class to manage fetching GIOS data API."""
         self.gios = Gios(station_id, session)
 
         super().__init__(hass, _LOGGER, name=DOMAIN, update_interval=SCAN_INTERVAL)
 
-    async def _async_update_data(self):
+    async def _async_update_data(self) -> dict[str, Any]:
         """Update data via library."""
         try:
             with timeout(API_TIMEOUT):
-                return await self.gios.async_update()
+                return cast(Dict[str, Any], await self.gios.async_update())
         except (
             ApiError,
             NoStationError,

homeassistant/components/gios/air_quality.py

@@ -1,8 +1,18 @@
 """Support for the GIOS service."""
+from __future__ import annotations
+
+from typing import Any, Optional, cast
+
 from homeassistant.components.air_quality import 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.entity_registry import async_get_registry
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
+from . import GiosDataUpdateCoordinator
 from .const import (
     API_AQI,
     API_CO,
@@ -23,111 +33,107 @@ from .const import (
 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:
     """Add a GIOS entities from 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([GiosAirQuality(coordinator, name)], False)
+    # We used to use int as entity unique_id, convert this to str.
+    entity_registry = await async_get_registry(hass)
+    old_entity_id = entity_registry.async_get_entity_id(
+        "air_quality", DOMAIN, coordinator.gios.station_id
+    )
+    if old_entity_id is not None:
+        entity_registry.async_update_entity(
+            old_entity_id, new_unique_id=str(coordinator.gios.station_id)
+        )
 
-
-def round_state(func):
-    """Round state."""
-
-    def _decorator(self):
-        res = func(self)
-        if isinstance(res, float):
-            return round(res)
-        return res
-
-    return _decorator
+    async_add_entities([GiosAirQuality(coordinator, name)])
 
 
 class GiosAirQuality(CoordinatorEntity, AirQualityEntity):
     """Define an GIOS sensor."""
 
-    def __init__(self, coordinator, name):
+    coordinator: GiosDataUpdateCoordinator
+
+    def __init__(self, coordinator: GiosDataUpdateCoordinator, name: str) -> None:
         """Initialize."""
         super().__init__(coordinator)
         self._name = name
-        self._attrs = {}
+        self._attrs: dict[str, Any] = {}
 
     @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."""
-        if self.air_quality_index in ICONS_MAP:
+        if self.air_quality_index is not None and self.air_quality_index in ICONS_MAP:
             return ICONS_MAP[self.air_quality_index]
         return "mdi:blur"
 
     @property
-    def air_quality_index(self):
+    def air_quality_index(self) -> str | None:
         """Return the air quality index."""
-        return self._get_sensor_value(API_AQI)
+        return cast(Optional[str], self.coordinator.data.get(API_AQI, {}).get("value"))
 
     @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._get_sensor_value(API_PM25)
+        return round_state(self._get_sensor_value(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._get_sensor_value(API_PM10)
+        return round_state(self._get_sensor_value(API_PM10))
 
     @property
-    @round_state
-    def ozone(self):
+    def ozone(self) -> float | None:
         """Return the O3 (ozone) level."""
-        return self._get_sensor_value(API_O3)
+        return round_state(self._get_sensor_value(API_O3))
 
     @property
-    @round_state
-    def carbon_monoxide(self):
+    def carbon_monoxide(self) -> float | None:
         """Return the CO (carbon monoxide) level."""
-        return self._get_sensor_value(API_CO)
+        return round_state(self._get_sensor_value(API_CO))
 
     @property
-    @round_state
-    def sulphur_dioxide(self):
+    def sulphur_dioxide(self) -> float | None:
         """Return the SO2 (sulphur dioxide) level."""
-        return self._get_sensor_value(API_SO2)
+        return round_state(self._get_sensor_value(API_SO2))
 
     @property
-    @round_state
-    def nitrogen_dioxide(self):
+    def nitrogen_dioxide(self) -> float | None:
         """Return the NO2 (nitrogen dioxide) level."""
-        return self._get_sensor_value(API_NO2)
+        return round_state(self._get_sensor_value(API_NO2))
 
     @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 self.coordinator.gios.station_id
+        return str(self.coordinator.gios.station_id)
 
     @property
-    def device_info(self):
+    def device_info(self) -> DeviceInfo:
         """Return the device info."""
         return {
-            "identifiers": {(DOMAIN, self.coordinator.gios.station_id)},
+            "identifiers": {(DOMAIN, str(self.coordinator.gios.station_id))},
             "name": DEFAULT_NAME,
             "manufacturer": MANUFACTURER,
             "entry_type": "service",
         }
 
     @property
-    def extra_state_attributes(self):
+    def extra_state_attributes(self) -> dict[str, Any] | None:
         """Return the state attributes."""
         # Different measuring stations have different sets of sensors. We don't know
         # what data we will get.
@@ -139,8 +145,13 @@ class GiosAirQuality(CoordinatorEntity, AirQualityEntity):
         self._attrs[ATTR_STATION] = self.coordinator.gios.station_name
         return self._attrs
 
-    def _get_sensor_value(self, sensor):
+    def _get_sensor_value(self, sensor: str) -> float | None:
         """Return value of specified sensor."""
         if sensor in self.coordinator.data:
-            return self.coordinator.data[sensor]["value"]
+            return cast(float, self.coordinator.data[sensor]["value"])
         return None
+
+
+def round_state(state: float | None) -> float | None:
+    """Round state."""
+    return round(state) if state is not None else None

homeassistant/components/gios/config_flow.py

@@ -1,5 +1,8 @@
 """Adds config flow for GIOS."""
+from __future__ import annotations
+
 import asyncio
+from typing import Any
 
 from aiohttp.client_exceptions import ClientConnectorError
 from async_timeout import timeout
@@ -8,6 +11,7 @@ import voluptuous as vol
 
 from homeassistant import config_entries
 from homeassistant.const import CONF_NAME
+from homeassistant.data_entry_flow import FlowResult
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
 from .const import API_TIMEOUT, CONF_STATION_ID, DEFAULT_NAME, DOMAIN
@@ -25,14 +29,16 @@ class GiosFlowHandler(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 = {}
 
         if user_input is not None:
             try:
                 await self.async_set_unique_id(
-                    user_input[CONF_STATION_ID], raise_on_progress=False
+                    str(user_input[CONF_STATION_ID]), raise_on_progress=False
                 )
                 self._abort_if_unique_id_configured()
 

homeassistant/components/gios/const.py

@@ -1,5 +1,8 @@
 """Constants for GIOS integration."""
+from __future__ import annotations
+
 from datetime import timedelta
+from typing import Final
 
 from homeassistant.components.air_quality import (
     ATTR_CO,
@@ -10,33 +13,33 @@ from homeassistant.components.air_quality import (
     ATTR_SO2,
 )
 
-ATTRIBUTION = "Data provided by GIOŚ"
+ATTRIBUTION: Final = "Data provided by GIOŚ"
 
-ATTR_STATION = "station"
-CONF_STATION_ID = "station_id"
-DEFAULT_NAME = "GIOŚ"
+ATTR_STATION: Final = "station"
+CONF_STATION_ID: Final = "station_id"
+DEFAULT_NAME: Final = "GIOŚ"
 # Term of service GIOŚ allow downloading data no more than twice an hour.
-SCAN_INTERVAL = timedelta(minutes=30)
-DOMAIN = "gios"
-MANUFACTURER = "Główny Inspektorat Ochrony Środowiska"
+SCAN_INTERVAL: Final = timedelta(minutes=30)
+DOMAIN: Final = "gios"
+MANUFACTURER: Final = "Główny Inspektorat Ochrony Środowiska"
 
-API_AQI = "aqi"
-API_CO = "co"
-API_NO2 = "no2"
-API_O3 = "o3"
-API_PM10 = "pm10"
-API_PM25 = "pm2.5"
-API_SO2 = "so2"
+API_AQI: Final = "aqi"
+API_CO: Final = "co"
+API_NO2: Final = "no2"
+API_O3: Final = "o3"
+API_PM10: Final = "pm10"
+API_PM25: Final = "pm2.5"
+API_SO2: Final = "so2"
 
-API_TIMEOUT = 30
+API_TIMEOUT: Final = 30
 
-AQI_GOOD = "dobry"
-AQI_MODERATE = "umiarkowany"
-AQI_POOR = "dostateczny"
-AQI_VERY_GOOD = "bardzo dobry"
-AQI_VERY_POOR = "zły"
+AQI_GOOD: Final = "dobry"
+AQI_MODERATE: Final = "umiarkowany"
+AQI_POOR: Final = "dostateczny"
+AQI_VERY_GOOD: Final = "bardzo dobry"
+AQI_VERY_POOR: Final = "zły"
 
-ICONS_MAP = {
+ICONS_MAP: Final[dict[str, str]] = {
     AQI_VERY_GOOD: "mdi:emoticon-excited",
     AQI_GOOD: "mdi:emoticon-happy",
     AQI_MODERATE: "mdi:emoticon-neutral",
@@ -44,7 +47,7 @@ ICONS_MAP = {
     AQI_VERY_POOR: "mdi:emoticon-dead",
 }
 
-SENSOR_MAP = {
+SENSOR_MAP: Final[dict[str, str]] = {
     API_CO: ATTR_CO,
     API_NO2: ATTR_NO2,
     API_O3: ATTR_OZONE,

homeassistant/components/gios/system_health.py

@@ -1,8 +1,12 @@
 """Provide info to system health."""
+from __future__ import annotations
+
+from typing import Any, Final
+
 from homeassistant.components import system_health
 from homeassistant.core import HomeAssistant, callback
 
-API_ENDPOINT = "http://api.gios.gov.pl/"
+API_ENDPOINT: Final = "http://api.gios.gov.pl/"
 
 
 @callback
@@ -13,7 +17,7 @@ def async_register(
     register.async_register_info(system_health_info)
 
 
-async def system_health_info(hass):
+async def system_health_info(hass: HomeAssistant) -> dict[str, Any]:
     """Get info for the info page."""
     return {
         "can_reach_server": system_health.async_check_can_reach_url(hass, API_ENDPOINT)

mypy.ini

@@ -176,6 +176,17 @@ no_implicit_optional = true
 warn_return_any = true
 warn_unreachable = true
 
+[mypy-homeassistant.components.gios.*]
+check_untyped_defs = true
+disallow_incomplete_defs = true
+disallow_subclassing_any = true
+disallow_untyped_calls = true
+disallow_untyped_decorators = true
+disallow_untyped_defs = true
+no_implicit_optional = true
+warn_return_any = true
+warn_unreachable = true
+
 [mypy-homeassistant.components.group.*]
 check_untyped_defs = true
 disallow_incomplete_defs = true
@@ -760,9 +771,6 @@ ignore_errors = true
 [mypy-homeassistant.components.geniushub.*]
 ignore_errors = true
 
-[mypy-homeassistant.components.gios.*]
-ignore_errors = true
-
 [mypy-homeassistant.components.glances.*]
 ignore_errors = true
 

script/hassfest/mypy_config.py

@@ -73,7 +73,6 @@ IGNORED_MODULES: Final[list[str]] = [
     "homeassistant.components.fritzbox.*",
     "homeassistant.components.garmin_connect.*",
     "homeassistant.components.geniushub.*",
-    "homeassistant.components.gios.*",
     "homeassistant.components.glances.*",
     "homeassistant.components.gogogate2.*",
     "homeassistant.components.google_assistant.*",

tests/components/gios/__init__.py

@@ -17,7 +17,7 @@ async def init_integration(hass, incomplete_data=False) -> MockConfigEntry:
     entry = MockConfigEntry(
         domain=DOMAIN,
         title="Home",
-        unique_id=123,
+        unique_id="123",
         data={"station_id": 123, "name": "Home"},
     )
 

tests/components/gios/test_air_quality.py

@@ -13,9 +13,10 @@ from homeassistant.components.air_quality import (
     ATTR_PM_2_5,
     ATTR_PM_10,
     ATTR_SO2,
+    DOMAIN as AIR_QUALITY_DOMAIN,
 )
 from homeassistant.components.gios.air_quality import ATTRIBUTION
-from homeassistant.components.gios.const import AQI_GOOD
+from homeassistant.components.gios.const import AQI_GOOD, DOMAIN
 from homeassistant.const import (
     ATTR_ATTRIBUTION,
     ATTR_ICON,
@@ -55,7 +56,7 @@ async def test_air_quality(hass):
 
     entry = registry.async_get("air_quality.home")
     assert entry
-    assert entry.unique_id == 123
+    assert entry.unique_id == "123"
 
 
 async def test_air_quality_with_incomplete_data(hass):
@@ -83,7 +84,7 @@ async def test_air_quality_with_incomplete_data(hass):
 
     entry = registry.async_get("air_quality.home")
     assert entry
-    assert entry.unique_id == 123
+    assert entry.unique_id == "123"
 
 
 async def test_availability(hass):
@@ -122,3 +123,23 @@ async def test_availability(hass):
         assert state
         assert state.state != STATE_UNAVAILABLE
         assert state.state == "4"
+
+
+async def test_migrate_unique_id(hass):
+    """Test migrate unique_id of the air_quality entity."""
+    registry = er.async_get(hass)
+
+    # Pre-create registry entries for disabled by default sensors
+    registry.async_get_or_create(
+        AIR_QUALITY_DOMAIN,
+        DOMAIN,
+        123,
+        suggested_object_id="home",
+        disabled_by=None,
+    )
+
+    await init_integration(hass)
+
+    entry = registry.async_get("air_quality.home")
+    assert entry
+    assert entry.unique_id == "123"

tests/components/gios/test_config_flow.py

@@ -102,4 +102,4 @@ async def test_create_entry(hass):
         assert result["title"] == CONFIG[CONF_STATION_ID]
         assert result["data"][CONF_STATION_ID] == CONFIG[CONF_STATION_ID]
 
-        assert flow.context["unique_id"] == CONFIG[CONF_STATION_ID]
+        assert flow.context["unique_id"] == "123"

tests/components/gios/test_init.py

@@ -1,4 +1,5 @@
 """Test init of GIOS integration."""
+import json
 from unittest.mock import patch
 
 from homeassistant.components.gios.const import DOMAIN
@@ -9,7 +10,9 @@ from homeassistant.config_entries import (
 )
 from homeassistant.const import STATE_UNAVAILABLE
 
-from tests.common import MockConfigEntry
+from . import STATIONS
+
+from tests.common import MockConfigEntry, load_fixture, mock_device_registry
 from tests.components.gios import init_integration
 
 
@@ -53,3 +56,46 @@ async def test_unload_entry(hass):
 
     assert entry.state == ENTRY_STATE_NOT_LOADED
     assert not hass.data.get(DOMAIN)
+
+
+async def test_migrate_device_and_config_entry(hass):
+    """Test device_info identifiers and config entry migration."""
+    config_entry = MockConfigEntry(
+        domain=DOMAIN,
+        title="Home",
+        unique_id=123,
+        data={
+            "station_id": 123,
+            "name": "Home",
+        },
+    )
+
+    indexes = json.loads(load_fixture("gios/indexes.json"))
+    station = json.loads(load_fixture("gios/station.json"))
+    sensors = json.loads(load_fixture("gios/sensors.json"))
+
+    with patch(
+        "homeassistant.components.gios.Gios._get_stations", return_value=STATIONS
+    ), patch(
+        "homeassistant.components.gios.Gios._get_station",
+        return_value=station,
+    ), patch(
+        "homeassistant.components.gios.Gios._get_all_sensors",
+        return_value=sensors,
+    ), patch(
+        "homeassistant.components.gios.Gios._get_indexes", return_value=indexes
+    ):
+        config_entry.add_to_hass(hass)
+
+        device_reg = mock_device_registry(hass)
+        device_entry = device_reg.async_get_or_create(
+            config_entry_id=config_entry.entry_id, identifiers={(DOMAIN, 123)}
+        )
+
+        await hass.config_entries.async_setup(config_entry.entry_id)
+        await hass.async_block_till_done()
+
+        migrated_device_entry = device_reg.async_get_or_create(
+            config_entry_id=config_entry.entry_id, identifiers={(DOMAIN, "123")}
+        )
+        assert device_entry.id == migrated_device_entry.id