Skip to content

Commit

Permalink
Add entity description mixin to transmission switches (#103843)
Browse files Browse the repository at this point in the history
* Add entity description mixin for transmission switches

* minor fix
  • Loading branch information
engrbm87 committed Nov 12, 2023
1 parent 04a4973 commit 0c421b7
Showing 1 changed file with 43 additions and 27 deletions.
70 changes: 43 additions & 27 deletions homeassistant/components/transmission/switch.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
"""Support for setting the Transmission BitTorrent client Turtle Mode."""
from collections.abc import Callable
from dataclasses import dataclass
import logging
from typing import Any

Expand All @@ -14,9 +16,38 @@

_LOGGING = logging.getLogger(__name__)

SWITCH_TYPES: tuple[SwitchEntityDescription, ...] = (
SwitchEntityDescription(key="on_off", translation_key="on_off"),
SwitchEntityDescription(key="turtle_mode", translation_key="turtle_mode"),

@dataclass
class TransmissionSwitchEntityDescriptionMixin:
"""Mixin for required keys."""

is_on_func: Callable[[TransmissionDataUpdateCoordinator], bool | None]
on_func: Callable[[TransmissionDataUpdateCoordinator], None]
off_func: Callable[[TransmissionDataUpdateCoordinator], None]


@dataclass
class TransmissionSwitchEntityDescription(
SwitchEntityDescription, TransmissionSwitchEntityDescriptionMixin
):
"""Entity description class for Transmission switches."""


SWITCH_TYPES: tuple[TransmissionSwitchEntityDescription, ...] = (
TransmissionSwitchEntityDescription(
key="on_off",
translation_key="on_off",
is_on_func=lambda coordinator: coordinator.data.active_torrent_count > 0,
on_func=lambda coordinator: coordinator.start_torrents(),
off_func=lambda coordinator: coordinator.stop_torrents(),
),
TransmissionSwitchEntityDescription(
key="turtle_mode",
translation_key="turtle_mode",
is_on_func=lambda coordinator: coordinator.get_alt_speed_enabled(),
on_func=lambda coordinator: coordinator.set_alt_speed_enabled(True),
off_func=lambda coordinator: coordinator.set_alt_speed_enabled(False),
),
)


Expand All @@ -41,12 +72,13 @@ class TransmissionSwitch(
):
"""Representation of a Transmission switch."""

entity_description: TransmissionSwitchEntityDescription
_attr_has_entity_name = True

def __init__(
self,
coordinator: TransmissionDataUpdateCoordinator,
entity_description: SwitchEntityDescription,
entity_description: TransmissionSwitchEntityDescription,
) -> None:
"""Initialize the Transmission switch."""
super().__init__(coordinator)
Expand All @@ -63,34 +95,18 @@ def __init__(
@property
def is_on(self) -> bool:
"""Return true if device is on."""
active = None
if self.entity_description.key == "on_off":
active = self.coordinator.data.active_torrent_count > 0
elif self.entity_description.key == "turtle_mode":
active = self.coordinator.get_alt_speed_enabled()

return bool(active)
return bool(self.entity_description.is_on_func(self.coordinator))

async def async_turn_on(self, **kwargs: Any) -> None:
"""Turn the device on."""
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.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
)
await self.hass.async_add_executor_job(
self.entity_description.on_func, self.coordinator
)
await self.coordinator.async_request_refresh()

async def async_turn_off(self, **kwargs: Any) -> None:
"""Turn the device 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.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
)
await self.hass.async_add_executor_job(
self.entity_description.off_func, self.coordinator
)
await self.coordinator.async_request_refresh()

0 comments on commit 0c421b7

Please sign in to comment.