Improve Transmission Entity description class (#105924)

* Add entity mixin for transmission sensors

* use kw_only in EntityDescription class

* minor fix

* Update sensor.py uom

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, ...] = (
+MODES: dict[str, list[str] | None] = {
-    SensorEntityDescription(key="download", translation_key="download_speed"),
+    "started_torrents": ["downloading"],
-    SensorEntityDescription(key="upload", translation_key="upload_speed"),
+    "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, ...] = (
+@dataclass(kw_only=True)
-    SensorEntityDescription(key="active_torrents", translation_key="active_torrents"),
+class TransmissionSensorEntityDescription(SensorEntityDescription):
-    SensorEntityDescription(key="paused_torrents", translation_key="paused_torrents"),
+    """Entity description class for Transmission sensors."""
-    SensorEntityDescription(key="total_torrents", translation_key="total_torrents"),
+
-    SensorEntityDescription(
+    val_func: Callable[[TransmissionDataUpdateCoordinator], StateType]
-        key="completed_torrents", translation_key="completed_torrents"
+    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] = []
+    async_add_entities(
-
+        TransmissionSensor(coordinator, description) for description in SENSOR_TYPES
-    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)
 
 
 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,37 +172,21 @@ class TransmissionSensor(
             manufacturer="Transmission",
         )
 
-
+    @property
-class TransmissionSpeedSensor(TransmissionSensor):
+    def native_value(self) -> StateType:
-    """Representation of a Transmission speed sensor."""
+        """Return the value of the sensor."""
-
+        return self.entity_description.val_func(self.coordinator)
-    _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) -> float:
+    def extra_state_attributes(self) -> dict[str, Any] | None:
-        """Return the speed of the sensor."""
+        """Return the state attributes, if any."""
-        data = self.coordinator.data
+        if attr_func := self.entity_description.extra_state_attr_func:
-        return (
+            return attr_func(self.coordinator)
-            float(data.download_speed)
+        return None
-            if self.entity_description.key == "download"
-            else float(data.upload_speed)
-        )
 
 
-class TransmissionStatusSensor(TransmissionSensor):
+def get_state(upload: int, download: int) -> str:
-    """Representation of a Transmission status sensor."""
+    """Get current download/upload state."""
-
-    _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:
@@ -141,47 +196,6 @@ class TransmissionStatusSensor(TransmissionSensor):
     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]:
-        """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,
-        }
-
-    @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 _filter_torrents(
     torrents: list[Torrent], statuses: list[str] | None = None
 ) -> list[Torrent]:
@@ -192,13 +206,13 @@ def _filter_torrents(
     ]
 
 
-def _torrents_info(
+def _torrents_info_attr(
-    torrents: list[Torrent], order: str, limit: int, statuses: list[str] | None = None
+    coordinator: TransmissionDataUpdateCoordinator, key: str
 ) -> dict[str, Any]:
     infos = {}
-    torrents = _filter_torrents(torrents, statuses)
+    torrents = _filter_torrents(coordinator.torrents, MODES[key])
-    torrents = SUPPORTED_ORDER_MODES[order](torrents)
+    torrents = SUPPORTED_ORDER_MODES[coordinator.order](torrents)
-    for torrent in torrents[:limit]:
+    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}