mirror of
https://github.com/home-assistant/core.git
synced 2025-07-25 14:17:45 +00:00
Improve Transmission Entity description class (#105924)
* Add entity mixin for transmission sensors * use kw_only in EntityDescription class * minor fix * Update sensor.py uom
This commit is contained in:
Rami Mosleh
GitHub
parent
5f2a13fec6
commit
ab40ba80a9
@ -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}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user