Reduce complexity of shelly button setup (#109625)

This commit is contained in:
J. Nick Koston 2024-02-04 13:38:36 -06:00 committed by GitHub
parent 7572a73c16
commit a7c074e388
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -3,6 +3,7 @@ from __future__ import annotations
from collections.abc import Callable, Coroutine
from dataclasses import dataclass
from functools import partial
from typing import TYPE_CHECKING, Any, Final, Generic, TypeVar
from aioshelly.const import RPC_GENERATIONS
@ -83,8 +84,8 @@ BUTTONS: Final[list[ShellyButtonDescription[Any]]] = [
@callback
def async_migrate_unique_ids(
entity_entry: er.RegistryEntry,
coordinator: ShellyRpcCoordinator | ShellyBlockCoordinator,
entity_entry: er.RegistryEntry,
) -> dict[str, Any] | None:
"""Migrate button unique IDs."""
if not entity_entry.entity_id.startswith("button"):
@ -117,35 +118,25 @@ async def async_setup_entry(
async_add_entities: AddEntitiesCallback,
) -> None:
"""Set buttons for device."""
@callback
def _async_migrate_unique_ids(
entity_entry: er.RegistryEntry,
) -> dict[str, Any] | None:
"""Migrate button unique IDs."""
if TYPE_CHECKING:
assert coordinator is not None
return async_migrate_unique_ids(entity_entry, coordinator)
coordinator: ShellyRpcCoordinator | ShellyBlockCoordinator | None = None
entry_data = get_entry_data(hass)[config_entry.entry_id]
coordinator: ShellyRpcCoordinator | ShellyBlockCoordinator | None
if get_device_entry_gen(config_entry) in RPC_GENERATIONS:
coordinator = get_entry_data(hass)[config_entry.entry_id].rpc
coordinator = entry_data.rpc
else:
coordinator = get_entry_data(hass)[config_entry.entry_id].block
coordinator = entry_data.block
if coordinator is not None:
await er.async_migrate_entries(
hass, config_entry.entry_id, _async_migrate_unique_ids
)
if TYPE_CHECKING:
assert coordinator is not None
entities: list[ShellyButton] = []
await er.async_migrate_entries(
hass, config_entry.entry_id, partial(async_migrate_unique_ids, coordinator)
)
for button in BUTTONS:
if not button.supported(coordinator):
continue
entities.append(ShellyButton(coordinator, button))
async_add_entities(entities)
async_add_entities(
ShellyButton(coordinator, button)
for button in BUTTONS
if button.supported(coordinator)
)
class ShellyButton(