core/homeassistant/components/shelly/button.py

146 lines
4.7 KiB
Python

"""Button for Shelly."""
from __future__ import annotations
from collections.abc import Callable
from dataclasses import dataclass
from typing import Final, cast
from homeassistant.components.button import (
ButtonDeviceClass,
ButtonEntity,
ButtonEntityDescription,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant
from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
from homeassistant.helpers.entity import DeviceInfo, EntityCategory
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.util import slugify
from . import BlockDeviceWrapper, RpcDeviceWrapper
from .const import BLOCK, DATA_CONFIG_ENTRY, DOMAIN, RPC, SHELLY_GAS_MODELS
from .utils import get_block_device_name, get_device_entry_gen, get_rpc_device_name
@dataclass
class ShellyButtonDescriptionMixin:
"""Mixin to describe a Button entity."""
press_action: Callable
@dataclass
class ShellyButtonDescription(ButtonEntityDescription, ShellyButtonDescriptionMixin):
"""Class to describe a Button entity."""
supported: Callable = lambda _: True
BUTTONS: Final = [
ShellyButtonDescription(
key="ota_update",
name="OTA Update",
device_class=ButtonDeviceClass.UPDATE,
entity_category=EntityCategory.CONFIG,
press_action=lambda wrapper: wrapper.async_trigger_ota_update(),
),
ShellyButtonDescription(
key="ota_update_beta",
name="OTA Update Beta",
device_class=ButtonDeviceClass.UPDATE,
entity_registry_enabled_default=False,
entity_category=EntityCategory.CONFIG,
press_action=lambda wrapper: wrapper.async_trigger_ota_update(beta=True),
),
ShellyButtonDescription(
key="reboot",
name="Reboot",
device_class=ButtonDeviceClass.RESTART,
entity_category=EntityCategory.CONFIG,
press_action=lambda wrapper: wrapper.device.trigger_reboot(),
),
ShellyButtonDescription(
key="self_test",
name="Self Test",
icon="mdi:progress-wrench",
entity_category=EntityCategory.DIAGNOSTIC,
press_action=lambda wrapper: wrapper.device.trigger_shelly_gas_self_test(),
supported=lambda wrapper: wrapper.device.model in SHELLY_GAS_MODELS,
),
ShellyButtonDescription(
key="mute",
name="Mute",
icon="mdi:volume-mute",
entity_category=EntityCategory.CONFIG,
press_action=lambda wrapper: wrapper.device.trigger_shelly_gas_mute(),
supported=lambda wrapper: wrapper.device.model in SHELLY_GAS_MODELS,
),
ShellyButtonDescription(
key="unmute",
name="Unmute",
icon="mdi:volume-high",
entity_category=EntityCategory.CONFIG,
press_action=lambda wrapper: wrapper.device.trigger_shelly_gas_unmute(),
supported=lambda wrapper: wrapper.device.model in SHELLY_GAS_MODELS,
),
]
async def async_setup_entry(
hass: HomeAssistant,
config_entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
) -> None:
"""Set buttons for device."""
wrapper: RpcDeviceWrapper | BlockDeviceWrapper | None = None
if get_device_entry_gen(config_entry) == 2:
if rpc_wrapper := hass.data[DOMAIN][DATA_CONFIG_ENTRY][
config_entry.entry_id
].get(RPC):
wrapper = cast(RpcDeviceWrapper, rpc_wrapper)
else:
if block_wrapper := hass.data[DOMAIN][DATA_CONFIG_ENTRY][
config_entry.entry_id
].get(BLOCK):
wrapper = cast(BlockDeviceWrapper, block_wrapper)
if wrapper is not None:
entities = []
for button in BUTTONS:
if not button.supported(wrapper):
continue
entities.append(ShellyButton(wrapper, button))
async_add_entities(entities)
class ShellyButton(ButtonEntity):
"""Defines a Shelly base button."""
entity_description: ShellyButtonDescription
def __init__(
self,
wrapper: RpcDeviceWrapper | BlockDeviceWrapper,
description: ShellyButtonDescription,
) -> None:
"""Initialize Shelly button."""
self.entity_description = description
self.wrapper = wrapper
if isinstance(wrapper, RpcDeviceWrapper):
device_name = get_rpc_device_name(wrapper.device)
else:
device_name = get_block_device_name(wrapper.device)
self._attr_name = f"{device_name} {description.name}"
self._attr_unique_id = slugify(self._attr_name)
self._attr_device_info = DeviceInfo(
connections={(CONNECTION_NETWORK_MAC, wrapper.mac)}
)
async def async_press(self) -> None:
"""Triggers the Shelly button press service."""
await self.entity_description.press_action(self.wrapper)