sensor.py

Go to the documentation of this file.

"""Support for monitoring the rtorrent BitTorrent client API."""
 
from __future__ import annotations
 
import logging
import xmlrpc.client
 
import voluptuous as vol
 
from homeassistant.components.sensor import (
    PLATFORM_SCHEMA as SENSOR_PLATFORM_SCHEMA,
    SensorDeviceClass,
    SensorEntity,
    SensorEntityDescription,
)
from homeassistant.const import (
    CONF_MONITORED_VARIABLES,
    CONF_NAME,
    CONF_URL,
    STATE_IDLE,
    UnitOfDataRate,
)
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import PlatformNotReady
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
 
_LOGGER = logging.getLogger(__name__)
 
SENSOR_TYPE_CURRENT_STATUS = "current_status"
SENSOR_TYPE_DOWNLOAD_SPEED = "download_speed"
SENSOR_TYPE_UPLOAD_SPEED = "upload_speed"
SENSOR_TYPE_ALL_TORRENTS = "all_torrents"
SENSOR_TYPE_STOPPED_TORRENTS = "stopped_torrents"
SENSOR_TYPE_COMPLETE_TORRENTS = "complete_torrents"
SENSOR_TYPE_UPLOADING_TORRENTS = "uploading_torrents"
SENSOR_TYPE_DOWNLOADING_TORRENTS = "downloading_torrents"
SENSOR_TYPE_ACTIVE_TORRENTS = "active_torrents"
 
DEFAULT_NAME = "rtorrent"
SENSOR_TYPES: tuple[SensorEntityDescription, ...] = (
    SensorEntityDescription(
        key=SENSOR_TYPE_CURRENT_STATUS,
        name="Status",
    ),
    SensorEntityDescription(
        key=SENSOR_TYPE_DOWNLOAD_SPEED,
        name="Down Speed",
        device_class=SensorDeviceClass.DATA_RATE,
        native_unit_of_measurement=UnitOfDataRate.KILOBYTES_PER_SECOND,
    ),
    SensorEntityDescription(
        key=SENSOR_TYPE_UPLOAD_SPEED,
        name="Up Speed",
        device_class=SensorDeviceClass.DATA_RATE,
        native_unit_of_measurement=UnitOfDataRate.KILOBYTES_PER_SECOND,
    ),
    SensorEntityDescription(
        key=SENSOR_TYPE_ALL_TORRENTS,
        name="All Torrents",
    ),
    SensorEntityDescription(
        key=SENSOR_TYPE_STOPPED_TORRENTS,
        name="Stopped Torrents",
    ),
    SensorEntityDescription(
        key=SENSOR_TYPE_COMPLETE_TORRENTS,
        name="Complete Torrents",
    ),
    SensorEntityDescription(
        key=SENSOR_TYPE_UPLOADING_TORRENTS,
        name="Uploading Torrents",
    ),
    SensorEntityDescription(
        key=SENSOR_TYPE_DOWNLOADING_TORRENTS,
        name="Downloading Torrents",
    ),
    SensorEntityDescription(
        key=SENSOR_TYPE_ACTIVE_TORRENTS,
        name="Active Torrents",
    ),
)
 
SENSOR_KEYS: list[str] = [desc.key for desc in SENSOR_TYPES]
 
PLATFORM_SCHEMA = SENSOR_PLATFORM_SCHEMA.extend(
    {
        vol.Required(CONF_URL): cv.url,
        vol.Optional(CONF_NAME, default=DEFAULT_NAME): cv.string,
        vol.Optional(CONF_MONITORED_VARIABLES, default=SENSOR_KEYS): vol.All(
            cv.ensure_list, [vol.In(SENSOR_KEYS)]
        ),
    }
)
 
 
def setup_platform(
    hass: HomeAssistant,
    config: ConfigType,
    add_entities: AddEntitiesCallback,
    discovery_info: DiscoveryInfoType | None = None,
) -> None:
    """Set up the rtorrent sensors."""
    url = config[CONF_URL]
    name = config[CONF_NAME]
 
    try:
        rtorrent = xmlrpc.client.ServerProxy(url)
    except (xmlrpc.client.ProtocolError, ConnectionRefusedError) as ex:
        _LOGGER.error("Connection to rtorrent daemon failed")
        raise PlatformNotReady from ex
    monitored_variables = config[CONF_MONITORED_VARIABLES]
    entities = [
        RTorrentSensor(rtorrent, name, description)
        for description in SENSOR_TYPES
        if description.key in monitored_variables
    ]
 
    add_entities(entities)
 
 
