Update coordinator improvements (#38366)

* Make generic

* Add type info to bunch of uses

* Recognize requests exceptions

* Recognize urllib exceptions

homeassistant/components/cert_expiry/__init__.py

@@ -1,6 +1,7 @@
 """The cert_expiry component."""
-from datetime import timedelta
+from datetime import datetime, timedelta
 import logging
+from typing import Optional
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, CONF_PORT
@@ -50,7 +51,7 @@ async def async_unload_entry(hass, entry):
     return await hass.config_entries.async_forward_entry_unload(entry, "sensor")
 
 
-class CertExpiryDataUpdateCoordinator(DataUpdateCoordinator):
+class CertExpiryDataUpdateCoordinator(DataUpdateCoordinator[datetime]):
     """Class to manage fetching Cert Expiry data from single endpoint."""
 
     def __init__(self, hass, host, port):
@@ -67,7 +68,7 @@ class CertExpiryDataUpdateCoordinator(DataUpdateCoordinator):
             hass, _LOGGER, name=name, update_interval=SCAN_INTERVAL,
         )
 
-    async def _async_update_data(self):
+    async def _async_update_data(self) -> Optional[datetime]:
         """Fetch certificate."""
         try:
             timestamp = await get_cert_expiry_timestamp(self.hass, self.host, self.port)

homeassistant/components/guardian/util.py

@@ -14,7 +14,7 @@ from .const import LOGGER
 DEFAULT_UPDATE_INTERVAL = timedelta(seconds=30)
 
 
