diff --git a/homeassistant/components/nextdns/__init__.py b/homeassistant/components/nextdns/__init__.py index f9cb1f60cd1..7e7f5ff2dd8 100644 --- a/homeassistant/components/nextdns/__init__.py +++ b/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: diff --git a/homeassistant/components/nextdns/sensor.py b/homeassistant/components/nextdns/sensor.py index ef0176b071a..71d71ff287b 100644 --- a/homeassistant/components/nextdns/sensor.py +++ b/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."""