"""Support for Radarr.""" from __future__ import annotations from collections.abc import Callable import dataclasses from datetime import UTC, datetime from typing import Any, Generic from aiopyarr import Diskspace, RootFolder, SystemStatus from homeassistant.components.sensor import ( SensorDeviceClass, SensorEntity, SensorEntityDescription, SensorStateClass, ) from homeassistant.const import EntityCategory, UnitOfInformation from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback from . import RadarrConfigEntry, RadarrEntity from .coordinator import RadarrDataUpdateCoordinator, T def get_space(data: list[Diskspace], name: str) -> str: """Get space.""" space = [ mount.freeSpace / 1024 ** BYTE_SIZES.index(UnitOfInformation.GIGABYTES) for mount in data if name in mount.path ] return f"{space[0]:.2f}" def get_modified_description( description: RadarrSensorEntityDescription[T], mount: RootFolder ) -> tuple[RadarrSensorEntityDescription[T], str]: """Return modified description and folder name.""" name = mount.path.rsplit("/")[-1].rsplit("\\")[-1] desc = dataclasses.replace( description, key=f"{description.key}_{name}", name=f"{description.name} {name}".capitalize(), ) return desc, name @dataclasses.dataclass(frozen=True) class RadarrSensorEntityDescriptionMixIn(Generic[T]): """Mixin for required keys.""" value_fn: Callable[[T, str], str | int | datetime] @dataclasses.dataclass(frozen=True) class RadarrSensorEntityDescription( SensorEntityDescription, RadarrSensorEntityDescriptionMixIn[T], Generic[T] ): """Class to describe a Radarr sensor.""" description_fn: ( Callable[ [RadarrSensorEntityDescription[T], RootFolder], tuple[RadarrSensorEntityDescription[T], str] | None, ] | None ) = None SENSOR_TYPES: dict[str, RadarrSensorEntityDescription[Any]] = { "disk_space": RadarrSensorEntityDescription( key="disk_space", name="Disk space", native_unit_of_measurement=UnitOfInformation.GIGABYTES, device_class=SensorDeviceClass.DATA_SIZE, icon="mdi:harddisk", value_fn=get_space, description_fn=get_modified_description, ), "movie": RadarrSensorEntityDescription[int]( key="movies", translation_key="movies", native_unit_of_measurement="Movies", entity_registry_enabled_default=False, value_fn=lambda data, _: data, ), "queue": RadarrSensorEntityDescription[int]( key="queue", translation_key="queue", native_unit_of_measurement="Movies", entity_registry_enabled_default=False, state_class=SensorStateClass.TOTAL, value_fn=lambda data, _: data, ), "status": RadarrSensorEntityDescription[SystemStatus]( key="start_time", translation_key="start_time", device_class=SensorDeviceClass.TIMESTAMP, entity_category=EntityCategory.DIAGNOSTIC, entity_registry_enabled_default=False, value_fn=lambda data, _: data.startTime.replace(tzinfo=UTC), ), } BYTE_SIZES = [ UnitOfInformation.BYTES, UnitOfInformation.KILOBYTES, UnitOfInformation.MEGABYTES, UnitOfInformation.GIGABYTES, ] PARALLEL_UPDATES = 1 async def async_setup_entry( hass: HomeAssistant, entry: RadarrConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up Radarr sensors based on a config entry.""" entities: list[RadarrSensor[Any]] = [] for coordinator_type, description in SENSOR_TYPES.items(): coordinator = getattr(entry.runtime_data, coordinator_type) if coordinator_type != "disk_space": entities.append(RadarrSensor(coordinator, description)) else: entities.extend( RadarrSensor(coordinator, *get_modified_description(description, mount)) for mount in coordinator.data if description.description_fn ) async_add_entities(entities) class RadarrSensor(RadarrEntity[T], SensorEntity): """Implementation of the Radarr sensor.""" coordinator: RadarrDataUpdateCoordinator[T] entity_description: RadarrSensorEntityDescription[T] def __init__( self, coordinator: RadarrDataUpdateCoordinator[T], description: RadarrSensorEntityDescription[T], folder_name: str = "", ) -> None: """Create Radarr entity.""" super().__init__(coordinator, description) self.folder_name = folder_name @property def native_value(self) -> str | int | datetime: """Return the state of the sensor.""" return self.entity_description.value_fn(self.coordinator.data, self.folder_name)