image.py

Go to the documentation of this file.

"""Image platform for UniFi Network integration.
 
Support for QR code for guest WLANs.
"""
 
from __future__ import annotations
 
from collections.abc import Callable
from dataclasses import dataclass
 
from aiounifi.interfaces.api_handlers import ItemEvent
from aiounifi.interfaces.wlans import Wlans
from aiounifi.models.api import ApiItemT
from aiounifi.models.wlan import Wlan
 
from homeassistant.components.image import ImageEntity, ImageEntityDescription
from homeassistant.const import EntityCategory
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.entity_platform import AddEntitiesCallback
import homeassistant.util.dt as dt_util
 
from . import UnifiConfigEntry
from .entity import (
    HandlerT,
    UnifiEntity,
    UnifiEntityDescription,
    async_wlan_available_fn,
    async_wlan_device_info_fn,
)
from .hub import UnifiHub
 
 
@callback
def async_wlan_qr_code_image_fn(hub: UnifiHub, wlan: Wlan) -> bytes:
    """Calculate receiving data transfer value."""
    return hub.api.wlans.generate_wlan_qr_code(wlan)
 
 
@dataclass(frozen=True, kw_only=True)
class UnifiImageEntityDescription(
    ImageEntityDescription, UnifiEntityDescription[HandlerT, ApiItemT]
):
    """Class describing UniFi image entity."""
 
    image_fn: Callable[[UnifiHub, ApiItemT], bytes]
    value_fn: Callable[[ApiItemT], str | None]
 
 
ENTITY_DESCRIPTIONS: tuple[UnifiImageEntityDescription, ...] = (
    UnifiImageEntityDescription[Wlans, Wlan](
        key="WLAN QR Code",
        translation_key="wlan_qr_code",
        entity_category=EntityCategory.DIAGNOSTIC,
        entity_registry_enabled_default=False,
        api_handler_fn=lambda api: api.wlans,
        available_fn=async_wlan_available_fn,
        device_info_fn=async_wlan_device_info_fn,
        name_fn=lambda wlan: "QR Code",
        object_fn=lambda api, obj_id: api.wlans[obj_id],
        unique_id_fn=lambda hub, obj_id: f"qr_code-{obj_id}",
        image_fn=async_wlan_qr_code_image_fn,
        value_fn=lambda obj: obj.x_passphrase,
    ),
)
 
 
async def async_setup_entry(
    hass: HomeAssistant,
    config_entry: UnifiConfigEntry,
    async_add_entities: AddEntitiesCallback,
) -> None:
    """Set up image platform for UniFi Network integration."""
    config_entry.runtime_data.entity_loader.register_platform(
        async_add_entities, UnifiImageEntity, ENTITY_DESCRIPTIONS, requires_admin=True
    )
 
 
class UnifiImageEntity(UnifiEntity[HandlerT, ApiItemT], ImageEntity):
    """Base representation of a UniFi image."""
 
    entity_description: UnifiImageEntityDescription[HandlerT, ApiItemT]
    _attr_content_type = "image/png"
 
    current_image: bytes | None = None
    previous_value: str | None = None
 
    def __init__(
        self,
        obj_id: str,
        hub: UnifiHub,
        description: UnifiEntityDescription[HandlerT, ApiItemT],
    ) -> None:
        """Initiatlize UniFi Image entity."""
        super().__init__(obj_id, hub, description)
        ImageEntity.__init__(self, hub.hass)
 
    def image(self) -> bytes | None:
        """Return bytes of image."""
        if self.current_imagecurrent_image is None:
            description = self.entity_descriptionentity_description
            obj = description.object_fn(self.apiapi, self._obj_id_obj_id)
            self.current_imagecurrent_image = description.image_fn(self.hubhub, obj)
        return self.current_imagecurrent_image
 
    @callback
    def async_update_state(self, event: ItemEvent, obj_id: str) -> None:
        """Update entity state."""
        description = self.entity_descriptionentity_description
        obj = description.object_fn(self.apiapi, self._obj_id_obj_id)
        if (value := description.value_fn(obj)) != self.previous_valueprevious_value:
            self.previous_valueprevious_value = value
            self.current_imagecurrent_image = None
            self._attr_image_last_updated_attr_image_last_updated = dt_util.utcnow()