Use EntityDescription - qbittorrent (#54428)

This commit is contained in:
Marc Mueller 2021-08-16 22:52:47 +02:00 committed by GitHub
parent bb4a36c877
commit f9fbcd4aec
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -1,11 +1,17 @@
"""Support for monitoring the qBittorrent API.""" """Support for monitoring the qBittorrent API."""
from __future__ import annotations
import logging import logging
from qbittorrent.client import Client, LoginRequired from qbittorrent.client import Client, LoginRequired
from requests.exceptions import RequestException from requests.exceptions import RequestException
import voluptuous as vol import voluptuous as vol
from homeassistant.components.sensor import PLATFORM_SCHEMA, SensorEntity from homeassistant.components.sensor import (
PLATFORM_SCHEMA,
SensorEntity,
SensorEntityDescription,
)
from homeassistant.const import ( from homeassistant.const import (
CONF_NAME, CONF_NAME,
CONF_PASSWORD, CONF_PASSWORD,
@ -25,11 +31,22 @@ SENSOR_TYPE_UPLOAD_SPEED = "upload_speed"
DEFAULT_NAME = "qBittorrent" DEFAULT_NAME = "qBittorrent"
SENSOR_TYPES = { SENSOR_TYPES: tuple[SensorEntityDescription, ...] = (
SENSOR_TYPE_CURRENT_STATUS: ["Status", None], SensorEntityDescription(
SENSOR_TYPE_DOWNLOAD_SPEED: ["Down Speed", DATA_RATE_KILOBYTES_PER_SECOND], key=SENSOR_TYPE_CURRENT_STATUS,
SENSOR_TYPE_UPLOAD_SPEED: ["Up Speed", DATA_RATE_KILOBYTES_PER_SECOND], name="Status",
} ),
SensorEntityDescription(
key=SENSOR_TYPE_DOWNLOAD_SPEED,
name="Down Speed",
native_unit_of_measurement=DATA_RATE_KILOBYTES_PER_SECOND,
),
SensorEntityDescription(
key=SENSOR_TYPE_UPLOAD_SPEED,
name="Up Speed",
native_unit_of_measurement=DATA_RATE_KILOBYTES_PER_SECOND,
),
)
PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend( PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
{ {
@ -56,12 +73,12 @@ def setup_platform(hass, config, add_entities, discovery_info=None):
name = config.get(CONF_NAME) name = config.get(CONF_NAME)
dev = [] entities = [
for sensor_type in SENSOR_TYPES: QBittorrentSensor(description, client, name, LoginRequired)
sensor = QBittorrentSensor(sensor_type, client, name, LoginRequired) for description in SENSOR_TYPES
dev.append(sensor) ]
add_entities(dev, True) add_entities(entities, True)
def format_speed(speed): def format_speed(speed):
@ -73,45 +90,29 @@ def format_speed(speed):
class QBittorrentSensor(SensorEntity): class QBittorrentSensor(SensorEntity):
"""Representation of an qBittorrent sensor.""" """Representation of an qBittorrent sensor."""
def __init__(self, sensor_type, qbittorrent_client, client_name, exception): def __init__(
self,
description: SensorEntityDescription,
qbittorrent_client,
client_name,
exception,
):
"""Initialize the qBittorrent sensor.""" """Initialize the qBittorrent sensor."""
self._name = SENSOR_TYPES[sensor_type][0] self.entity_description = description
self.client = qbittorrent_client self.client = qbittorrent_client
self.type = sensor_type
self.client_name = client_name
self._state = None
self._unit_of_measurement = SENSOR_TYPES[sensor_type][1]
self._available = False
self._exception = exception self._exception = exception
@property self._attr_name = f"{client_name} {description.name}"
def name(self): self._attr_available = False
"""Return the name of the sensor."""
return f"{self.client_name} {self._name}"
@property
def native_value(self):
"""Return the state of the sensor."""
return self._state
@property
def available(self):
"""Return true if device is available."""
return self._available
@property
def native_unit_of_measurement(self):
"""Return the unit of measurement of this entity, if any."""
return self._unit_of_measurement
def update(self): def update(self):
"""Get the latest data from qBittorrent and updates the state.""" """Get the latest data from qBittorrent and updates the state."""
try: try:
data = self.client.sync_main_data() data = self.client.sync_main_data()
self._available = True self._attr_available = True
except RequestException: except RequestException:
_LOGGER.error("Connection lost") _LOGGER.error("Connection lost")
self._available = False self._attr_available = False
return return
except self._exception: except self._exception:
_LOGGER.error("Invalid authentication") _LOGGER.error("Invalid authentication")
@ -123,17 +124,18 @@ class QBittorrentSensor(SensorEntity):
download = data["server_state"]["dl_info_speed"] download = data["server_state"]["dl_info_speed"]
upload = data["server_state"]["up_info_speed"] upload = data["server_state"]["up_info_speed"]
if self.type == SENSOR_TYPE_CURRENT_STATUS: sensor_type = self.entity_description.key
if sensor_type == SENSOR_TYPE_CURRENT_STATUS:
if upload > 0 and download > 0: if upload > 0 and download > 0:
self._state = "up_down" self._attr_native_value = "up_down"
elif upload > 0 and download == 0: elif upload > 0 and download == 0:
self._state = "seeding" self._attr_native_value = "seeding"
elif upload == 0 and download > 0: elif upload == 0 and download > 0:
self._state = "downloading" self._attr_native_value = "downloading"
else: else:
self._state = STATE_IDLE self._attr_native_value = STATE_IDLE
elif self.type == SENSOR_TYPE_DOWNLOAD_SPEED: elif sensor_type == SENSOR_TYPE_DOWNLOAD_SPEED:
self._state = format_speed(download) self._attr_native_value = format_speed(download)
elif self.type == SENSOR_TYPE_UPLOAD_SPEED: elif sensor_type == SENSOR_TYPE_UPLOAD_SPEED:
self._state = format_speed(upload) self._attr_native_value = format_speed(upload)