"""Support for Imou button controls.""" from pyimouapi.exceptions import ImouException from pyimouapi.ha_device import ImouHaDevice from homeassistant.components.button import ButtonDeviceClass, ButtonEntity from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback from .const import PTZ_MOVE_DURATION_MS, imou_device_identifier from .coordinator import ImouConfigEntry, ImouDataUpdateCoordinator from .entity import ImouEntity PARALLEL_UPDATES = 1 # Button types PARAM_RESTART_DEVICE = "restart_device" PARAM_MUTE = "mute" PARAM_PTZ_UP = "ptz_up" PARAM_PTZ_DOWN = "ptz_down" PARAM_PTZ_LEFT = "ptz_left" PARAM_PTZ_RIGHT = "ptz_right" BUTTON_TYPES = ( PARAM_RESTART_DEVICE, PARAM_MUTE, PARAM_PTZ_UP, PARAM_PTZ_DOWN, PARAM_PTZ_LEFT, PARAM_PTZ_RIGHT, ) PTZ_BUTTON_TYPES = ( PARAM_PTZ_UP, PARAM_PTZ_DOWN, PARAM_PTZ_LEFT, PARAM_PTZ_RIGHT, ) BUTTON_DEVICE_CLASS: dict[str, ButtonDeviceClass] = { PARAM_RESTART_DEVICE: ButtonDeviceClass.RESTART, } def _iter_buttons( coordinator: ImouDataUpdateCoordinator, ) -> list[tuple[str, ImouHaDevice]]: """Return (button_type, device) pairs for supported buttons.""" return [ (button_type, device) for device in coordinator.devices for button_type in device.buttons if button_type in BUTTON_TYPES ] async def async_setup_entry( hass: HomeAssistant, entry: ImouConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up Imou button entities.""" coordinator = entry.runtime_data def _add_buttons(new_devices: list[ImouHaDevice]) -> None: device_keys = {imou_device_identifier(device) for device in new_devices} async_add_entities( ImouButton(coordinator, button_type, device) for button_type, device in _iter_buttons(coordinator) if imou_device_identifier(device) in device_keys ) coordinator.new_device_callbacks.append(_add_buttons) @callback def _remove_new_device_callback() -> None: if _add_buttons in coordinator.new_device_callbacks: coordinator.new_device_callbacks.remove(_add_buttons) entry.async_on_unload(_remove_new_device_callback) _add_buttons(coordinator.devices) class ImouButton(ImouEntity, ButtonEntity): """Imou button entity.""" def __init__( self, coordinator: ImouDataUpdateCoordinator, entity_type: str, device: ImouHaDevice, ) -> None: """Initialize the Imou button entity.""" super().__init__(coordinator, entity_type, device) if device_class := BUTTON_DEVICE_CLASS.get(entity_type): self._attr_device_class = device_class self._attr_translation_key = None async def async_press(self) -> None: """Handle button press.""" duration = PTZ_MOVE_DURATION_MS if self._entity_type in PTZ_BUTTON_TYPES else 0 try: await self.coordinator.device_manager.async_press_button( self.device, self._entity_type, duration, ) except ImouException as e: raise HomeAssistantError(str(e)) from e