Use EntityDescription for Transmission entities (#103581)

pull/103628/head
Rami Mosleh 2023-11-08 01:05:17 +02:00 committed by GitHub
parent 36011d0384
commit 77a2f1664e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 73 additions and 73 deletions

View File

@ -7,8 +7,6 @@ from transmission_rpc import Torrent
DOMAIN = "transmission"
SWITCH_TYPES = {"on_off": "Switch", "turtle_mode": "Turtle mode"}
ORDER_NEWEST_FIRST = "newest_first"
ORDER_OLDEST_FIRST = "oldest_first"
ORDER_BEST_RATIO_FIRST = "best_ratio_first"

View File

@ -6,7 +6,11 @@ from typing import Any
from transmission_rpc.torrent import Torrent
from homeassistant.components.sensor import SensorDeviceClass, SensorEntity
from homeassistant.components.sensor import (
SensorDeviceClass,
SensorEntity,
SensorEntityDescription,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import STATE_IDLE, UnitOfDataRate
from homeassistant.core import HomeAssistant
@ -24,6 +28,25 @@ 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"),
)
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"
),
SensorEntityDescription(key="started_torrents", translation_key="started_torrents"),
)
async def async_setup_entry(
hass: HomeAssistant,
@ -36,47 +59,19 @@ async def async_setup_entry(
config_entry.entry_id
]
entities: list[TransmissionSensor] = []
entities = [
TransmissionSpeedSensor(
coordinator,
"download_speed",
"download",
),
TransmissionSpeedSensor(
coordinator,
"upload_speed",
"upload",
),
TransmissionStatusSensor(
coordinator,
"transmission_status",
"status",
),
TransmissionTorrentsSensor(
coordinator,
"active_torrents",
"active_torrents",
),
TransmissionTorrentsSensor(
coordinator,
"paused_torrents",
"paused_torrents",
),
TransmissionTorrentsSensor(
coordinator,
"total_torrents",
"total_torrents",
),
TransmissionTorrentsSensor(
coordinator,
"completed_torrents",
"completed_torrents",
),
TransmissionTorrentsSensor(
coordinator,
"started_torrents",
"started_torrents",
),
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)
@ -88,19 +83,18 @@ class TransmissionSensor(
"""A base class for all Transmission sensors."""
_attr_has_entity_name = True
_attr_should_poll = False
def __init__(
self,
coordinator: TransmissionDataUpdateCoordinator,
sensor_translation_key: str,
key: str,
entity_description: SensorEntityDescription,
) -> None:
"""Initialize the sensor."""
super().__init__(coordinator)
self._attr_translation_key = sensor_translation_key
self._key = key
self._attr_unique_id = f"{coordinator.config_entry.entry_id}-{key}"
self.entity_description = entity_description
self._attr_unique_id = (
f"{coordinator.config_entry.entry_id}-{entity_description.key}"
)
self._attr_device_info = DeviceInfo(
entry_type=DeviceEntryType.SERVICE,
identifiers={(DOMAIN, coordinator.config_entry.entry_id)},
@ -122,7 +116,7 @@ class TransmissionSpeedSensor(TransmissionSensor):
data = self.coordinator.data
return (
float(data.download_speed)
if self._key == "download"
if self.entity_description.key == "download"
else float(data.upload_speed)
)
@ -173,7 +167,7 @@ class TransmissionTorrentsSensor(TransmissionSensor):
torrents=self.coordinator.torrents,
order=self.coordinator.order,
limit=self.coordinator.limit,
statuses=self.MODES[self._key],
statuses=self.MODES[self.entity_description.key],
)
return {
STATE_ATTR_TORRENT_INFO: info,
@ -183,7 +177,7 @@ class TransmissionTorrentsSensor(TransmissionSensor):
def native_value(self) -> int:
"""Return the count of the sensor."""
torrents = _filter_torrents(
self.coordinator.torrents, statuses=self.MODES[self._key]
self.coordinator.torrents, statuses=self.MODES[self.entity_description.key]
)
return len(torrents)

View File

@ -69,6 +69,14 @@
"started_torrents": {
"name": "Started torrents"
}
},
"switch": {
"on_off": {
"name": "Switch"
},
"turtle_mode": {
"name": "Turtle mode"
}
}
},
"services": {

View File

@ -1,20 +1,24 @@
"""Support for setting the Transmission BitTorrent client Turtle Mode."""
from collections.abc import Callable
import logging
from typing import Any
from homeassistant.components.switch import SwitchEntity
from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant
from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.update_coordinator import CoordinatorEntity
from .const import DOMAIN, SWITCH_TYPES
from .const import DOMAIN
from .coordinator import TransmissionDataUpdateCoordinator
_LOGGING = logging.getLogger(__name__)
SWITCH_TYPES: tuple[SwitchEntityDescription, ...] = (
SwitchEntityDescription(key="on_off", translation_key="on_off"),
SwitchEntityDescription(key="turtle_mode", translation_key="turtle_mode"),
)
async def async_setup_entry(
hass: HomeAssistant,
@ -27,11 +31,9 @@ async def async_setup_entry(
config_entry.entry_id
]
entities = []
for switch_type, switch_name in SWITCH_TYPES.items():
entities.append(TransmissionSwitch(switch_type, switch_name, coordinator))
async_add_entities(entities)
async_add_entities(
TransmissionSwitch(coordinator, description) for description in SWITCH_TYPES
)
class TransmissionSwitch(
@ -40,20 +42,18 @@ class TransmissionSwitch(
"""Representation of a Transmission switch."""
_attr_has_entity_name = True
_attr_should_poll = False
def __init__(
self,
switch_type: str,
switch_name: str,
coordinator: TransmissionDataUpdateCoordinator,
entity_description: SwitchEntityDescription,
) -> None:
"""Initialize the Transmission switch."""
super().__init__(coordinator)
self._attr_name = switch_name
self.type = switch_type
self.unsub_update: Callable[[], None] | None = None
self._attr_unique_id = f"{coordinator.config_entry.entry_id}-{switch_type}"
self.entity_description = entity_description
self._attr_unique_id = (
f"{coordinator.config_entry.entry_id}-{entity_description.key}"
)
self._attr_device_info = DeviceInfo(
entry_type=DeviceEntryType.SERVICE,
identifiers={(DOMAIN, coordinator.config_entry.entry_id)},
@ -64,19 +64,19 @@ class TransmissionSwitch(
def is_on(self) -> bool:
"""Return true if device is on."""
active = None
if self.type == "on_off":
if self.entity_description.key == "on_off":
active = self.coordinator.data.active_torrent_count > 0
elif self.type == "turtle_mode":
elif self.entity_description.key == "turtle_mode":
active = self.coordinator.get_alt_speed_enabled()
return bool(active)
async def async_turn_on(self, **kwargs: Any) -> None:
"""Turn the device on."""
if self.type == "on_off":
if self.entity_description.key == "on_off":
_LOGGING.debug("Starting all torrents")
await self.hass.async_add_executor_job(self.coordinator.start_torrents)
elif self.type == "turtle_mode":
elif self.entity_description.key == "turtle_mode":
_LOGGING.debug("Turning Turtle Mode of Transmission on")
await self.hass.async_add_executor_job(
self.coordinator.set_alt_speed_enabled, True
@ -85,10 +85,10 @@ class TransmissionSwitch(
async def async_turn_off(self, **kwargs: Any) -> None:
"""Turn the device off."""
if self.type == "on_off":
if self.entity_description.key == "on_off":
_LOGGING.debug("Stopping all torrents")
await self.hass.async_add_executor_job(self.coordinator.stop_torrents)
if self.type == "turtle_mode":
if self.entity_description.key == "turtle_mode":
_LOGGING.debug("Turning Turtle Mode of Transmission off")
await self.hass.async_add_executor_job(
self.coordinator.set_alt_speed_enabled, False