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

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):
+    def check_removed_torrent(self):
-        """Get started torrent info functionality."""
+        """Get removed torrent functionality."""
-        all_torrents = self._api.get_torrents()
+        actual_torrents = self.torrents
-        current_down = {}
+        actual_all_torrents = [var.name for var in actual_torrents]
 
-        for torrent in all_torrents:
+        removed_torrents = list(set(self.all_torrents).difference(actual_all_torrents))
-            if torrent.status == "downloading":
+        for var in removed_torrents:
-                info = self.started_torrent_dict[torrent.name] = {
+            self.hass.bus.fire(EVENT_REMOVED_TORRENT, {"name": var})
-                    "added_date": torrent.addedDate,
+        self.all_torrents = actual_all_torrents
-                    "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)
 
     def start_torrents(self):
         """Start all torrents."""

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"

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 = []
+    dev = [
-    for sensor_type in SENSOR_TYPES:
+        TransmissionSpeedSensor(tm_client, name, "Down Speed", "download"),
-        dev.append(
+        TransmissionSpeedSensor(tm_client, name, "Up Speed", "upload"),
-            TransmissionSensor(
+        TransmissionStatusSensor(tm_client, name, "Status"),
-                sensor_type,
+        TransmissionTorrentsSensor(tm_client, name, "Active Torrents", "active"),
-                tm_client,
+        TransmissionTorrentsSensor(tm_client, name, "Paused Torrents", "paused"),
-                name,
+        TransmissionTorrentsSensor(tm_client, name, "Total Torrents", "total"),
-                SENSOR_TYPES[sensor_type][0],
+        TransmissionTorrentsSensor(tm_client, name, "Completed Torrents", "completed"),
-                SENSOR_TYPES[sensor_type][1],
+        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__(
+    def __init__(self, tm_client, client_name, sensor_name, sub_type=None):
-        self, sensor_type, tm_client, client_name, sensor_name, unit_of_measurement
-    ):
         """Initialize the sensor."""
-        self._name = sensor_name
-        self._state = None
         self._tm_client = tm_client
-        self._unit_of_measurement = unit_of_measurement
+        self._client_name = client_name
-        self._data = None
+        self._name = sensor_name
-        self.client_name = client_name
+        self._sub_type = sub_type
-        self.type = sensor_type
+        self._state = None
-        self.unsub_update = 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,
+        @callback
-            self._tm_client.api.signal_update,
+        def update():
-            self._schedule_immediate_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):
+class TransmissionSpeedSensor(TransmissionSensor):
-        """Unsubscribe from update dispatcher."""
+    """Representation of a Transmission speed sensor."""
-        if self.unsub_update:
+
-            self.unsub_update()
+    @property
-            self.unsub_update = None
+    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":
+class TransmissionStatusSensor(TransmissionSensor):
-            if self._data:
+    """Representation of a Transmission status sensor."""
-                upload = self._data.uploadSpeed
+
-                download = self._data.downloadSpeed
+    def update(self):
-                if upload > 0 and download > 0:
+        """Get the latest data from Transmission and updates the state."""
-                    self._state = "Up/Down"
+        data = self._tm_client.api.data
-                elif upload > 0 and download == 0:
+        if data:
-                    self._state = "Seeding"
+            upload = data.uploadSpeed
-                elif upload == 0 and download > 0:
+            download = data.downloadSpeed
-                    self._state = "Downloading"
+            if upload > 0 and download > 0:
-                else:
+                self._state = "Up/Down"
-                    self._state = STATE_IDLE
+            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":
+class TransmissionTorrentsSensor(TransmissionSensor):
-                mb_spd = float(self._data.downloadSpeed)
+    """Representation of a Transmission torrents sensor."""
-                mb_spd = mb_spd / 1024 / 1024
+
-                self._state = round(mb_spd, 2 if mb_spd < 0.1 else 1)
+    SUBTYPE_MODES = {
-            elif self.type == "upload_speed":
+        "started": ("downloading"),
-                mb_spd = float(self._data.uploadSpeed)
+        "completed": ("seeding"),
-                mb_spd = mb_spd / 1024 / 1024
+        "paused": ("stopped"),
-                self._state = round(mb_spd, 2 if mb_spd < 0.1 else 1)
+        "active": ("seeding", "downloading"),
-            elif self.type == "active_torrents":
+        "total": None,
-                self._state = self._data.activeTorrentCount
+    }
-            elif self.type == "paused_torrents":
+
-                self._state = self._data.pausedTorrentCount
+    @property
-            elif self.type == "total_torrents":
+    def unit_of_measurement(self):
-                self._state = self._data.torrentCount
+        """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

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