-class GuardianDataUpdateCoordinator(DataUpdateCoordinator):
+class GuardianDataUpdateCoordinator(DataUpdateCoordinator[dict]):
     """Define an extended DataUpdateCoordinator with some Guardian goodies."""
 
     def __init__(

homeassistant/components/ipp/__init__.py

@@ -86,7 +86,7 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return unload_ok
 
 
-class IPPDataUpdateCoordinator(DataUpdateCoordinator):
+class IPPDataUpdateCoordinator(DataUpdateCoordinator[IPPPrinter]):
     """Class to manage fetching IPP data from single endpoint."""
 
     def __init__(

homeassistant/components/roku/__init__.py

@@ -108,7 +108,7 @@ def roku_exception_handler(func):
     return handler
 
 
-class RokuDataUpdateCoordinator(DataUpdateCoordinator):
+class RokuDataUpdateCoordinator(DataUpdateCoordinator[Device]):
     """Class to manage fetching Roku data."""
 
     def __init__(

homeassistant/components/toon/coordinator.py

@@ -21,8 +21,8 @@ from .const import CONF_CLOUDHOOK_URL, DEFAULT_SCAN_INTERVAL, DOMAIN
 _LOGGER = logging.getLogger(__name__)
 
 
-class ToonDataUpdateCoordinator(DataUpdateCoordinator):
+class ToonDataUpdateCoordinator(DataUpdateCoordinator[Status]):
-    """Class to manage fetching WLED data from single endpoint."""
+    """Class to manage fetching Toon data from single endpoint."""
 
     def __init__(
         self, hass: HomeAssistant, *, entry: ConfigEntry, session: OAuth2Session

homeassistant/components/updater/__init__.py

@@ -64,7 +64,7 @@ async def async_setup(hass, config):
 
     include_components = conf.get(CONF_COMPONENT_REPORTING)
 
-    async def check_new_version():
+    async def check_new_version() -> Updater:
         """Check if a new version is available and report if one is."""
         newest, release_notes = await get_newest_version(
             hass, huuid, include_components
@@ -98,7 +98,7 @@ async def async_setup(hass, config):
 
         return Updater(update_available, newest, release_notes)
 
-    coordinator = hass.data[DOMAIN] = update_coordinator.DataUpdateCoordinator(
+    coordinator = hass.data[DOMAIN] = update_coordinator.DataUpdateCoordinator[Updater](
         hass,
         _LOGGER,
         name="Home Assistant update",

homeassistant/components/upnp/sensor.py

@@ -1,6 +1,6 @@
 """Support for UPnP/IGD Sensors."""
 from datetime import timedelta
-from typing import Mapping
+from typing import Any, Mapping
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import DATA_BYTES, DATA_RATE_KIBIBYTES_PER_SECOND
@@ -94,7 +94,7 @@ async def async_setup_entry(
     update_interval = timedelta(seconds=update_interval_sec)
     _LOGGER.debug("update_interval: %s", update_interval)
     _LOGGER.debug("Adding sensors")
-    coordinator = DataUpdateCoordinator(
+    coordinator = DataUpdateCoordinator[Mapping[str, Any]](
         hass,
         _LOGGER,
         name=device.name,
@@ -122,7 +122,7 @@ class UpnpSensor(Entity):
 
     def __init__(
         self,
-        coordinator: DataUpdateCoordinator,
+        coordinator: DataUpdateCoordinator[Mapping[str, Any]],
         device: Device,
         sensor_type: Mapping[str, str],
         update_multiplier: int = 2,
@@ -169,7 +169,7 @@ class UpnpSensor(Entity):
         return self._sensor_type["unit"]
 
     @property
-    def device_info(self) -> Mapping[str, any]:
+    def device_info(self) -> Mapping[str, Any]:
         """Get device info."""
         return {
             "connections": {(dr.CONNECTION_UPNP, self._device.udn)},

homeassistant/components/withings/common.py

@@ -582,7 +582,9 @@ class DataManager:
             update_interval=timedelta(minutes=120),
             update_method=self.async_subscribe_webhook,
         )
-        self.poll_data_update_coordinator = DataUpdateCoordinator(
+        self.poll_data_update_coordinator = DataUpdateCoordinator[
+            Dict[MeasureType, Any]
+        ](
             hass,
             _LOGGER,
             name="poll_data_update_coordinator",

homeassistant/components/wled/__init__.py

@@ -107,7 +107,7 @@ def wled_exception_handler(func):
     return handler
 
 
-class WLEDDataUpdateCoordinator(DataUpdateCoordinator):
+class WLEDDataUpdateCoordinator(DataUpdateCoordinator[WLEDDevice]):
     """Class to manage fetching WLED data from single endpoint."""
 
     def __init__(

homeassistant/helpers/update_coordinator.py

@@ -3,9 +3,11 @@ import asyncio
 from datetime import datetime, timedelta
 import logging
 from time import monotonic
-from typing import Any, Awaitable, Callable, List, Optional
+from typing import Awaitable, Callable, Generic, List, Optional, TypeVar
+import urllib.error
 
 import aiohttp
+import requests
 
 from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
 from homeassistant.helpers import event
@@ -16,12 +18,14 @@ from .debounce import Debouncer
 REQUEST_REFRESH_DEFAULT_COOLDOWN = 10
 REQUEST_REFRESH_DEFAULT_IMMEDIATE = True
 
+T = TypeVar("T")
+
 
 class UpdateFailed(Exception):
     """Raised when an update has failed."""
 
 
-class DataUpdateCoordinator:
+class DataUpdateCoordinator(Generic[T]):
     """Class to manage fetching data from single endpoint."""
 
     def __init__(
@@ -31,7 +35,7 @@ class DataUpdateCoordinator:
         *,
         name: str,
         update_interval: Optional[timedelta] = None,
-        update_method: Optional[Callable[[], Awaitable]] = None,
+        update_method: Optional[Callable[[], Awaitable[T]]] = None,
         request_refresh_debouncer: Optional[Debouncer] = None,
     ):
         """Initialize global data updater."""
@@ -41,7 +45,7 @@ class DataUpdateCoordinator:
         self.update_method = update_method
         self.update_interval = update_interval
 
-        self.data: Optional[Any] = None
+        self.data: Optional[T] = None
 
         self._listeners: List[CALLBACK_TYPE] = []
         self._unsub_refresh: Optional[CALLBACK_TYPE] = None
@@ -120,7 +124,7 @@ class DataUpdateCoordinator:
         """
         await self._debounced_refresh.async_call()
 
-    async def _async_update_data(self) -> Optional[Any]:
+    async def _async_update_data(self) -> Optional[T]:
         """Fetch the latest data from the source."""
         if self.update_method is None:
             raise NotImplementedError("Update method not implemented")
@@ -138,16 +142,24 @@ class DataUpdateCoordinator:
             start = monotonic()
             self.data = await self._async_update_data()
 
-        except asyncio.TimeoutError:
+        except (asyncio.TimeoutError, requests.exceptions.Timeout):
             if self.last_update_success:
                 self.logger.error("Timeout fetching %s data", self.name)
                 self.last_update_success = False
 
-        except aiohttp.ClientError as err:
+        except (aiohttp.ClientError, requests.exceptions.RequestException) as err:
             if self.last_update_success:
                 self.logger.error("Error requesting %s data: %s", self.name, err)
                 self.last_update_success = False
 
+        except urllib.error.URLError as err:
+            if self.last_update_success:
+                if err.reason == "timed out":
+                    self.logger.error("Timeout fetching %s data", self.name)
+                else:
+                    self.logger.error("Error requesting %s data: %s", self.name, err)
+                self.last_update_success = False
+
         except UpdateFailed as err:
             if self.last_update_success:
                 self.logger.error("Error fetching %s data: %s", self.name, err)

tests/helpers/test_update_coordinator.py

@@ -2,9 +2,11 @@
 import asyncio
 from datetime import timedelta
 import logging
+import urllib.error
 
 import aiohttp
 import pytest
+import requests
 
 from homeassistant.helpers import update_coordinator
 from homeassistant.util.dt import utcnow
@@ -19,12 +21,12 @@ def get_crd(hass, update_interval):
     """Make coordinator mocks."""
     calls = 0
 
-    async def refresh():
+    async def refresh() -> int:
         nonlocal calls
         calls += 1
         return calls
 
-    crd = update_coordinator.DataUpdateCoordinator(
+    crd = update_coordinator.DataUpdateCoordinator[int](
         hass,
         LOGGER,
         name="test",
@@ -111,7 +113,11 @@ async def test_request_refresh_no_auto_update(crd_without_update_interval):
     "err_msg",
     [
         (asyncio.TimeoutError, "Timeout fetching test data"),
+        (requests.exceptions.Timeout, "Timeout fetching test data"),
+        (urllib.error.URLError("timed out"), "Timeout fetching test data"),
         (aiohttp.ClientError, "Error requesting test data"),
+        (requests.exceptions.RequestException, "Error requesting test data"),
+        (urllib.error.URLError("something"), "Error requesting test data"),
         (update_coordinator.UpdateFailed, "Error fetching test data"),
     ],
 )