Improve shelly generic typing (#84889)

Improve shelly generic typing

homeassistant/components/shelly/button.py

@@ -1,9 +1,9 @@
 """Button for Shelly."""
 from __future__ import annotations
 
-from collections.abc import Callable
+from collections.abc import Callable, Coroutine
 from dataclasses import dataclass
-from typing import Final
+from typing import Any, Final, Generic, TypeVar, Union
 
 from homeassistant.components.button import (
     ButtonDeviceClass,
@@ -22,30 +22,36 @@ from .const import SHELLY_GAS_MODELS
 from .coordinator import ShellyBlockCoordinator, ShellyRpcCoordinator, get_entry_data
 from .utils import get_block_device_name, get_device_entry_gen, get_rpc_device_name
 
+_ShellyCoordinatorT = TypeVar(
+    "_ShellyCoordinatorT", bound=Union[ShellyBlockCoordinator, ShellyRpcCoordinator]
+)
+
 
 @dataclass
-class ShellyButtonDescriptionMixin:
+class ShellyButtonDescriptionMixin(Generic[_ShellyCoordinatorT]):
     """Mixin to describe a Button entity."""
 
-    press_action: Callable
+    press_action: Callable[[_ShellyCoordinatorT], Coroutine[Any, Any, None]]
 
 
 @dataclass
-class ShellyButtonDescription(ButtonEntityDescription, ShellyButtonDescriptionMixin):
+class ShellyButtonDescription(
+    ButtonEntityDescription, ShellyButtonDescriptionMixin[_ShellyCoordinatorT]
+):
     """Class to describe a Button entity."""
 
-    supported: Callable = lambda _: True
+    supported: Callable[[_ShellyCoordinatorT], bool] = lambda _: True
 
 
-BUTTONS: Final = [
-    ShellyButtonDescription(
+BUTTONS: Final[list[ShellyButtonDescription[Any]]] = [
+    ShellyButtonDescription[Union[ShellyBlockCoordinator, ShellyRpcCoordinator]](
         key="reboot",
         name="Reboot",
         device_class=ButtonDeviceClass.RESTART,
         entity_category=EntityCategory.CONFIG,
         press_action=lambda coordinator: coordinator.device.trigger_reboot(),
     ),
-    ShellyButtonDescription(
+    ShellyButtonDescription[ShellyBlockCoordinator](
         key="self_test",
         name="Self Test",
         icon="mdi:progress-wrench",
@@ -53,7 +59,7 @@ BUTTONS: Final = [
         press_action=lambda coordinator: coordinator.device.trigger_shelly_gas_self_test(),
         supported=lambda coordinator: coordinator.device.model in SHELLY_GAS_MODELS,
     ),
-    ShellyButtonDescription(
+    ShellyButtonDescription[ShellyBlockCoordinator](
         key="mute",
         name="Mute",
         icon="mdi:volume-mute",
@@ -61,7 +67,7 @@ BUTTONS: Final = [
         press_action=lambda coordinator: coordinator.device.trigger_shelly_gas_mute(),
         supported=lambda coordinator: coordinator.device.model in SHELLY_GAS_MODELS,
     ),
-    ShellyButtonDescription(
+    ShellyButtonDescription[ShellyBlockCoordinator](
         key="unmute",
         name="Unmute",
         icon="mdi:volume-high",
@@ -85,7 +91,7 @@ async def async_setup_entry(
         coordinator = get_entry_data(hass)[config_entry.entry_id].block
 
     if coordinator is not None:
-        entities = []
+        entities: list[ShellyButton] = []
 
         for button in BUTTONS:
             if not button.supported(coordinator):
@@ -95,15 +101,21 @@ async def async_setup_entry(
         async_add_entities(entities)
 
 
-class ShellyButton(CoordinatorEntity, ButtonEntity):
+class ShellyButton(
+    CoordinatorEntity[Union[ShellyRpcCoordinator, ShellyBlockCoordinator]], ButtonEntity
+):
     """Defines a Shelly base button."""
 
-    entity_description: ShellyButtonDescription
+    entity_description: ShellyButtonDescription[
+        ShellyRpcCoordinator | ShellyBlockCoordinator
+    ]
 
     def __init__(
         self,
         coordinator: ShellyRpcCoordinator | ShellyBlockCoordinator,
-        description: ShellyButtonDescription,
+        description: ShellyButtonDescription[
+            ShellyRpcCoordinator | ShellyBlockCoordinator
+        ],
     ) -> None:
         """Initialize Shelly button."""
         super().__init__(coordinator)