"""Support for Litter-Robot sensors.""" from __future__ import annotations from collections.abc import Callable from dataclasses import dataclass from datetime import datetime from typing import Any, Generic from pylitterbot import FeederRobot, LitterRobot, LitterRobot4, Robot from homeassistant.components.sensor import ( SensorDeviceClass, SensorEntity, SensorEntityDescription, SensorStateClass, ) from homeassistant.const import PERCENTAGE, EntityCategory, UnitOfMass from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback from .coordinator import LitterRobotConfigEntry from .entity import LitterRobotEntity, _RobotT def icon_for_gauge_level(gauge_level: int | None = None, offset: int = 0) -> str: """Return a gauge icon valid identifier.""" if gauge_level is None or gauge_level <= 0 + offset: return "mdi:gauge-empty" if gauge_level > 70 + offset: return "mdi:gauge-full" if gauge_level > 30 + offset: return "mdi:gauge" return "mdi:gauge-low" @dataclass(frozen=True, kw_only=True) class RobotSensorEntityDescription(SensorEntityDescription, Generic[_RobotT]): """A class that describes robot sensor entities.""" icon_fn: Callable[[Any], str | None] = lambda _: None value_fn: Callable[[_RobotT], float | datetime | str | None] ROBOT_SENSOR_MAP: dict[type[Robot], list[RobotSensorEntityDescription]] = { LitterRobot: [ # type: ignore[type-abstract] # only used for isinstance check RobotSensorEntityDescription[LitterRobot]( key="waste_drawer_level", translation_key="waste_drawer", native_unit_of_measurement=PERCENTAGE, icon_fn=lambda state: icon_for_gauge_level(state, 10), state_class=SensorStateClass.MEASUREMENT, value_fn=lambda robot: robot.waste_drawer_level, ), RobotSensorEntityDescription[LitterRobot]( key="sleep_mode_start_time", translation_key="sleep_mode_start_time", device_class=SensorDeviceClass.TIMESTAMP, value_fn=( lambda robot: robot.sleep_mode_start_time if robot.sleep_mode_enabled else None ), ), RobotSensorEntityDescription[LitterRobot]( key="sleep_mode_end_time", translation_key="sleep_mode_end_time", device_class=SensorDeviceClass.TIMESTAMP, value_fn=( lambda robot: robot.sleep_mode_end_time if robot.sleep_mode_enabled else None ), ), RobotSensorEntityDescription[LitterRobot]( key="last_seen", translation_key="last_seen", device_class=SensorDeviceClass.TIMESTAMP, entity_category=EntityCategory.DIAGNOSTIC, value_fn=lambda robot: robot.last_seen, ), RobotSensorEntityDescription[LitterRobot]( key="status_code", translation_key="status_code", entity_category=EntityCategory.DIAGNOSTIC, device_class=SensorDeviceClass.ENUM, options=[ "br", "ccc", "ccp", "cd", "csf", "csi", "cst", "df1", "df2", "dfs", "dhf", "dpf", "ec", "hpf", "off", "offline", "otf", "p", "pd", "pwrd", "pwru", "rdy", "scf", "sdf", "spf", ], value_fn=( lambda robot: status.lower() if (status := robot.status_code) else None ), ), ], LitterRobot4: [ RobotSensorEntityDescription[LitterRobot4]( key="litter_level", translation_key="litter_level", native_unit_of_measurement=PERCENTAGE, icon_fn=lambda state: icon_for_gauge_level(state, 10), state_class=SensorStateClass.MEASUREMENT, value_fn=lambda robot: robot.litter_level, ), RobotSensorEntityDescription[LitterRobot4]( key="pet_weight", translation_key="pet_weight", native_unit_of_measurement=UnitOfMass.POUNDS, device_class=SensorDeviceClass.WEIGHT, state_class=SensorStateClass.MEASUREMENT, value_fn=lambda robot: robot.pet_weight, ), ], FeederRobot: [ RobotSensorEntityDescription[FeederRobot]( key="food_level", translation_key="food_level", native_unit_of_measurement=PERCENTAGE, icon_fn=lambda state: icon_for_gauge_level(state, 10), state_class=SensorStateClass.MEASUREMENT, value_fn=lambda robot: robot.food_level, ) ], } async def async_setup_entry( hass: HomeAssistant, entry: LitterRobotConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up Litter-Robot sensors using config entry.""" coordinator = entry.runtime_data async_add_entities( LitterRobotSensorEntity( robot=robot, coordinator=coordinator, description=description ) for robot in coordinator.account.robots for robot_type, entity_descriptions in ROBOT_SENSOR_MAP.items() if isinstance(robot, robot_type) for description in entity_descriptions ) class LitterRobotSensorEntity(LitterRobotEntity[_RobotT], SensorEntity): """Litter-Robot sensor entity.""" entity_description: RobotSensorEntityDescription[_RobotT] @property def native_value(self) -> float | datetime | str | None: """Return the state.""" return self.entity_description.value_fn(self.robot) @property def icon(self) -> str | None: """Return the icon to use in the frontend, if any.""" if (icon := self.entity_description.icon_fn(self.state)) is not None: return icon return super().icon