From ab40ba80a936cdb92e596131f1b40d11e9726f83 Mon Sep 17 00:00:00 2001 From: Rami Mosleh Date: Mon, 18 Dec 2023 21:09:01 +0300 Subject: [PATCH] Improve Transmission Entity description class (#105924) * Add entity mixin for transmission sensors * use kw_only in EntityDescription class * minor fix * Update sensor.py uom --- .../components/transmission/sensor.py | 238 +++++++++--------- 1 file changed, 126 insertions(+), 112 deletions(-) diff --git a/homeassistant/components/transmission/sensor.py b/homeassistant/components/transmission/sensor.py index 20f4fc95c87..29d37f28bad 100644 --- a/homeassistant/components/transmission/sensor.py +++ b/homeassistant/components/transmission/sensor.py @@ -1,7 +1,9 @@ """Support for monitoring the Transmission BitTorrent client API.""" from __future__ import annotations +from collections.abc import Callable from contextlib import suppress +from dataclasses import dataclass from typing import Any from transmission_rpc.torrent import Torrent @@ -16,6 +18,7 @@ from homeassistant.const import STATE_IDLE, UnitOfDataRate from homeassistant.core import HomeAssistant from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo from homeassistant.helpers.entity_platform import AddEntitiesCallback +from homeassistant.helpers.typing import StateType from homeassistant.helpers.update_coordinator import CoordinatorEntity from .const import ( @@ -28,23 +31,103 @@ from .const import ( ) from .coordinator import TransmissionDataUpdateCoordinator -SPEED_SENSORS: tuple[SensorEntityDescription, ...] = ( - SensorEntityDescription(key="download", translation_key="download_speed"), - SensorEntityDescription(key="upload", translation_key="upload_speed"), -) +MODES: dict[str, list[str] | None] = { + "started_torrents": ["downloading"], + "completed_torrents": ["seeding"], + "paused_torrents": ["stopped"], + "active_torrents": [ + "seeding", + "downloading", + ], + "total_torrents": None, +} -STATUS_SENSORS: tuple[SensorEntityDescription, ...] = ( - SensorEntityDescription(key="status", translation_key="transmission_status"), -) -TORRENT_SENSORS: tuple[SensorEntityDescription, ...] = ( - SensorEntityDescription(key="active_torrents", translation_key="active_torrents"), - SensorEntityDescription(key="paused_torrents", translation_key="paused_torrents"), - SensorEntityDescription(key="total_torrents", translation_key="total_torrents"), - SensorEntityDescription( - key="completed_torrents", translation_key="completed_torrents" +@dataclass(kw_only=True) +class TransmissionSensorEntityDescription(SensorEntityDescription): + """Entity description class for Transmission sensors.""" + + val_func: Callable[[TransmissionDataUpdateCoordinator], StateType] + extra_state_attr_func: Callable[[Any], dict[str, str]] | None = None + + +SENSOR_TYPES: tuple[TransmissionSensorEntityDescription, ...] = ( + TransmissionSensorEntityDescription( + key="download", + translation_key="download_speed", + device_class=SensorDeviceClass.DATA_RATE, + native_unit_of_measurement=UnitOfDataRate.BYTES_PER_SECOND, + suggested_display_precision=2, + suggested_unit_of_measurement=UnitOfDataRate.MEGABYTES_PER_SECOND, + val_func=lambda coordinator: float(coordinator.data.download_speed), + ), + TransmissionSensorEntityDescription( + key="upload", + translation_key="upload_speed", + device_class=SensorDeviceClass.DATA_RATE, + native_unit_of_measurement=UnitOfDataRate.BYTES_PER_SECOND, + suggested_display_precision=2, + suggested_unit_of_measurement=UnitOfDataRate.MEGABYTES_PER_SECOND, + val_func=lambda coordinator: float(coordinator.data.upload_speed), + ), + TransmissionSensorEntityDescription( + key="status", + translation_key="transmission_status", + device_class=SensorDeviceClass.ENUM, + options=[STATE_IDLE, STATE_UP_DOWN, STATE_SEEDING, STATE_DOWNLOADING], + val_func=lambda coordinator: get_state( + coordinator.data.upload_speed, coordinator.data.download_speed + ), + ), + TransmissionSensorEntityDescription( + key="active_torrents", + translation_key="active_torrents", + native_unit_of_measurement="torrents", + val_func=lambda coordinator: coordinator.data.active_torrent_count, + extra_state_attr_func=lambda coordinator: _torrents_info_attr( + coordinator=coordinator, key="active_torrents" + ), + ), + TransmissionSensorEntityDescription( + key="paused_torrents", + translation_key="paused_torrents", + native_unit_of_measurement="torrents", + val_func=lambda coordinator: coordinator.data.paused_torrent_count, + extra_state_attr_func=lambda coordinator: _torrents_info_attr( + coordinator=coordinator, key="paused_torrents" + ), + ), + TransmissionSensorEntityDescription( + key="total_torrents", + translation_key="total_torrents", + native_unit_of_measurement="torrents", + val_func=lambda coordinator: coordinator.data.torrent_count, + extra_state_attr_func=lambda coordinator: _torrents_info_attr( + coordinator=coordinator, key="total_torrents" + ), + ), + TransmissionSensorEntityDescription( + key="completed_torrents", + translation_key="completed_torrents", + native_unit_of_measurement="torrents", + val_func=lambda coordinator: len( + _filter_torrents(coordinator.torrents, MODES["completed_torrents"]) + ), + extra_state_attr_func=lambda coordinator: _torrents_info_attr( + coordinator=coordinator, key="completed_torrents" + ), + ), + TransmissionSensorEntityDescription( + key="started_torrents", + translation_key="started_torrents", + native_unit_of_measurement="torrents", + val_func=lambda coordinator: len( + _filter_torrents(coordinator.torrents, MODES["started_torrents"]) + ), + extra_state_attr_func=lambda coordinator: _torrents_info_attr( + coordinator=coordinator, key="started_torrents" + ), ), - SensorEntityDescription(key="started_torrents", translation_key="started_torrents"), ) @@ -59,22 +142,9 @@ async def async_setup_entry( config_entry.entry_id ] - entities: list[TransmissionSensor] = [] - - entities = [ - TransmissionSpeedSensor(coordinator, description) - for description in SPEED_SENSORS - ] - entities += [ - TransmissionStatusSensor(coordinator, description) - for description in STATUS_SENSORS - ] - entities += [ - TransmissionTorrentsSensor(coordinator, description) - for description in TORRENT_SENSORS - ] - - async_add_entities(entities) + async_add_entities( + TransmissionSensor(coordinator, description) for description in SENSOR_TYPES + ) class TransmissionSensor( @@ -82,12 +152,13 @@ class TransmissionSensor( ): """A base class for all Transmission sensors.""" + entity_description: TransmissionSensorEntityDescription _attr_has_entity_name = True def __init__( self, coordinator: TransmissionDataUpdateCoordinator, - entity_description: SensorEntityDescription, + entity_description: TransmissionSensorEntityDescription, ) -> None: """Initialize the sensor.""" super().__init__(coordinator) @@ -101,85 +172,28 @@ class TransmissionSensor( manufacturer="Transmission", ) - -class TransmissionSpeedSensor(TransmissionSensor): - """Representation of a Transmission speed sensor.""" - - _attr_device_class = SensorDeviceClass.DATA_RATE - _attr_native_unit_of_measurement = UnitOfDataRate.BYTES_PER_SECOND - _attr_suggested_display_precision = 2 - _attr_suggested_unit_of_measurement = UnitOfDataRate.MEGABYTES_PER_SECOND + @property + def native_value(self) -> StateType: + """Return the value of the sensor.""" + return self.entity_description.val_func(self.coordinator) @property - def native_value(self) -> float: - """Return the speed of the sensor.""" - data = self.coordinator.data - return ( - float(data.download_speed) - if self.entity_description.key == "download" - else float(data.upload_speed) - ) - - -class TransmissionStatusSensor(TransmissionSensor): - """Representation of a Transmission status sensor.""" - - _attr_device_class = SensorDeviceClass.ENUM - _attr_options = [STATE_IDLE, STATE_UP_DOWN, STATE_SEEDING, STATE_DOWNLOADING] - - @property - def native_value(self) -> str: - """Return the value of the status sensor.""" - upload = self.coordinator.data.upload_speed - download = self.coordinator.data.download_speed - if upload > 0 and download > 0: - return STATE_UP_DOWN - if upload > 0 and download == 0: - return STATE_SEEDING - if upload == 0 and download > 0: - return STATE_DOWNLOADING - return STATE_IDLE - - -class TransmissionTorrentsSensor(TransmissionSensor): - """Representation of a Transmission torrents sensor.""" - - MODES: dict[str, list[str] | None] = { - "started_torrents": ["downloading"], - "completed_torrents": ["seeding"], - "paused_torrents": ["stopped"], - "active_torrents": [ - "seeding", - "downloading", - ], - "total_torrents": None, - } - - @property - def native_unit_of_measurement(self) -> str: - """Return the unit of measurement of this entity, if any.""" - return "Torrents" - - @property - def extra_state_attributes(self) -> dict[str, Any]: + def extra_state_attributes(self) -> dict[str, Any] | None: """Return the state attributes, if any.""" - info = _torrents_info( - torrents=self.coordinator.torrents, - order=self.coordinator.order, - limit=self.coordinator.limit, - statuses=self.MODES[self.entity_description.key], - ) - return { - STATE_ATTR_TORRENT_INFO: info, - } + if attr_func := self.entity_description.extra_state_attr_func: + return attr_func(self.coordinator) + return None - @property - def native_value(self) -> int: - """Return the count of the sensor.""" - torrents = _filter_torrents( - self.coordinator.torrents, statuses=self.MODES[self.entity_description.key] - ) - return len(torrents) + +def get_state(upload: int, download: int) -> str: + """Get current download/upload state.""" + if upload > 0 and download > 0: + return STATE_UP_DOWN + if upload > 0 and download == 0: + return STATE_SEEDING + if upload == 0 and download > 0: + return STATE_DOWNLOADING + return STATE_IDLE def _filter_torrents( @@ -192,13 +206,13 @@ def _filter_torrents( ] -def _torrents_info( - torrents: list[Torrent], order: str, limit: int, statuses: list[str] | None = None +def _torrents_info_attr( + coordinator: TransmissionDataUpdateCoordinator, key: str ) -> dict[str, Any]: infos = {} - torrents = _filter_torrents(torrents, statuses) - torrents = SUPPORTED_ORDER_MODES[order](torrents) - for torrent in torrents[:limit]: + torrents = _filter_torrents(coordinator.torrents, MODES[key]) + torrents = SUPPORTED_ORDER_MODES[coordinator.order](torrents) + for torrent in torrents[: coordinator.limit]: info = infos[torrent.name] = { "added_date": torrent.added_date, "percent_done": f"{torrent.percent_done * 100:.2f}", @@ -207,4 +221,4 @@ def _torrents_info( } with suppress(ValueError): info["eta"] = str(torrent.eta) - return infos + return {STATE_ATTR_TORRENT_INFO: infos}