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

Refactor Shelly wrapper to coordinator (#79628)

Browse source
This commit is contained in:
Shay Levy 2022-10-05 15:39:58 +03:00 • committed by GitHub
parent 4d3d22320f
commit 22c68b95bf
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
14 changed files with 336 additions and 321 deletions
Download patch file Download diff file Expand all files Collapse all files

188
homeassistant/components/shelly/entity.py
View file

@ -11,23 +11,13 @@ import async_timeout
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers import (
device_registry,
entity,
entity_registry,
update_coordinator,
)
from homeassistant.helpers import device_registry, entity, entity_registry
from homeassistant.helpers.entity import DeviceInfo, EntityDescription
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.restore_state import RestoreEntity
from homeassistant.helpers.typing import StateType
from homeassistant.helpers.update_coordinator import CoordinatorEntity
from . import (
BlockDeviceWrapper,
RpcDeviceWrapper,
RpcPollingWrapper,
ShellyDeviceRestWrapper,
)
from .const import (
AIOSHELLY_DEVICE_TIMEOUT_SEC,
BLOCK,
@ -38,6 +28,12 @@ from .const import (
RPC,
RPC_POLL,
)
from .coordinator import (
ShellyBlockCoordinator,
ShellyRestCoordinator,
ShellyRpcCoordinator,
ShellyRpcPollingCoordinator,
)
from .utils import (
async_remove_shelly_entity,
get_block_entity_name,
@ -58,20 +54,20 @@ def async_setup_entry_attribute_entities(
],
) -> None:
"""Set up entities for attributes."""
wrapper: BlockDeviceWrapper = hass.data[DOMAIN][DATA_CONFIG_ENTRY][
coordinator: ShellyBlockCoordinator = hass.data[DOMAIN][DATA_CONFIG_ENTRY][
config_entry.entry_id
][BLOCK]
if wrapper.device.initialized:
if coordinator.device.initialized:
async_setup_block_attribute_entities(
hass, async_add_entities, wrapper, sensors, sensor_class
hass, async_add_entities, coordinator, sensors, sensor_class
)
else:
async_restore_block_attribute_entities(
hass,
config_entry,
async_add_entities,
wrapper,
coordinator,
sensors,
sensor_class,
description_class,
@ -82,16 +78,16 @@ def async_setup_entry_attribute_entities(
def async_setup_block_attribute_entities(
hass: HomeAssistant,
async_add_entities: AddEntitiesCallback,
wrapper: BlockDeviceWrapper,
coordinator: ShellyBlockCoordinator,
sensors: Mapping[tuple[str, str], BlockEntityDescription],
sensor_class: Callable,
) -> None:
"""Set up entities for block attributes."""
blocks = []
assert wrapper.device.blocks
assert coordinator.device.blocks
for block in wrapper.device.blocks:
for block in coordinator.device.blocks:
for sensor_id in block.sensor_ids:
description = sensors.get((block.type, sensor_id))
if description is None:
@ -103,10 +99,10 @@ def async_setup_block_attribute_entities(
# Filter and remove entities that according to settings should not create an entity
if description.removal_condition and description.removal_condition(
wrapper.device.settings, block
coordinator.device.settings, block
):
domain = sensor_class.__module__.split(".")[-1]
unique_id = f"{wrapper.mac}-{block.description}-{sensor_id}"
unique_id = f"{coordinator.mac}-{block.description}-{sensor_id}"
async_remove_shelly_entity(hass, domain, unique_id)
else:
blocks.append((block, sensor_id, description))
@ -116,7 +112,7 @@ def async_setup_block_attribute_entities(
async_add_entities(
[
sensor_class(wrapper, block, sensor_id, description)
sensor_class(coordinator, block, sensor_id, description)
for block, sensor_id, description in blocks
]
)
@ -127,7 +123,7 @@ def async_restore_block_attribute_entities(
hass: HomeAssistant,
config_entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
wrapper: BlockDeviceWrapper,
coordinator: ShellyBlockCoordinator,
sensors: Mapping[tuple[str, str], BlockEntityDescription],
sensor_class: Callable,
description_class: Callable[
@ -152,7 +148,7 @@ def async_restore_block_attribute_entities(
description = description_class(entry)
entities.append(
sensor_class(wrapper, None, attribute, description, entry, sensors)
sensor_class(coordinator, None, attribute, description, entry, sensors)
)
if not entities:
@ -170,40 +166,44 @@ def async_setup_entry_rpc(
sensor_class: Callable,
) -> None:
"""Set up entities for REST sensors."""
wrapper: RpcDeviceWrapper = hass.data[DOMAIN][DATA_CONFIG_ENTRY][
coordinator: ShellyRpcCoordinator = hass.data[DOMAIN][DATA_CONFIG_ENTRY][
config_entry.entry_id
][RPC]
polling_wrapper: RpcPollingWrapper = hass.data[DOMAIN][DATA_CONFIG_ENTRY][
config_entry.entry_id
][RPC_POLL]
polling_coordinator: ShellyRpcPollingCoordinator = hass.data[DOMAIN][
DATA_CONFIG_ENTRY
][config_entry.entry_id][RPC_POLL]
entities = []
for sensor_id in sensors:
description = sensors[sensor_id]
key_instances = get_rpc_key_instances(wrapper.device.status, description.key)
key_instances = get_rpc_key_instances(
coordinator.device.status, description.key
)
for key in key_instances:
# Filter non-existing sensors
if description.sub_key not in wrapper.device.status[
if description.sub_key not in coordinator.device.status[
key
] and not description.supported(wrapper.device.status[key]):
] and not description.supported(coordinator.device.status[key]):
continue
# Filter and remove entities that according to settings/status should not create an entity
if description.removal_condition and description.removal_condition(
wrapper.device.config, wrapper.device.status, key
coordinator.device.config, coordinator.device.status, key
):
domain = sensor_class.__module__.split(".")[-1]
unique_id = f"{wrapper.mac}-{key}-{sensor_id}"
unique_id = f"{coordinator.mac}-{key}-{sensor_id}"
async_remove_shelly_entity(hass, domain, unique_id)
else:
if description.use_polling_wrapper:
if description.use_polling_coordinator:
entities.append(
sensor_class(polling_wrapper, key, sensor_id, description)
sensor_class(polling_coordinator, key, sensor_id, description)
)
else:
entities.append(sensor_class(wrapper, key, sensor_id, description))
entities.append(
sensor_class(coordinator, key, sensor_id, description)
)
if not entities:
return
@ -220,7 +220,7 @@ def async_setup_entry_rest(
sensor_class: Callable,
) -> None:
"""Set up entities for REST sensors."""
wrapper: ShellyDeviceRestWrapper = hass.data[DOMAIN][DATA_CONFIG_ENTRY][
coordinator: ShellyRestCoordinator = hass.data[DOMAIN][DATA_CONFIG_ENTRY][
config_entry.entry_id
][REST]
@ -228,7 +228,7 @@ def async_setup_entry_rest(
for sensor_id in sensors:
description = sensors.get(sensor_id)
if not wrapper.device.settings.get("sleep_mode"):
if not coordinator.device.settings.get("sleep_mode"):
entities.append((sensor_id, description))
if not entities:
@ -236,7 +236,7 @@ def async_setup_entry_rest(
async_add_entities(
[
sensor_class(wrapper, sensor_id, description)
sensor_class(coordinator, sensor_id, description)
for sensor_id, description in entities
]
)
@ -270,7 +270,7 @@ class RpcEntityDescription(EntityDescription, RpcEntityRequiredKeysMixin):
available: Callable[[dict], bool] | None = None
removal_condition: Callable[[dict, dict, str], bool] | None = None
extra_state_attributes: Callable[[dict, dict], dict | None] | None = None
use_polling_wrapper: bool = False
use_polling_coordinator: bool = False
supported: Callable = lambda _: False
@ -282,32 +282,32 @@ class RestEntityDescription(EntityDescription):
extra_state_attributes: Callable[[dict], dict | None] | None = None
class ShellyBlockEntity(entity.Entity):
class ShellyBlockEntity(CoordinatorEntity[ShellyBlockCoordinator]):
"""Helper class to represent a block entity."""
def __init__(self, wrapper: BlockDeviceWrapper, block: Block) -> None:
def __init__(self, coordinator: ShellyBlockCoordinator, block: Block) -> None:
"""Initialize Shelly entity."""
self.wrapper = wrapper
super().__init__(coordinator)
self.block = block
self._attr_name = get_block_entity_name(wrapper.device, block)
self._attr_name = get_block_entity_name(coordinator.device, block)
self._attr_should_poll = False
self._attr_device_info = DeviceInfo(
connections={(device_registry.CONNECTION_NETWORK_MAC, wrapper.mac)}
connections={(device_registry.CONNECTION_NETWORK_MAC, coordinator.mac)}
)
self._attr_unique_id = f"{wrapper.mac}-{block.description}"
self._attr_unique_id = f"{coordinator.mac}-{block.description}"
@property
def available(self) -> bool:
"""Available."""
return self.wrapper.last_update_success
return self.coordinator.last_update_success
async def async_added_to_hass(self) -> None:
"""When entity is added to HASS."""
self.async_on_remove(self.wrapper.async_add_listener(self._update_callback))
self.async_on_remove(self.coordinator.async_add_listener(self._update_callback))
async def async_update(self) -> None:
"""Update entity with latest info."""
await self.wrapper.async_request_refresh()
await self.coordinator.async_request_refresh()
@callback
def _update_callback(self) -> None:
@ -327,7 +327,7 @@ class ShellyBlockEntity(entity.Entity):
kwargs,
repr(err),
)
self.wrapper.last_update_success = False
self.coordinator.last_update_success = False
return None
@ -336,36 +336,36 @@ class ShellyRpcEntity(entity.Entity):
def __init__(
self,
wrapper: RpcDeviceWrapper | RpcPollingWrapper,
coordinator: ShellyRpcCoordinator | ShellyRpcPollingCoordinator,
key: str,
) -> None:
"""Initialize Shelly entity."""
self.wrapper = wrapper
self.coordinator = coordinator
self.key = key
self._attr_should_poll = False
self._attr_device_info = {
"connections": {(device_registry.CONNECTION_NETWORK_MAC, wrapper.mac)}
"connections": {(device_registry.CONNECTION_NETWORK_MAC, coordinator.mac)}
}
self._attr_unique_id = f"{wrapper.mac}-{key}"
self._attr_name = get_rpc_entity_name(wrapper.device, key)
self._attr_unique_id = f"{coordinator.mac}-{key}"
self._attr_name = get_rpc_entity_name(coordinator.device, key)
@property
def available(self) -> bool:
"""Available."""
return self.wrapper.device.connected
return self.coordinator.device.connected
@property
def status(self) -> dict:
"""Device status by entity key."""
return cast(dict, self.wrapper.device.status[self.key])
return cast(dict, self.coordinator.device.status[self.key])
async def async_added_to_hass(self) -> None:
"""When entity is added to HASS."""
self.async_on_remove(self.wrapper.async_add_listener(self._update_callback))
self.async_on_remove(self.coordinator.async_add_listener(self._update_callback))
async def async_update(self) -> None:
"""Update entity with latest info."""
await self.wrapper.async_request_refresh()
await self.coordinator.async_request_refresh()
@callback
def _update_callback(self) -> None:
@ -382,7 +382,7 @@ class ShellyRpcEntity(entity.Entity):
)
try:
async with async_timeout.timeout(AIOSHELLY_DEVICE_TIMEOUT_SEC):
return await self.wrapper.device.call_rpc(method, params)
return await self.coordinator.device.call_rpc(method, params)
except asyncio.TimeoutError as err:
LOGGER.error(
"Call RPC for entity %s failed, method: %s, params: %s, error: %s",
@ -391,7 +391,7 @@ class ShellyRpcEntity(entity.Entity):
params,
repr(err),
)
self.wrapper.last_update_success = False
self.coordinator.last_update_success = False
return None
@ -402,18 +402,20 @@ class ShellyBlockAttributeEntity(ShellyBlockEntity, entity.Entity):
def __init__(
self,
wrapper: BlockDeviceWrapper,
coordinator: ShellyBlockCoordinator,
block: Block,
attribute: str,
description: BlockEntityDescription,
) -> None:
"""Initialize sensor."""
super().__init__(wrapper, block)
super().__init__(coordinator, block)
self.attribute = attribute
self.entity_description = description
self._attr_unique_id: str = f"{super().unique_id}-{self.attribute}"
self._attr_name = get_block_entity_name(wrapper.device, block, description.name)
self._attr_name = get_block_entity_name(
coordinator.device, block, description.name
)
@property
def attribute_value(self) -> StateType:
@ -442,40 +444,42 @@ class ShellyBlockAttributeEntity(ShellyBlockEntity, entity.Entity):
return self.entity_description.extra_state_attributes(self.block)
class ShellyRestAttributeEntity(update_coordinator.CoordinatorEntity):
class ShellyRestAttributeEntity(CoordinatorEntity[ShellyBlockCoordinator]):
"""Class to load info from REST."""
entity_description: RestEntityDescription
def __init__(
self,
wrapper: BlockDeviceWrapper,
coordinator: ShellyBlockCoordinator,
attribute: str,
description: RestEntityDescription,
) -> None:
"""Initialize sensor."""
super().__init__(wrapper)
self.wrapper = wrapper
super().__init__(coordinator)
self.block_coordinator = coordinator
self.attribute = attribute
self.entity_description = description
self._attr_name = get_block_entity_name(wrapper.device, None, description.name)
self._attr_unique_id = f"{wrapper.mac}-{attribute}"
self._attr_name = get_block_entity_name(
coordinator.device, None, description.name
)
self._attr_unique_id = f"{coordinator.mac}-{attribute}"
self._attr_device_info = DeviceInfo(
connections={(device_registry.CONNECTION_NETWORK_MAC, wrapper.mac)}
connections={(device_registry.CONNECTION_NETWORK_MAC, coordinator.mac)}
)
self._last_value = None
@property
def available(self) -> bool:
"""Available."""
return self.wrapper.last_update_success
return self.block_coordinator.last_update_success
@property
def attribute_value(self) -> StateType:
"""Value of sensor."""
if callable(self.entity_description.value):
self._last_value = self.entity_description.value(
self.wrapper.device.status, self._last_value
self.block_coordinator.device.status, self._last_value
)
return self._last_value
@ -486,7 +490,7 @@ class ShellyRestAttributeEntity(update_coordinator.CoordinatorEntity):
return None
return self.entity_description.extra_state_attributes(
self.wrapper.device.status
self.block_coordinator.device.status
)
@ -497,18 +501,18 @@ class ShellyRpcAttributeEntity(ShellyRpcEntity, entity.Entity):
def __init__(
self,
wrapper: RpcDeviceWrapper,
coordinator: ShellyRpcCoordinator,
key: str,
attribute: str,
description: RpcEntityDescription,
) -> None:
"""Initialize sensor."""
super().__init__(wrapper, key)
super().__init__(coordinator, key)
self.attribute = attribute
self.entity_description = description
self._attr_unique_id = f"{super().unique_id}-{attribute}"
self._attr_name = get_rpc_entity_name(wrapper.device, key, description.name)
self._attr_name = get_rpc_entity_name(coordinator.device, key, description.name)
self._last_value = None
@property
@ -516,13 +520,13 @@ class ShellyRpcAttributeEntity(ShellyRpcEntity, entity.Entity):
"""Value of sensor."""
if callable(self.entity_description.value):
self._last_value = self.entity_description.value(
self.wrapper.device.status[self.key].get(
self.coordinator.device.status[self.key].get(
self.entity_description.sub_key
),
self._last_value,
)
else:
self._last_value = self.wrapper.device.status[self.key][
self._last_value = self.coordinator.device.status[self.key][
self.entity_description.sub_key
]
@ -537,7 +541,7 @@ class ShellyRpcAttributeEntity(ShellyRpcEntity, entity.Entity):
return available
return self.entity_description.available(
self.wrapper.device.status[self.key][self.entity_description.sub_key]
self.coordinator.device.status[self.key][self.entity_description.sub_key]
)
@property
@ -546,11 +550,11 @@ class ShellyRpcAttributeEntity(ShellyRpcEntity, entity.Entity):
if self.entity_description.extra_state_attributes is None:
return None
assert self.wrapper.device.shelly
assert self.coordinator.device.shelly
return self.entity_description.extra_state_attributes(
self.wrapper.device.status[self.key][self.entity_description.sub_key],
self.wrapper.device.shelly,
self.coordinator.device.status[self.key][self.entity_description.sub_key],
self.coordinator.device.shelly,
)
@ -560,7 +564,7 @@ class ShellySleepingBlockAttributeEntity(ShellyBlockAttributeEntity, RestoreEnti
# pylint: disable=super-init-not-called
def __init__(
self,
wrapper: BlockDeviceWrapper,
coordinator: ShellyBlockCoordinator,
block: Block | None,
attribute: str,
description: BlockEntityDescription,
@ -570,20 +574,22 @@ class ShellySleepingBlockAttributeEntity(ShellyBlockAttributeEntity, RestoreEnti
"""Initialize the sleeping sensor."""
self.sensors = sensors
self.last_state: StateType = None
self.wrapper = wrapper
self.coordinator = coordinator
self.attribute = attribute
self.block: Block | None = block # type: ignore[assignment]
self.entity_description = description
self._attr_should_poll = False
self._attr_device_info = DeviceInfo(
connections={(device_registry.CONNECTION_NETWORK_MAC, wrapper.mac)}
connections={(device_registry.CONNECTION_NETWORK_MAC, coordinator.mac)}
)
if block is not None:
self._attr_unique_id = f"{self.wrapper.mac}-{block.description}-{attribute}"
self._attr_unique_id = (
f"{self.coordinator.mac}-{block.description}-{attribute}"
)
self._attr_name = get_block_entity_name(
self.wrapper.device, block, self.entity_description.name
self.coordinator.device, block, self.entity_description.name
)
elif entry is not None:
self._attr_unique_id = entry.unique_id
@ -603,7 +609,7 @@ class ShellySleepingBlockAttributeEntity(ShellyBlockAttributeEntity, RestoreEnti
"""Handle device update."""
if (
self.block is not None
or not self.wrapper.device.initialized
or not self.coordinator.device.initialized
or self.sensors is None
):
super()._update_callback()
@ -611,9 +617,9 @@ class ShellySleepingBlockAttributeEntity(ShellyBlockAttributeEntity, RestoreEnti
_, entity_block, entity_sensor = self._attr_unique_id.split("-")
assert self.wrapper.device.blocks
assert self.coordinator.device.blocks
for block in self.wrapper.device.blocks:
for block in self.coordinator.device.blocks:
if block.description != entity_block:
continue