Add more type hints to Transmission (#99190)

* More type hints of transmssion

* More type hints
This commit is contained in:
Rami Mosleh 2023-08-28 12:20:18 +03:00 committed by GitHub
parent b0f3b7bb76
commit 8edae37082
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 54 additions and 33 deletions

View file

@ -14,6 +14,7 @@ from transmission_rpc.error import (
TransmissionConnectError,
TransmissionError,
)
from transmission_rpc.session import SessionStats
import voluptuous as vol
from homeassistant.config_entries import ConfigEntry, ConfigEntryState
@ -340,7 +341,7 @@ class TransmissionClient:
def set_scan_interval(self, scan_interval: float) -> None:
"""Update scan interval."""
def refresh(event_time: datetime):
def refresh(event_time: datetime) -> None:
"""Get the latest data from Transmission."""
self.api.update()
@ -367,22 +368,22 @@ class TransmissionData:
"""Initialize the Transmission RPC API."""
self.hass = hass
self.config = config
self.data: transmission_rpc.Session = None
self.available: bool = True
self._all_torrents: list[transmission_rpc.Torrent] = []
self._api: transmission_rpc.Client = api
self.data: SessionStats | None = None
self.available: bool = True
self._session: transmission_rpc.Session | None = None
self._all_torrents: list[transmission_rpc.Torrent] = []
self._completed_torrents: list[transmission_rpc.Torrent] = []
self._session: transmission_rpc.Session = None
self._started_torrents: list[transmission_rpc.Torrent] = []
self._torrents: list[transmission_rpc.Torrent] = []
@property
def host(self):
def host(self) -> str:
"""Return the host name."""
return self.config.data[CONF_HOST]
@property
def signal_update(self):
def signal_update(self) -> str:
"""Update signal per transmission entry."""
return f"{DATA_UPDATED}-{self.host}"
@ -391,7 +392,7 @@ class TransmissionData:
"""Get the list of torrents."""
return self._torrents
def update(self):
def update(self) -> None:
"""Get the latest data from Transmission instance."""
try:
self.data = self._api.session_stats()
@ -409,7 +410,7 @@ class TransmissionData:
_LOGGER.error("Unable to connect to Transmission client %s", self.host)
dispatcher_send(self.hass, self.signal_update)
def init_torrent_list(self):
def init_torrent_list(self) -> None:
"""Initialize torrent lists."""
self._torrents = self._api.get_torrents()
self._completed_torrents = [
@ -419,7 +420,7 @@ class TransmissionData:
torrent for torrent in self._torrents if torrent.status == "downloading"
]
def check_completed_torrent(self):
def check_completed_torrent(self) -> None:
"""Get completed torrent functionality."""
old_completed_torrent_names = {
torrent.name for torrent in self._completed_torrents
@ -437,7 +438,7 @@ class TransmissionData:
self._completed_torrents = current_completed_torrents
def check_started_torrent(self):
def check_started_torrent(self) -> None:
"""Get started torrent functionality."""
old_started_torrent_names = {torrent.name for torrent in self._started_torrents}
@ -453,7 +454,7 @@ class TransmissionData:
self._started_torrents = current_started_torrents
def check_removed_torrent(self):
def check_removed_torrent(self) -> None:
"""Get removed torrent functionality."""
current_torrent_names = {torrent.name for torrent in self._torrents}
@ -465,24 +466,24 @@ class TransmissionData:
self._all_torrents = self._torrents.copy()
def start_torrents(self):
def start_torrents(self) -> None:
"""Start all torrents."""
if not self._torrents:
return
self._api.start_all()
def stop_torrents(self):
def stop_torrents(self) -> None:
"""Stop all active torrents."""
if not self._torrents:
return
torrent_ids = [torrent.id for torrent in self._torrents]
self._api.stop_torrent(torrent_ids)
def set_alt_speed_enabled(self, is_enabled):
def set_alt_speed_enabled(self, is_enabled: bool) -> None:
"""Set the alternative speed flag."""
self._api.set_session(alt_speed_enabled=is_enabled)
def get_alt_speed_enabled(self):
def get_alt_speed_enabled(self) -> bool | None:
"""Get the alternative speed flag."""
if self._session is None:
return None

View file

@ -2,6 +2,7 @@
from __future__ import annotations
from contextlib import suppress
from typing import Any
from transmission_rpc.torrent import Torrent
@ -12,6 +13,7 @@ from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
from homeassistant.helpers.dispatcher import async_dispatcher_connect
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.typing import StateType
from . import TransmissionClient
from .const import (
@ -33,8 +35,8 @@ async def async_setup_entry(
) -> None:
"""Set up the Transmission sensors."""
tm_client = hass.data[DOMAIN][config_entry.entry_id]
name = config_entry.data[CONF_NAME]
tm_client: TransmissionClient = hass.data[DOMAIN][config_entry.entry_id]
name: str = config_entry.data[CONF_NAME]
dev = [
TransmissionSpeedSensor(
@ -96,12 +98,18 @@ class TransmissionSensor(SensorEntity):
_attr_has_entity_name = True
_attr_should_poll = False
def __init__(self, tm_client, client_name, sensor_translation_key, key):
def __init__(
self,
tm_client: TransmissionClient,
client_name: str,
sensor_translation_key: str,
key: str,
) -> None:
"""Initialize the sensor."""
self._tm_client: TransmissionClient = tm_client
self._tm_client = tm_client
self._attr_translation_key = sensor_translation_key
self._key = key
self._state = None
self._state: StateType = None
self._attr_unique_id = f"{tm_client.config_entry.entry_id}-{key}"
self._attr_device_info = DeviceInfo(
entry_type=DeviceEntryType.SERVICE,
@ -111,7 +119,7 @@ class TransmissionSensor(SensorEntity):
)
@property
def native_value(self):
def native_value(self) -> StateType:
"""Return the state of the sensor."""
return self._state
@ -192,12 +200,12 @@ class TransmissionTorrentsSensor(TransmissionSensor):
}
@property
def native_unit_of_measurement(self):
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):
def extra_state_attributes(self) -> dict[str, Any]:
"""Return the state attributes, if any."""
info = _torrents_info(
torrents=self._tm_client.api.torrents,
@ -217,7 +225,9 @@ class TransmissionTorrentsSensor(TransmissionSensor):
self._state = len(torrents)
def _filter_torrents(torrents: list[Torrent], statuses=None) -> list[Torrent]:
def _filter_torrents(
torrents: list[Torrent], statuses: list[str] | None = None
) -> list[Torrent]:
return [
torrent
for torrent in torrents
@ -225,7 +235,9 @@ def _filter_torrents(torrents: list[Torrent], statuses=None) -> list[Torrent]:
]
def _torrents_info(torrents, order, limit, statuses=None):
def _torrents_info(
torrents: list[Torrent], order: str, limit: int, statuses: list[str] | None = None
) -> dict[str, Any]:
infos = {}
torrents = _filter_torrents(torrents, statuses)
torrents = SUPPORTED_ORDER_MODES[order](torrents)

View file

@ -1,4 +1,5 @@
"""Support for setting the Transmission BitTorrent client Turtle Mode."""
from collections.abc import Callable
import logging
from typing import Any
@ -10,6 +11,7 @@ from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
from homeassistant.helpers.dispatcher import async_dispatcher_connect
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from . import TransmissionClient
from .const import DOMAIN, SWITCH_TYPES
_LOGGING = logging.getLogger(__name__)
@ -22,8 +24,8 @@ async def async_setup_entry(
) -> None:
"""Set up the Transmission switch."""
tm_client = hass.data[DOMAIN][config_entry.entry_id]
name = config_entry.data[CONF_NAME]
tm_client: TransmissionClient = hass.data[DOMAIN][config_entry.entry_id]
name: str = config_entry.data[CONF_NAME]
dev = []
for switch_type, switch_name in SWITCH_TYPES.items():
@ -38,14 +40,20 @@ class TransmissionSwitch(SwitchEntity):
_attr_has_entity_name = True
_attr_should_poll = False
def __init__(self, switch_type, switch_name, tm_client, client_name):
def __init__(
self,
switch_type: str,
switch_name: str,
tm_client: TransmissionClient,
client_name: str,
) -> None:
"""Initialize the Transmission switch."""
self._attr_name = switch_name
self.type = switch_type
self._tm_client = tm_client
self._state = STATE_OFF
self._data = None
self.unsub_update = None
self.unsub_update: Callable[[], None] | None = None
self._attr_unique_id = f"{tm_client.config_entry.entry_id}-{switch_type}"
self._attr_device_info = DeviceInfo(
entry_type=DeviceEntryType.SERVICE,
@ -55,7 +63,7 @@ class TransmissionSwitch(SwitchEntity):
)
@property
def is_on(self):
def is_on(self) -> bool:
"""Return true if device is on."""
return self._state == STATE_ON
@ -93,10 +101,10 @@ class TransmissionSwitch(SwitchEntity):
)
@callback
def _schedule_immediate_update(self):
def _schedule_immediate_update(self) -> None:
self.async_schedule_update_ha_state(True)
async def will_remove_from_hass(self):
async def will_remove_from_hass(self) -> None:
"""Unsubscribe from update dispatcher."""
if self.unsub_update:
self.unsub_update()