Use generics in NextDNS (#74517)

Use generics in nextdns

homeassistant/components/nextdns/__init__.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 import asyncio
 from datetime import timedelta
 import logging
+from typing import TypeVar
 
 from aiohttp.client_exceptions import ClientConnectorError
 from async_timeout import timeout
@@ -39,8 +40,10 @@ from .const import (
     UPDATE_INTERVAL_ANALYTICS,
 )
 
+TCoordinatorData = TypeVar("TCoordinatorData", bound=NextDnsData)
 
-class NextDnsUpdateCoordinator(DataUpdateCoordinator):
+
+class NextDnsUpdateCoordinator(DataUpdateCoordinator[TCoordinatorData]):
     """Class to manage fetching NextDNS data API."""
 
     def __init__(
@@ -64,12 +67,12 @@ class NextDnsUpdateCoordinator(DataUpdateCoordinator):
 
         super().__init__(hass, _LOGGER, name=DOMAIN, update_interval=update_interval)
 
-    async def _async_update_data(self) -> NextDnsData:
+    async def _async_update_data(self) -> TCoordinatorData:
         """Update data via library."""
         raise NotImplementedError("Update method not implemented")
 
 
-class NextDnsStatusUpdateCoordinator(NextDnsUpdateCoordinator):
+class NextDnsStatusUpdateCoordinator(NextDnsUpdateCoordinator[AnalyticsStatus]):
     """Class to manage fetching NextDNS analytics status data from API."""
 
     async def _async_update_data(self) -> AnalyticsStatus:
@@ -81,7 +84,7 @@ class NextDnsStatusUpdateCoordinator(NextDnsUpdateCoordinator):
             raise UpdateFailed(err) from err
 
 
-class NextDnsDnssecUpdateCoordinator(NextDnsUpdateCoordinator):
+class NextDnsDnssecUpdateCoordinator(NextDnsUpdateCoordinator[AnalyticsDnssec]):
     """Class to manage fetching NextDNS analytics Dnssec data from API."""
 
     async def _async_update_data(self) -> AnalyticsDnssec:
@@ -93,7 +96,7 @@ class NextDnsDnssecUpdateCoordinator(NextDnsUpdateCoordinator):
             raise UpdateFailed(err) from err
 
 
-class NextDnsEncryptionUpdateCoordinator(NextDnsUpdateCoordinator):
+class NextDnsEncryptionUpdateCoordinator(NextDnsUpdateCoordinator[AnalyticsEncryption]):
     """Class to manage fetching NextDNS analytics encryption data from API."""
 
     async def _async_update_data(self) -> AnalyticsEncryption:
@@ -105,7 +108,7 @@ class NextDnsEncryptionUpdateCoordinator(NextDnsUpdateCoordinator):
             raise UpdateFailed(err) from err
 
 
-class NextDnsIpVersionsUpdateCoordinator(NextDnsUpdateCoordinator):
+class NextDnsIpVersionsUpdateCoordinator(NextDnsUpdateCoordinator[AnalyticsIpVersions]):
     """Class to manage fetching NextDNS analytics IP versions data from API."""
 
     async def _async_update_data(self) -> AnalyticsIpVersions:
@@ -117,7 +120,7 @@ class NextDnsIpVersionsUpdateCoordinator(NextDnsUpdateCoordinator):
             raise UpdateFailed(err) from err
 
 
-class NextDnsProtocolsUpdateCoordinator(NextDnsUpdateCoordinator):
+class NextDnsProtocolsUpdateCoordinator(NextDnsUpdateCoordinator[AnalyticsProtocols]):
     """Class to manage fetching NextDNS analytics protocols data from API."""
 
     async def _async_update_data(self) -> AnalyticsProtocols:

homeassistant/components/nextdns/sensor.py

@@ -3,7 +3,15 @@ from __future__ import annotations
 
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, cast
+from typing import Generic, cast
+
+from nextdns import (
+    AnalyticsDnssec,
+    AnalyticsEncryption,
+    AnalyticsIpVersions,
+    AnalyticsProtocols,
+    AnalyticsStatus,
+)
 
 from homeassistant.components.sensor import (
     SensorEntity,
@@ -18,7 +26,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from . import NextDnsUpdateCoordinator
+from . import NextDnsUpdateCoordinator, TCoordinatorData
 from .const import (
     ATTR_DNSSEC,
     ATTR_ENCRYPTION,
@@ -32,23 +40,26 @@ PARALLEL_UPDATES = 1
 
 
 @dataclass
-class NextDnsSensorRequiredKeysMixin:
+class NextDnsSensorRequiredKeysMixin(Generic[TCoordinatorData]):
     """Class for NextDNS entity required keys."""
 
+    coordinator_class: type[TCoordinatorData]
     coordinator_type: str
-    value: Callable[[Any], StateType]
+    value: Callable[[TCoordinatorData], StateType]
 
 
 @dataclass
 class NextDnsSensorEntityDescription(
-    SensorEntityDescription, NextDnsSensorRequiredKeysMixin
+    SensorEntityDescription,
+    NextDnsSensorRequiredKeysMixin[TCoordinatorData],
 ):
     """NextDNS sensor entity description."""
 
 
-SENSORS = (
+SENSORS: tuple[NextDnsSensorEntityDescription, ...] = (
     NextDnsSensorEntityDescription(
         key="all_queries",
+        coordinator_class=AnalyticsStatus,
         coordinator_type=ATTR_STATUS,
         entity_category=EntityCategory.DIAGNOSTIC,
         icon="mdi:dns",
@@ -59,6 +70,7 @@ SENSORS = (
     ),
     NextDnsSensorEntityDescription(
         key="blocked_queries",
+        coordinator_class=AnalyticsStatus,
         coordinator_type=ATTR_STATUS,
         entity_category=EntityCategory.DIAGNOSTIC,
         icon="mdi:dns",
@@ -69,6 +81,7 @@ SENSORS = (
     ),
     NextDnsSensorEntityDescription(
         key="relayed_queries",
+        coordinator_class=AnalyticsStatus,
         coordinator_type=ATTR_STATUS,
         entity_category=EntityCategory.DIAGNOSTIC,
         icon="mdi:dns",
@@ -79,6 +92,7 @@ SENSORS = (
     ),
     NextDnsSensorEntityDescription(
         key="blocked_queries_ratio",
+        coordinator_class=AnalyticsStatus,
         coordinator_type=ATTR_STATUS,
         entity_category=EntityCategory.DIAGNOSTIC,
         icon="mdi:dns",
@@ -89,6 +103,7 @@ SENSORS = (
     ),
     NextDnsSensorEntityDescription(
         key="doh_queries",
+        coordinator_class=AnalyticsProtocols,
         coordinator_type=ATTR_PROTOCOLS,
         entity_category=EntityCategory.DIAGNOSTIC,
         entity_registry_enabled_default=False,
@@ -100,6 +115,7 @@ SENSORS = (
     ),
     NextDnsSensorEntityDescription(
         key="dot_queries",
+        coordinator_class=AnalyticsProtocols,
         coordinator_type=ATTR_PROTOCOLS,
         entity_category=EntityCategory.DIAGNOSTIC,
         entity_registry_enabled_default=False,
@@ -111,6 +127,7 @@ SENSORS = (
     ),
     NextDnsSensorEntityDescription(
         key="doq_queries",
+        coordinator_class=AnalyticsProtocols,
         coordinator_type=ATTR_PROTOCOLS,
         entity_category=EntityCategory.DIAGNOSTIC,
         entity_registry_enabled_default=False,
@@ -122,6 +139,7 @@ SENSORS = (
     ),
     NextDnsSensorEntityDescription(
         key="udp_queries",
+        coordinator_class=AnalyticsProtocols,
         coordinator_type=ATTR_PROTOCOLS,
         entity_category=EntityCategory.DIAGNOSTIC,
         entity_registry_enabled_default=False,
@@ -133,6 +151,7 @@ SENSORS = (
     ),
     NextDnsSensorEntityDescription(
         key="doh_queries_ratio",
+        coordinator_class=AnalyticsProtocols,
         coordinator_type=ATTR_PROTOCOLS,
         entity_registry_enabled_default=False,
         icon="mdi:dns",
@@ -144,6 +163,7 @@ SENSORS = (
     ),
     NextDnsSensorEntityDescription(
         key="dot_queries_ratio",
+        coordinator_class=AnalyticsProtocols,
         coordinator_type=ATTR_PROTOCOLS,
         entity_category=EntityCategory.DIAGNOSTIC,
         entity_registry_enabled_default=False,
@@ -155,6 +175,7 @@ SENSORS = (
     ),
     NextDnsSensorEntityDescription(
         key="doq_queries_ratio",
+        coordinator_class=AnalyticsProtocols,
         coordinator_type=ATTR_PROTOCOLS,
         entity_registry_enabled_default=False,
         icon="mdi:dns",
@@ -166,6 +187,7 @@ SENSORS = (
     ),
     NextDnsSensorEntityDescription(
         key="udp_queries_ratio",
+        coordinator_class=AnalyticsProtocols,
         coordinator_type=ATTR_PROTOCOLS,
         entity_category=EntityCategory.DIAGNOSTIC,
         entity_registry_enabled_default=False,
@@ -177,6 +199,7 @@ SENSORS = (
     ),
     NextDnsSensorEntityDescription(
         key="encrypted_queries",
+        coordinator_class=AnalyticsEncryption,
         coordinator_type=ATTR_ENCRYPTION,
         entity_category=EntityCategory.DIAGNOSTIC,
         entity_registry_enabled_default=False,
@@ -188,6 +211,7 @@ SENSORS = (
     ),
     NextDnsSensorEntityDescription(
         key="unencrypted_queries",
+        coordinator_class=AnalyticsEncryption,
         coordinator_type=ATTR_ENCRYPTION,
         entity_category=EntityCategory.DIAGNOSTIC,
         entity_registry_enabled_default=False,
@@ -199,6 +223,7 @@ SENSORS = (
     ),
     NextDnsSensorEntityDescription(
         key="encrypted_queries_ratio",
+        coordinator_class=AnalyticsEncryption,
         coordinator_type=ATTR_ENCRYPTION,
         entity_category=EntityCategory.DIAGNOSTIC,
         entity_registry_enabled_default=False,
@@ -210,6 +235,7 @@ SENSORS = (
     ),
     NextDnsSensorEntityDescription(
         key="ipv4_queries",
+        coordinator_class=AnalyticsIpVersions,
         coordinator_type=ATTR_IP_VERSIONS,
         entity_category=EntityCategory.DIAGNOSTIC,
         entity_registry_enabled_default=False,
@@ -221,6 +247,7 @@ SENSORS = (
     ),
     NextDnsSensorEntityDescription(
         key="ipv6_queries",
+        coordinator_class=AnalyticsIpVersions,
         coordinator_type=ATTR_IP_VERSIONS,
         entity_category=EntityCategory.DIAGNOSTIC,
         entity_registry_enabled_default=False,
@@ -232,6 +259,7 @@ SENSORS = (
     ),
     NextDnsSensorEntityDescription(
         key="ipv6_queries_ratio",
+        coordinator_class=AnalyticsIpVersions,
         coordinator_type=ATTR_IP_VERSIONS,
         entity_category=EntityCategory.DIAGNOSTIC,
         entity_registry_enabled_default=False,
@@ -243,6 +271,7 @@ SENSORS = (
     ),
     NextDnsSensorEntityDescription(
         key="validated_queries",
+        coordinator_class=AnalyticsDnssec,
         coordinator_type=ATTR_DNSSEC,
         entity_category=EntityCategory.DIAGNOSTIC,
         entity_registry_enabled_default=False,
@@ -254,6 +283,7 @@ SENSORS = (
     ),
     NextDnsSensorEntityDescription(
         key="not_validated_queries",
+        coordinator_class=AnalyticsDnssec,
         coordinator_type=ATTR_DNSSEC,
         entity_category=EntityCategory.DIAGNOSTIC,
         entity_registry_enabled_default=False,
@@ -265,6 +295,7 @@ SENSORS = (
     ),
     NextDnsSensorEntityDescription(
         key="validated_queries_ratio",
+        coordinator_class=AnalyticsDnssec,
         coordinator_type=ATTR_DNSSEC,
         entity_category=EntityCategory.DIAGNOSTIC,
         entity_registry_enabled_default=False,
@@ -294,12 +325,14 @@ async def async_setup_entry(
     async_add_entities(sensors)
 
 
-class NextDnsSensor(CoordinatorEntity[NextDnsUpdateCoordinator], SensorEntity):
+class NextDnsSensor(
+    CoordinatorEntity[NextDnsUpdateCoordinator[TCoordinatorData]], SensorEntity
+):
     """Define an NextDNS sensor."""
 
     def __init__(
         self,
-        coordinator: NextDnsUpdateCoordinator,
+        coordinator: NextDnsUpdateCoordinator[TCoordinatorData],
         description: NextDnsSensorEntityDescription,
     ) -> None:
         """Initialize."""