core/homeassistant/components/mystrom/sensor.py

"""Support for myStrom sensors of switches/plugs."""

from collections.abc import Callable, Coroutine
from dataclasses import dataclass
from datetime import datetime, timedelta
import logging
from typing import Any, override

from pymystrom.exceptions import MyStromConnectionError
from pymystrom.pir import MyStromPir
from pymystrom.switch import MyStromSwitch

from homeassistant.components.sensor import (
    SensorDeviceClass,
    SensorEntity,
    SensorEntityDescription,
    SensorStateClass,
)
from homeassistant.const import (
    LIGHT_LUX,
    EntityCategory,
    UnitOfEnergy,
    UnitOfPower,
    UnitOfTemperature,
)
from homeassistant.core import HomeAssistant
from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC, DeviceInfo
from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
from homeassistant.util.dt import utcnow

from .const import DOMAIN, MANUFACTURER
from .models import MyStromConfigEntry

_LOGGER = logging.getLogger(__name__)


@dataclass(frozen=True, kw_only=True)
class MyStromSensorEntityDescription[_DeviceT](SensorEntityDescription):
    """Class describing mystrom sensor entities."""

    value_fn: Callable[[_DeviceT], float | None]
    # Only needed where nothing else on the device polls; a switch is kept
    # fresh by its own entity refreshing the shared device.
    update_fn: Callable[[_DeviceT], Coroutine[Any, Any, None]] | None = None


SENSOR_TYPES_PIR: tuple[MyStromSensorEntityDescription[MyStromPir], ...] = (
    MyStromSensorEntityDescription(
        key="temperature",
        device_class=SensorDeviceClass.TEMPERATURE,
        state_class=SensorStateClass.MEASUREMENT,
        native_unit_of_measurement=UnitOfTemperature.CELSIUS,
        value_fn=(
            lambda device: (
                float(device.temperature_compensated)
                if device.temperature_compensated is not None
                else None
            )
        ),
        update_fn=lambda device: device.get_temperatures(),
    ),
    MyStromSensorEntityDescription(
        key="illuminance",
        device_class=SensorDeviceClass.ILLUMINANCE,
        state_class=SensorStateClass.MEASUREMENT,
        native_unit_of_measurement=LIGHT_LUX,
        value_fn=(
            lambda device: (
                float(device.intensity) if device.intensity is not None else None
            )
        ),
        update_fn=lambda device: device.get_light(),
    ),
)


SENSOR_TYPES_SWITCH: tuple[MyStromSensorEntityDescription[MyStromSwitch], ...] = (
    MyStromSensorEntityDescription(
        key="avg_consumption",
        translation_key="avg_consumption",
        device_class=SensorDeviceClass.POWER,
        native_unit_of_measurement=UnitOfPower.WATT,
        value_fn=lambda device: device.consumedWs,
    ),
    MyStromSensorEntityDescription(
        key="consumption",
        device_class=SensorDeviceClass.POWER,
        state_class=SensorStateClass.MEASUREMENT,
        native_unit_of_measurement=UnitOfPower.WATT,
        value_fn=lambda device: device.consumption,
    ),
    MyStromSensorEntityDescription(
        key="energy_since_boot",
        translation_key="energy_since_boot",
        device_class=SensorDeviceClass.ENERGY,
        state_class=SensorStateClass.TOTAL_INCREASING,
        native_unit_of_measurement=UnitOfEnergy.JOULE,
        suggested_unit_of_measurement=UnitOfEnergy.KILO_WATT_HOUR,
        value_fn=lambda device: device.energy_since_boot,
    ),
    MyStromSensorEntityDescription(
        key="temperature",
        device_class=SensorDeviceClass.TEMPERATURE,
        state_class=SensorStateClass.MEASUREMENT,
        native_unit_of_measurement=UnitOfTemperature.CELSIUS,
        value_fn=lambda device: device.temperature,
    ),
)


