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