From a7c074e388073c911917e287818b9f3b99f9a034 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sun, 4 Feb 2024 13:38:36 -0600
Subject: [PATCH] Reduce complexity of shelly button setup (#109625)

---
 homeassistant/components/shelly/button.py | 41 +++++++++--------------
 1 file changed, 16 insertions(+), 25 deletions(-)

diff --git a/homeassistant/components/shelly/button.py b/homeassistant/components/shelly/button.py
index 17f60f566aa..6dd97731ab5 100644
--- a/homeassistant/components/shelly/button.py
+++ b/homeassistant/components/shelly/button.py
@@ -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(