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

Improve shelly generic typing (#84889)

Browse source 
Improve shelly generic typing
This commit is contained in:
Marc Mueller 2022-12-31 07:33:41 +01:00 • committed by GitHub
parent f6991cd4dc
commit 11b5de9463
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 29 additions and 17 deletions
Download patch file Download diff file Expand all files Collapse all files

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

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

4
homeassistant/components/shelly/coordinator.py
View file

@ -267,7 +267,7 @@ class ShellyBlockCoordinator(DataUpdateCoordinator[None]):
self.shutdown() self.shutdown()
class ShellyRestCoordinator(DataUpdateCoordinator): class ShellyRestCoordinator(DataUpdateCoordinator[None]):
"""Coordinator for a Shelly REST device.""" """Coordinator for a Shelly REST device."""
def __init__( def __init__(
@ -579,7 +579,7 @@ class ShellyRpcCoordinator(DataUpdateCoordinator[None]):
await self.shutdown() await self.shutdown()
class ShellyRpcPollingCoordinator(DataUpdateCoordinator): class ShellyRpcPollingCoordinator(DataUpdateCoordinator[None]):
"""Polling coordinator for a Shelly RPC based device.""" """Polling coordinator for a Shelly RPC based device."""
def __init__( def __init__(