async def async_setup_entry(
    hass: HomeAssistant,
    entry: MyStromConfigEntry,
    async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
    """Set up the myStrom entities."""
    device = entry.runtime_data.device
    info = entry.runtime_data.info

    entities: list[MyStromSensorBase] = []
    match device:
        case MyStromPir():
            entities = [
                MyStromSensor(device, entry.title, description, info["mac"])
                for description in SENSOR_TYPES_PIR
                if description.value_fn(device) is not None
            ]
        case MyStromSwitch():
            entities = [
                MyStromSensor(device, entry.title, description, info["mac"])
                for description in SENSOR_TYPES_SWITCH
                if description.value_fn(device) is not None
            ]
            if device.time_since_boot is not None:
                entities.append(
                    MyStromSwitchUptimeSensor(device, entry.title, info["mac"])
                )
        case _:
            entities = []

    async_add_entities(entities)


class MyStromSensorBase(SensorEntity):
    """Base class for myStrom sensors."""

    _attr_has_entity_name = True

    def __init__(
        self,
        device: MyStromSwitch,
        name: str,
        key: str,
        mac: str,
    ) -> None:
        """Initialize the sensor."""
        self.device = device
        self._attr_unique_id = f"{mac}-{key}"
        self._attr_device_info = DeviceInfo(
            identifiers={(DOMAIN, mac)},
            connections={(CONNECTION_NETWORK_MAC, mac)},
            name=name,
            manufacturer=MANUFACTURER,
            sw_version=getattr(device, "firmware", None),
        )


class MyStromSensor[_DeviceT](MyStromSensorBase):
    """Representation of the consumption or temperature of a myStrom switch/plug."""

    entity_description: MyStromSensorEntityDescription[_DeviceT]

    _attr_has_entity_name = True

    def __init__(
        self,
        device: _DeviceT,
        name: str,
        description: MyStromSensorEntityDescription[_DeviceT],
        mac: str,
    ) -> None:
        """Initialize the sensor."""
        super().__init__(device, name, description.key, mac)
        self.entity_description = description

    @property
    @override
    def native_value(self) -> float | None:
        """Return the value of the sensor."""
        return self.entity_description.value_fn(self.device)

    async def async_update(self) -> None:
        """Get the latest reading from the device."""
        if (update_fn := self.entity_description.update_fn) is None:
            return

        try:
            await update_fn(self.device)
        except MyStromConnectionError:
            if self.available:
                self._attr_available = False
                _LOGGER.error("No route to myStrom device")
        else:
            self._attr_available = True


class MyStromSwitchUptimeSensor(MyStromSensorBase):
    """Representation of a MyStrom Switch uptime sensor."""

    entity_description = SensorEntityDescription(
        key="time_since_boot",
        device_class=SensorDeviceClass.TIMESTAMP,
        entity_category=EntityCategory.DIAGNOSTIC,
        translation_key="time_since_boot",
    )

    def __init__(
        self,
        device: MyStromSwitch,
        name: str,
        mac: str,
    ) -> None:
        """Initialize the uptime sensor."""
        super().__init__(device, name, self.entity_description.key, mac)
        self._last_value: datetime | None = None
        self._last_attributes: dict[str, Any] = {}

    @property
    @override
    def native_value(self) -> datetime | None:
        """Return the uptime of the device as a datetime."""

        if self.device.time_since_boot is None or self.device.boot_id is None:
            return None

        # Return cached value if boot_id hasn't changed
        if (
            self._last_value is not None
            and self._last_attributes.get("boot_id") == self.device.boot_id
        ):
            return self._last_value

        self._last_value = utcnow() - timedelta(seconds=self.device.time_since_boot)

        return self._last_value

    @property
    @override
    def extra_state_attributes(self) -> dict[str, Any]:
        """Return the optional state attributes."""

        self._last_attributes = {
            "boot_id": self.device.boot_id,
        }

        return self._last_attributes