def format_speed(speed):
    """Return a bytes/s measurement as a human readable string."""
    kb_spd = float(speed) / 1024
    return round(kb_spd, 2 if kb_spd < 0.1 else 1)
 
 
class RTorrentSensor(SensorEntity):
    """Representation of an rtorrent sensor."""
 
    def __init__(
        self, rtorrent_client, client_name, description: SensorEntityDescription
    ) -> None:
        """Initialize the sensor."""
        self.entity_descriptionentity_description = description
        self.clientclient = rtorrent_client
        self.datadata = None
 
        self._attr_name_attr_name = f"{client_name} {description.name}"
        self._attr_available_attr_available = False
 
    def update(self):
        """Get the latest data from rtorrent and updates the state."""
        multicall = xmlrpc.client.MultiCall(self.clientclient)
        multicall.throttle.global_up.rate()
        multicall.throttle.global_down.rate()
        multicall.d.multicall2("", "main")
        multicall.d.multicall2("", "stopped")
        multicall.d.multicall2("", "complete")
        multicall.d.multicall2("", "seeding", "d.up.rate=")
        multicall.d.multicall2("", "leeching", "d.down.rate=")
 
        try:
            self.datadata = multicall()
            self._attr_available_attr_available = True
        except (xmlrpc.client.ProtocolError, OSError) as ex:
            _LOGGER.error("Connection to rtorrent failed (%s)", ex)
            self._attr_available_attr_available = False
            return
 
        upload = self.datadata[0]
        download = self.datadata[1]
        all_torrents = self.datadata[2]
        stopped_torrents = self.datadata[3]
        complete_torrents = self.datadata[4]
 
        uploading_torrents = 0
        for up_torrent in self.datadata[5]:
            if up_torrent[0]:
                uploading_torrents += 1
 
        downloading_torrents = 0
        for down_torrent in self.datadata[6]:
            if down_torrent[0]:
                downloading_torrents += 1
 
        active_torrents = uploading_torrents + downloading_torrents
 
        sensor_type = self.entity_descriptionentity_description.key
        if sensor_type == SENSOR_TYPE_CURRENT_STATUS:
            if self.datadata:
                if upload > 0 and download > 0:
                    self._attr_native_value_attr_native_value = "up_down"
                elif upload > 0 and download == 0:
                    self._attr_native_value_attr_native_value = "seeding"
                elif upload == 0 and download > 0:
                    self._attr_native_value_attr_native_value = "downloading"
                else:
                    self._attr_native_value_attr_native_value = STATE_IDLE
            else:
                self._attr_native_value_attr_native_value = None
 
        if self.datadata:
            if sensor_type == SENSOR_TYPE_DOWNLOAD_SPEED:
                self._attr_native_value_attr_native_value = format_speed(download)
            elif sensor_type == SENSOR_TYPE_UPLOAD_SPEED:
                self._attr_native_value_attr_native_value = format_speed(upload)
            elif sensor_type == SENSOR_TYPE_ALL_TORRENTS:
                self._attr_native_value_attr_native_value = len(all_torrents)
            elif sensor_type == SENSOR_TYPE_STOPPED_TORRENTS:
                self._attr_native_value_attr_native_value = len(stopped_torrents)
            elif sensor_type == SENSOR_TYPE_COMPLETE_TORRENTS:
                self._attr_native_value_attr_native_value = len(complete_torrents)
            elif sensor_type == SENSOR_TYPE_UPLOADING_TORRENTS:
                self._attr_native_value_attr_native_value = uploading_torrents
            elif sensor_type == SENSOR_TYPE_DOWNLOADING_TORRENTS:
                self._attr_native_value_attr_native_value = downloading_torrents
            elif sensor_type == SENSOR_TYPE_ACTIVE_TORRENTS:
                self._attr_native_value_attr_native_value = active_torrents