diff --git a/homeassistant/components/flume/binary_sensor.py b/homeassistant/components/flume/binary_sensor.py index 2f5ea4323d9..0ec3576b2f9 100644 --- a/homeassistant/components/flume/binary_sensor.py +++ b/homeassistant/components/flume/binary_sensor.py @@ -125,11 +125,12 @@ async def async_setup_entry( async_add_entities(flume_entity_list) -class FlumeNotificationBinarySensor(FlumeEntity, BinarySensorEntity): +class FlumeNotificationBinarySensor( + FlumeEntity[FlumeNotificationDataUpdateCoordinator], BinarySensorEntity +): """Binary sensor class.""" entity_description: FlumeBinarySensorEntityDescription - coordinator: FlumeNotificationDataUpdateCoordinator @property def is_on(self) -> bool: @@ -144,11 +145,12 @@ class FlumeNotificationBinarySensor(FlumeEntity, BinarySensorEntity): ) -class FlumeConnectionBinarySensor(FlumeEntity, BinarySensorEntity): +class FlumeConnectionBinarySensor( + FlumeEntity[FlumeDeviceConnectionUpdateCoordinator], BinarySensorEntity +): """Binary Sensor class for WIFI Connection status.""" entity_description: FlumeBinarySensorEntityDescription - coordinator: FlumeDeviceConnectionUpdateCoordinator _attr_entity_category = EntityCategory.DIAGNOSTIC _attr_device_class = BinarySensorDeviceClass.CONNECTIVITY diff --git a/homeassistant/components/flume/entity.py b/homeassistant/components/flume/entity.py index 7cd84127c64..b958510e733 100644 --- a/homeassistant/components/flume/entity.py +++ b/homeassistant/components/flume/entity.py @@ -1,16 +1,29 @@ """Platform for shared base classes for sensors.""" from __future__ import annotations +from typing import TypeVar, Union + from homeassistant.helpers.entity import DeviceInfo, EntityDescription -from homeassistant.helpers.update_coordinator import ( - CoordinatorEntity, - DataUpdateCoordinator, -) +from homeassistant.helpers.update_coordinator import CoordinatorEntity from .const import DOMAIN +from .coordinator import ( + FlumeDeviceConnectionUpdateCoordinator, + FlumeDeviceDataUpdateCoordinator, + FlumeNotificationDataUpdateCoordinator, +) + +_FlumeCoordinatorT = TypeVar( + "_FlumeCoordinatorT", + bound=Union[ + FlumeDeviceDataUpdateCoordinator, + FlumeDeviceConnectionUpdateCoordinator, + FlumeNotificationDataUpdateCoordinator, + ], +) -class FlumeEntity(CoordinatorEntity[DataUpdateCoordinator[None]]): +class FlumeEntity(CoordinatorEntity[_FlumeCoordinatorT]): """Base entity class.""" _attr_attribution = "Data provided by Flume API" @@ -18,7 +31,7 @@ class FlumeEntity(CoordinatorEntity[DataUpdateCoordinator[None]]): def __init__( self, - coordinator: DataUpdateCoordinator, + coordinator: _FlumeCoordinatorT, description: EntityDescription, device_id: str, location_name: str, diff --git a/homeassistant/components/flume/sensor.py b/homeassistant/components/flume/sensor.py index c7ebbd1f456..4b5fe5bc8e9 100644 --- a/homeassistant/components/flume/sensor.py +++ b/homeassistant/components/flume/sensor.py @@ -130,11 +130,9 @@ async def async_setup_entry( async_add_entities(flume_entity_list) -class FlumeSensor(FlumeEntity, SensorEntity): +class FlumeSensor(FlumeEntity[FlumeDeviceDataUpdateCoordinator], SensorEntity): """Representation of the Flume sensor.""" - coordinator: FlumeDeviceDataUpdateCoordinator - @property def native_value(self): """Return the state of the sensor."""