evlli/hass-core
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Refactor Shelly wrapper to coordinator (#79628)

Browse source
This commit is contained in:
Shay Levy 2022-10-05 15:39:58 +03:00 • committed by GitHub
parent 4d3d22320f
commit 22c68b95bf
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
14 changed files with 336 additions and 321 deletions
Download patch file Download diff file Expand all files Collapse all files

49
homeassistant/components/shelly/button.py
View file

@ -15,10 +15,11 @@ 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.helpers.update_coordinator import CoordinatorEntity
from homeassistant.util import slugify
from . import BlockDeviceWrapper, RpcDeviceWrapper
from .const import BLOCK, DATA_CONFIG_ENTRY, DOMAIN, RPC, SHELLY_GAS_MODELS
from .coordinator import ShellyBlockCoordinator, ShellyRpcCoordinator
from .utils import get_block_device_name, get_device_entry_gen, get_rpc_device_name
@ -42,31 +43,31 @@ BUTTONS: Final = [
name="Reboot",
device_class=ButtonDeviceClass.RESTART,
entity_category=EntityCategory.CONFIG,
press_action=lambda wrapper: wrapper.device.trigger_reboot(),
press_action=lambda coordinator: coordinator.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,
press_action=lambda coordinator: coordinator.device.trigger_shelly_gas_self_test(),
supported=lambda coordinator: coordinator.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,
press_action=lambda coordinator: coordinator.device.trigger_shelly_gas_mute(),
supported=lambda coordinator: coordinator.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,
press_action=lambda coordinator: coordinator.device.trigger_shelly_gas_unmute(),
supported=lambda coordinator: coordinator.device.model in SHELLY_GAS_MODELS,
),
]
@ -77,54 +78,54 @@ async def async_setup_entry(
async_add_entities: AddEntitiesCallback,
) -> None:
"""Set buttons for device."""
wrapper: RpcDeviceWrapper | BlockDeviceWrapper | None = None
coordinator: ShellyRpcCoordinator | ShellyBlockCoordinator | None = None
if get_device_entry_gen(config_entry) == 2:
if rpc_wrapper := hass.data[DOMAIN][DATA_CONFIG_ENTRY][
if rpc_coordinator := hass.data[DOMAIN][DATA_CONFIG_ENTRY][
config_entry.entry_id
].get(RPC):
wrapper = cast(RpcDeviceWrapper, rpc_wrapper)
coordinator = cast(ShellyRpcCoordinator, rpc_coordinator)
else:
if block_wrapper := hass.data[DOMAIN][DATA_CONFIG_ENTRY][
if block_coordinator := hass.data[DOMAIN][DATA_CONFIG_ENTRY][
config_entry.entry_id
].get(BLOCK):
wrapper = cast(BlockDeviceWrapper, block_wrapper)
coordinator = cast(ShellyBlockCoordinator, block_coordinator)
if wrapper is not None:
if coordinator is not None:
entities = []
for button in BUTTONS:
if not button.supported(wrapper):
if not button.supported(coordinator):
continue
entities.append(ShellyButton(wrapper, button))
entities.append(ShellyButton(coordinator, button))
async_add_entities(entities)
class ShellyButton(ButtonEntity):
class ShellyButton(CoordinatorEntity, ButtonEntity):
"""Defines a Shelly base button."""
entity_description: ShellyButtonDescription
def __init__(
self,
wrapper: RpcDeviceWrapper | BlockDeviceWrapper,
coordinator: ShellyRpcCoordinator | ShellyBlockCoordinator,
description: ShellyButtonDescription,
) -> None:
"""Initialize Shelly button."""
super().__init__(coordinator)
self.entity_description = description
self.wrapper = wrapper
if isinstance(wrapper, RpcDeviceWrapper):
device_name = get_rpc_device_name(wrapper.device)
if isinstance(coordinator, ShellyRpcCoordinator):
device_name = get_rpc_device_name(coordinator.device)
else:
device_name = get_block_device_name(wrapper.device)
device_name = get_block_device_name(coordinator.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)}
connections={(CONNECTION_NETWORK_MAC, coordinator.mac)}
)
async def async_press(self) -> None:
"""Triggers the Shelly button press service."""
await self.entity_description.press_action(self.wrapper)
await self.entity_description.press_action(self.coordinator)