From 9062d6e5e6fbb2a25f743ed3af7a76d10cb5b15c Mon Sep 17 00:00:00 2001 From: Gleb Sinyavskiy Date: Mon, 20 Apr 2020 15:07:26 +0200 Subject: [PATCH] Improve the transmission integration (#34223) * Update state after adding a new torrent * Use cached torrents list in check_started_torrent_info * Add torrent_info to all sensors * Add torrent_info for active torrents * Fix typo * Update codeowners * Do not set eta if it's unknown * Fix codeowners * Extract TransmissionSpeedSensor * Extract TransmissionStatusSensor * Extract TransmissionTorrentsSensor * Refactor device_state_attributes() and update() * Remove unused methods * Use async_on_remove * Fix sensor update * Add transmission.remove_torrent service * Add transmission_removed_torrent event * Fix naming * Fix typo in services.yaml --- .../components/transmission/__init__.py | 88 ++++---- .../components/transmission/const.py | 21 +- .../components/transmission/sensor.py | 190 ++++++++++-------- .../components/transmission/services.yaml | 13 ++ 4 files changed, 182 insertions(+), 130 deletions(-) diff --git a/homeassistant/components/transmission/__init__.py b/homeassistant/components/transmission/__init__.py index 08aa52e3a13..32177e91160 100644 --- a/homeassistant/components/transmission/__init__.py +++ b/homeassistant/components/transmission/__init__.py @@ -9,6 +9,7 @@ import voluptuous as vol from homeassistant.config_entries import SOURCE_IMPORT from homeassistant.const import ( CONF_HOST, + CONF_ID, CONF_NAME, CONF_PASSWORD, CONF_PORT, @@ -21,13 +22,19 @@ from homeassistant.helpers.dispatcher import dispatcher_send from homeassistant.helpers.event import async_track_time_interval from .const import ( + ATTR_DELETE_DATA, ATTR_TORRENT, DATA_UPDATED, + DEFAULT_DELETE_DATA, DEFAULT_NAME, DEFAULT_PORT, DEFAULT_SCAN_INTERVAL, DOMAIN, + EVENT_DOWNLOADED_TORRENT, + EVENT_REMOVED_TORRENT, + EVENT_STARTED_TORRENT, SERVICE_ADD_TORRENT, + SERVICE_REMOVE_TORRENT, ) from .errors import AuthenticationError, CannotConnect, UnknownError @@ -38,6 +45,14 @@ SERVICE_ADD_TORRENT_SCHEMA = vol.Schema( {vol.Required(ATTR_TORRENT): cv.string, vol.Required(CONF_NAME): cv.string} ) +SERVICE_REMOVE_TORRENT_SCHEMA = vol.Schema( + { + vol.Required(CONF_NAME): cv.string, + vol.Required(CONF_ID): cv.positive_int, + vol.Optional(ATTR_DELETE_DATA, default=DEFAULT_DELETE_DATA): cv.boolean, + } +) + TRANS_SCHEMA = vol.All( vol.Schema( { @@ -95,6 +110,7 @@ async def async_unload_entry(hass, config_entry): if not hass.data[DOMAIN]: hass.services.async_remove(DOMAIN, SERVICE_ADD_TORRENT) + hass.services.async_remove(DOMAIN, SERVICE_REMOVE_TORRENT) return True @@ -180,15 +196,38 @@ class TransmissionClient: ("http", "ftp:", "magnet:") ) or self.hass.config.is_allowed_path(torrent): tm_client.tm_api.add_torrent(torrent) + tm_client.api.update() else: _LOGGER.warning( "Could not add torrent: unsupported type or no permission" ) + def remove_torrent(service): + """Remove torrent.""" + tm_client = None + for entry in self.hass.config_entries.async_entries(DOMAIN): + if entry.data[CONF_NAME] == service.data[CONF_NAME]: + tm_client = self.hass.data[DOMAIN][entry.entry_id] + break + if tm_client is None: + _LOGGER.error("Transmission instance is not found") + return + torrent_id = service.data[CONF_ID] + delete_data = service.data[ATTR_DELETE_DATA] + tm_client.tm_api.remove_torrent(torrent_id, delete_data=delete_data) + tm_client.api.update() + self.hass.services.async_register( DOMAIN, SERVICE_ADD_TORRENT, add_torrent, schema=SERVICE_ADD_TORRENT_SCHEMA ) + self.hass.services.async_register( + DOMAIN, + SERVICE_REMOVE_TORRENT, + remove_torrent, + schema=SERVICE_REMOVE_TORRENT_SCHEMA, + ) + self.config_entry.add_update_listener(self.async_options_updated) return True @@ -234,13 +273,13 @@ class TransmissionData: self.hass = hass self.config = config self.data = None - self.torrents = None + self.torrents = [] self.session = None self.available = True self._api = api self.completed_torrents = [] self.started_torrents = [] - self.started_torrent_dict = {} + self.all_torrents = [] @property def host(self): @@ -259,9 +298,9 @@ class TransmissionData: self.torrents = self._api.get_torrents() self.session = self._api.get_session() - self.check_started_torrent_info() self.check_completed_torrent() self.check_started_torrent() + self.check_removed_torrent() _LOGGER.debug("Torrent Data for %s Updated", self.host) self.available = True @@ -292,7 +331,7 @@ class TransmissionData: ) for var in tmp_completed_torrents: - self.hass.bus.fire("transmission_downloaded_torrent", {"name": var}) + self.hass.bus.fire(EVENT_DOWNLOADED_TORRENT, {"name": var}) self.completed_torrents = actual_completed_torrents @@ -308,41 +347,18 @@ class TransmissionData: ) for var in tmp_started_torrents: - self.hass.bus.fire("transmission_started_torrent", {"name": var}) + self.hass.bus.fire(EVENT_STARTED_TORRENT, {"name": var}) self.started_torrents = actual_started_torrents - def check_started_torrent_info(self): - """Get started torrent info functionality.""" - all_torrents = self._api.get_torrents() - current_down = {} + def check_removed_torrent(self): + """Get removed torrent functionality.""" + actual_torrents = self.torrents + actual_all_torrents = [var.name for var in actual_torrents] - for torrent in all_torrents: - if torrent.status == "downloading": - info = self.started_torrent_dict[torrent.name] = { - "added_date": torrent.addedDate, - "percent_done": f"{torrent.percentDone * 100:.2f}", - } - try: - info["eta"] = str(torrent.eta) - except ValueError: - info["eta"] = "unknown" - - current_down[torrent.name] = True - - elif torrent.name in self.started_torrent_dict: - self.started_torrent_dict.pop(torrent.name) - - for torrent in list(self.started_torrent_dict): - if torrent not in current_down: - self.started_torrent_dict.pop(torrent) - - def get_started_torrent_count(self): - """Get the number of started torrents.""" - return len(self.started_torrents) - - def get_completed_torrent_count(self): - """Get the number of completed torrents.""" - return len(self.completed_torrents) + removed_torrents = list(set(self.all_torrents).difference(actual_all_torrents)) + for var in removed_torrents: + self.hass.bus.fire(EVENT_REMOVED_TORRENT, {"name": var}) + self.all_torrents = actual_all_torrents def start_torrents(self): """Start all torrents.""" diff --git a/homeassistant/components/transmission/const.py b/homeassistant/components/transmission/const.py index 659ef97d9de..8edbf944890 100644 --- a/homeassistant/components/transmission/const.py +++ b/homeassistant/components/transmission/const.py @@ -1,27 +1,24 @@ """Constants for the Transmission Bittorent Client component.""" -from homeassistant.const import DATA_RATE_MEGABYTES_PER_SECOND - DOMAIN = "transmission" -SENSOR_TYPES = { - "active_torrents": ["Active Torrents", "Torrents"], - "current_status": ["Status", None], - "download_speed": ["Down Speed", DATA_RATE_MEGABYTES_PER_SECOND], - "paused_torrents": ["Paused Torrents", "Torrents"], - "total_torrents": ["Total Torrents", "Torrents"], - "upload_speed": ["Up Speed", DATA_RATE_MEGABYTES_PER_SECOND], - "completed_torrents": ["Completed Torrents", "Torrents"], - "started_torrents": ["Started Torrents", "Torrents"], -} SWITCH_TYPES = {"on_off": "Switch", "turtle_mode": "Turtle Mode"} +DEFAULT_DELETE_DATA = False DEFAULT_NAME = "Transmission" DEFAULT_PORT = 9091 DEFAULT_SCAN_INTERVAL = 120 STATE_ATTR_TORRENT_INFO = "torrent_info" + +ATTR_DELETE_DATA = "delete_data" ATTR_TORRENT = "torrent" + SERVICE_ADD_TORRENT = "add_torrent" +SERVICE_REMOVE_TORRENT = "remove_torrent" DATA_UPDATED = "transmission_data_updated" + +EVENT_STARTED_TORRENT = "transmission_started_torrent" +EVENT_REMOVED_TORRENT = "transmission_removed_torrent" +EVENT_DOWNLOADED_TORRENT = "transmission_downloaded_torrent" diff --git a/homeassistant/components/transmission/sensor.py b/homeassistant/components/transmission/sensor.py index 0db731d6f01..812d63f24d8 100644 --- a/homeassistant/components/transmission/sensor.py +++ b/homeassistant/components/transmission/sensor.py @@ -1,12 +1,12 @@ """Support for monitoring the Transmission BitTorrent client API.""" import logging -from homeassistant.const import CONF_NAME, STATE_IDLE +from homeassistant.const import CONF_NAME, DATA_RATE_MEGABYTES_PER_SECOND, STATE_IDLE from homeassistant.core import callback from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity import Entity -from .const import DOMAIN, SENSOR_TYPES, STATE_ATTR_TORRENT_INFO +from .const import DOMAIN, STATE_ATTR_TORRENT_INFO _LOGGER = logging.getLogger(__name__) @@ -17,41 +17,35 @@ async def async_setup_entry(hass, config_entry, async_add_entities): tm_client = hass.data[DOMAIN][config_entry.entry_id] name = config_entry.data[CONF_NAME] - dev = [] - for sensor_type in SENSOR_TYPES: - dev.append( - TransmissionSensor( - sensor_type, - tm_client, - name, - SENSOR_TYPES[sensor_type][0], - SENSOR_TYPES[sensor_type][1], - ) - ) + dev = [ + TransmissionSpeedSensor(tm_client, name, "Down Speed", "download"), + TransmissionSpeedSensor(tm_client, name, "Up Speed", "upload"), + TransmissionStatusSensor(tm_client, name, "Status"), + TransmissionTorrentsSensor(tm_client, name, "Active Torrents", "active"), + TransmissionTorrentsSensor(tm_client, name, "Paused Torrents", "paused"), + TransmissionTorrentsSensor(tm_client, name, "Total Torrents", "total"), + TransmissionTorrentsSensor(tm_client, name, "Completed Torrents", "completed"), + TransmissionTorrentsSensor(tm_client, name, "Started Torrents", "started"), + ] async_add_entities(dev, True) class TransmissionSensor(Entity): - """Representation of a Transmission sensor.""" + """A base class for all Transmission sensors.""" - def __init__( - self, sensor_type, tm_client, client_name, sensor_name, unit_of_measurement - ): + def __init__(self, tm_client, client_name, sensor_name, sub_type=None): """Initialize the sensor.""" - self._name = sensor_name - self._state = None self._tm_client = tm_client - self._unit_of_measurement = unit_of_measurement - self._data = None - self.client_name = client_name - self.type = sensor_type - self.unsub_update = None + self._client_name = client_name + self._name = sensor_name + self._sub_type = sub_type + self._state = None @property def name(self): """Return the name of the sensor.""" - return f"{self.client_name} {self._name}" + return f"{self._client_name} {self._name}" @property def unique_id(self): @@ -68,77 +62,109 @@ class TransmissionSensor(Entity): """Return the polling requirement for this sensor.""" return False - @property - def unit_of_measurement(self): - """Return the unit of measurement of this entity, if any.""" - return self._unit_of_measurement - @property def available(self): """Could the device be accessed during the last update call.""" return self._tm_client.api.available - @property - def device_state_attributes(self): - """Return the state attributes, if any.""" - if self._tm_client.api.started_torrent_dict and self.type == "started_torrents": - return {STATE_ATTR_TORRENT_INFO: self._tm_client.api.started_torrent_dict} - return None - async def async_added_to_hass(self): """Handle entity which will be added.""" - self.unsub_update = async_dispatcher_connect( - self.hass, - self._tm_client.api.signal_update, - self._schedule_immediate_update, + + @callback + def update(): + """Update the state.""" + self.async_schedule_update_ha_state(True) + + self.async_on_remove( + async_dispatcher_connect( + self.hass, self._tm_client.api.signal_update, update + ) ) - @callback - def _schedule_immediate_update(self): - self.async_schedule_update_ha_state(True) - async def will_remove_from_hass(self): - """Unsubscribe from update dispatcher.""" - if self.unsub_update: - self.unsub_update() - self.unsub_update = None +class TransmissionSpeedSensor(TransmissionSensor): + """Representation of a Transmission speed sensor.""" + + @property + def unit_of_measurement(self): + """Return the unit of measurement of this entity, if any.""" + return DATA_RATE_MEGABYTES_PER_SECOND def update(self): """Get the latest data from Transmission and updates the state.""" - self._data = self._tm_client.api.data + data = self._tm_client.api.data + if data: + mb_spd = ( + float(data.downloadSpeed) + if self._sub_type == "download" + else float(data.uploadSpeed) + ) + mb_spd = mb_spd / 1024 / 1024 + self._state = round(mb_spd, 2 if mb_spd < 0.1 else 1) - if self.type == "completed_torrents": - self._state = self._tm_client.api.get_completed_torrent_count() - elif self.type == "started_torrents": - self._state = self._tm_client.api.get_started_torrent_count() - if self.type == "current_status": - if self._data: - upload = self._data.uploadSpeed - download = self._data.downloadSpeed - if upload > 0 and download > 0: - self._state = "Up/Down" - elif upload > 0 and download == 0: - self._state = "Seeding" - elif upload == 0 and download > 0: - self._state = "Downloading" - else: - self._state = STATE_IDLE +class TransmissionStatusSensor(TransmissionSensor): + """Representation of a Transmission status sensor.""" + + def update(self): + """Get the latest data from Transmission and updates the state.""" + data = self._tm_client.api.data + if data: + upload = data.uploadSpeed + download = data.downloadSpeed + if upload > 0 and download > 0: + self._state = "Up/Down" + elif upload > 0 and download == 0: + self._state = "Seeding" + elif upload == 0 and download > 0: + self._state = "Downloading" else: - self._state = None + self._state = STATE_IDLE + else: + self._state = None - if self._data: - if self.type == "download_speed": - mb_spd = float(self._data.downloadSpeed) - mb_spd = mb_spd / 1024 / 1024 - self._state = round(mb_spd, 2 if mb_spd < 0.1 else 1) - elif self.type == "upload_speed": - mb_spd = float(self._data.uploadSpeed) - mb_spd = mb_spd / 1024 / 1024 - self._state = round(mb_spd, 2 if mb_spd < 0.1 else 1) - elif self.type == "active_torrents": - self._state = self._data.activeTorrentCount - elif self.type == "paused_torrents": - self._state = self._data.pausedTorrentCount - elif self.type == "total_torrents": - self._state = self._data.torrentCount + +class TransmissionTorrentsSensor(TransmissionSensor): + """Representation of a Transmission torrents sensor.""" + + SUBTYPE_MODES = { + "started": ("downloading"), + "completed": ("seeding"), + "paused": ("stopped"), + "active": ("seeding", "downloading"), + "total": None, + } + + @property + def unit_of_measurement(self): + """Return the unit of measurement of this entity, if any.""" + return "Torrents" + + @property + def device_state_attributes(self): + """Return the state attributes, if any.""" + info = _torrents_info( + self._tm_client.api.torrents, self.SUBTYPE_MODES[self._sub_type] + ) + return {STATE_ATTR_TORRENT_INFO: info} + + def update(self): + """Get the latest data from Transmission and updates the state.""" + self._state = len(self.device_state_attributes[STATE_ATTR_TORRENT_INFO]) + + +def _torrents_info(torrents, statuses=None): + infos = {} + for torrent in torrents: + if statuses is None or torrent.status in statuses: + info = infos[torrent.name] = { + "added_date": torrent.addedDate, + "percent_done": f"{torrent.percentDone * 100:.2f}", + "status": torrent.status, + "id": torrent.id, + } + try: + info["eta"] = str(torrent.eta) + except ValueError: + pass + return infos diff --git a/homeassistant/components/transmission/services.yaml b/homeassistant/components/transmission/services.yaml index de3314e20f6..e8114b680ab 100644 --- a/homeassistant/components/transmission/services.yaml +++ b/homeassistant/components/transmission/services.yaml @@ -7,3 +7,16 @@ add_torrent: torrent: description: URL, magnet link or Base64 encoded file. example: http://releases.ubuntu.com/19.04/ubuntu-19.04-desktop-amd64.iso.torrent + +remove_torrent: + description: Remove a torrent + fields: + name: + description: Instance name as entered during entry config + example: Transmission + id: + description: ID of a torrent + example: 123 + delete_data: + description: Delete torrent data + example: false