Store runtime data inside the config entry in Shelly (#116763)

This commit is contained in:
Michael 2024-05-04 12:41:25 +02:00 committed by GitHub
parent e5543e3b95
commit 8238cd9f22
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 93 additions and 108 deletions

View file

@ -16,7 +16,6 @@ from aioshelly.exceptions import (
from aioshelly.rpc_device import RpcDevice from aioshelly.rpc_device import RpcDevice
import voluptuous as vol import voluptuous as vol
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME, Platform from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME, Platform
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
@ -35,7 +34,6 @@ from .const import (
BLOCK_WRONG_SLEEP_PERIOD, BLOCK_WRONG_SLEEP_PERIOD,
CONF_COAP_PORT, CONF_COAP_PORT,
CONF_SLEEP_PERIOD, CONF_SLEEP_PERIOD,
DATA_CONFIG_ENTRY,
DOMAIN, DOMAIN,
FIRMWARE_UNSUPPORTED_ISSUE_ID, FIRMWARE_UNSUPPORTED_ISSUE_ID,
LOGGER, LOGGER,
@ -44,11 +42,11 @@ from .const import (
) )
from .coordinator import ( from .coordinator import (
ShellyBlockCoordinator, ShellyBlockCoordinator,
ShellyConfigEntry,
ShellyEntryData, ShellyEntryData,
ShellyRestCoordinator, ShellyRestCoordinator,
ShellyRpcCoordinator, ShellyRpcCoordinator,
ShellyRpcPollingCoordinator, ShellyRpcPollingCoordinator,
get_entry_data,
) )
from .utils import ( from .utils import (
async_create_issue_unsupported_firmware, async_create_issue_unsupported_firmware,
@ -102,15 +100,13 @@ CONFIG_SCHEMA: Final = vol.Schema({DOMAIN: COAP_SCHEMA}, extra=vol.ALLOW_EXTRA)
async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
"""Set up the Shelly component.""" """Set up the Shelly component."""
hass.data[DOMAIN] = {DATA_CONFIG_ENTRY: {}}
if (conf := config.get(DOMAIN)) is not None: if (conf := config.get(DOMAIN)) is not None:
hass.data[DOMAIN][CONF_COAP_PORT] = conf[CONF_COAP_PORT] hass.data[DOMAIN] = {CONF_COAP_PORT: conf[CONF_COAP_PORT]}
return True return True
async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def async_setup_entry(hass: HomeAssistant, entry: ShellyConfigEntry) -> bool:
"""Set up Shelly from a config entry.""" """Set up Shelly from a config entry."""
# The custom component for Shelly devices uses shelly domain as well as core # The custom component for Shelly devices uses shelly domain as well as core
# integration. If the user removes the custom component but doesn't remove the # integration. If the user removes the custom component but doesn't remove the
@ -127,7 +123,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
) )
return False return False
get_entry_data(hass)[entry.entry_id] = ShellyEntryData() entry.runtime_data = ShellyEntryData()
if get_device_entry_gen(entry) in RPC_GENERATIONS: if get_device_entry_gen(entry) in RPC_GENERATIONS:
return await _async_setup_rpc_entry(hass, entry) return await _async_setup_rpc_entry(hass, entry)
@ -135,7 +131,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
return await _async_setup_block_entry(hass, entry) return await _async_setup_block_entry(hass, entry)
async def _async_setup_block_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def _async_setup_block_entry(
hass: HomeAssistant, entry: ShellyConfigEntry
) -> bool:
"""Set up Shelly block based device from a config entry.""" """Set up Shelly block based device from a config entry."""
options = ConnectionOptions( options = ConnectionOptions(
entry.data[CONF_HOST], entry.data[CONF_HOST],
@ -163,7 +161,7 @@ async def _async_setup_block_entry(hass: HomeAssistant, entry: ConfigEntry) -> b
device_entry = None device_entry = None
sleep_period = entry.data.get(CONF_SLEEP_PERIOD) sleep_period = entry.data.get(CONF_SLEEP_PERIOD)
shelly_entry_data = get_entry_data(hass)[entry.entry_id] shelly_entry_data = entry.runtime_data
# Some old firmware have a wrong sleep period hardcoded value. # Some old firmware have a wrong sleep period hardcoded value.
# Following code block will force the right value for affected devices # Following code block will force the right value for affected devices
@ -220,7 +218,7 @@ async def _async_setup_block_entry(hass: HomeAssistant, entry: ConfigEntry) -> b
return True return True
async def _async_setup_rpc_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def _async_setup_rpc_entry(hass: HomeAssistant, entry: ShellyConfigEntry) -> bool:
"""Set up Shelly RPC based device from a config entry.""" """Set up Shelly RPC based device from a config entry."""
options = ConnectionOptions( options = ConnectionOptions(
entry.data[CONF_HOST], entry.data[CONF_HOST],
@ -249,7 +247,7 @@ async def _async_setup_rpc_entry(hass: HomeAssistant, entry: ConfigEntry) -> boo
device_entry = None device_entry = None
sleep_period = entry.data.get(CONF_SLEEP_PERIOD) sleep_period = entry.data.get(CONF_SLEEP_PERIOD)
shelly_entry_data = get_entry_data(hass)[entry.entry_id] shelly_entry_data = entry.runtime_data
if sleep_period == 0: if sleep_period == 0:
# Not a sleeping device, finish setup # Not a sleeping device, finish setup
@ -290,9 +288,9 @@ async def _async_setup_rpc_entry(hass: HomeAssistant, entry: ConfigEntry) -> boo
return True return True
async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def async_unload_entry(hass: HomeAssistant, entry: ShellyConfigEntry) -> bool:
"""Unload a config entry.""" """Unload a config entry."""
shelly_entry_data = get_entry_data(hass)[entry.entry_id] shelly_entry_data = entry.runtime_data
platforms = RPC_SLEEPING_PLATFORMS platforms = RPC_SLEEPING_PLATFORMS
if not entry.data.get(CONF_SLEEP_PERIOD): if not entry.data.get(CONF_SLEEP_PERIOD):
@ -310,7 +308,6 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
# and if we setup again, we will fix anything that is # and if we setup again, we will fix anything that is
# in an inconsistent state at that time. # in an inconsistent state at that time.
await shelly_entry_data.rpc.shutdown() await shelly_entry_data.rpc.shutdown()
get_entry_data(hass).pop(entry.entry_id)
return unload_ok return unload_ok
@ -331,6 +328,5 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
if unload_ok := await hass.config_entries.async_unload_platforms(entry, platforms): if unload_ok := await hass.config_entries.async_unload_platforms(entry, platforms):
if shelly_entry_data.block: if shelly_entry_data.block:
shelly_entry_data.block.shutdown() shelly_entry_data.block.shutdown()
get_entry_data(hass).pop(entry.entry_id)
return unload_ok return unload_ok

View file

@ -12,13 +12,13 @@ from homeassistant.components.binary_sensor import (
BinarySensorEntity, BinarySensorEntity,
BinarySensorEntityDescription, BinarySensorEntityDescription,
) )
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import STATE_ON, EntityCategory from homeassistant.const import STATE_ON, EntityCategory
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.restore_state import RestoreEntity from homeassistant.helpers.restore_state import RestoreEntity
from .const import CONF_SLEEP_PERIOD from .const import CONF_SLEEP_PERIOD
from .coordinator import ShellyConfigEntry
from .entity import ( from .entity import (
BlockEntityDescription, BlockEntityDescription,
RestEntityDescription, RestEntityDescription,
@ -220,7 +220,7 @@ RPC_SENSORS: Final = {
async def async_setup_entry( async def async_setup_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up sensors for device.""" """Set up sensors for device."""

View file

@ -14,7 +14,6 @@ from homeassistant.components.button import (
ButtonEntity, ButtonEntity,
ButtonEntityDescription, ButtonEntityDescription,
) )
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import EntityCategory from homeassistant.const import EntityCategory
from homeassistant.core import HomeAssistant, callback from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers import entity_registry as er from homeassistant.helpers import entity_registry as er
@ -24,7 +23,7 @@ from homeassistant.helpers.update_coordinator import CoordinatorEntity
from homeassistant.util import slugify from homeassistant.util import slugify
from .const import LOGGER, SHELLY_GAS_MODELS from .const import LOGGER, SHELLY_GAS_MODELS
from .coordinator import ShellyBlockCoordinator, ShellyRpcCoordinator, get_entry_data from .coordinator import ShellyBlockCoordinator, ShellyConfigEntry, ShellyRpcCoordinator
from .utils import get_device_entry_gen from .utils import get_device_entry_gen
_ShellyCoordinatorT = TypeVar( _ShellyCoordinatorT = TypeVar(
@ -108,11 +107,11 @@ def async_migrate_unique_ids(
async def async_setup_entry( async def async_setup_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
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] entry_data = config_entry.runtime_data
coordinator: ShellyRpcCoordinator | ShellyBlockCoordinator | None coordinator: ShellyRpcCoordinator | ShellyBlockCoordinator | None
if get_device_entry_gen(config_entry) in RPC_GENERATIONS: if get_device_entry_gen(config_entry) in RPC_GENERATIONS:
coordinator = entry_data.rpc coordinator = entry_data.rpc

View file

@ -18,7 +18,6 @@ from homeassistant.components.climate import (
HVACAction, HVACAction,
HVACMode, HVACMode,
) )
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import ATTR_TEMPERATURE, UnitOfTemperature from homeassistant.const import ATTR_TEMPERATURE, UnitOfTemperature
from homeassistant.core import HomeAssistant, State, callback from homeassistant.core import HomeAssistant, State, callback
from homeassistant.exceptions import HomeAssistantError from homeassistant.exceptions import HomeAssistantError
@ -42,7 +41,7 @@ from .const import (
RPC_THERMOSTAT_SETTINGS, RPC_THERMOSTAT_SETTINGS,
SHTRV_01_TEMPERATURE_SETTINGS, SHTRV_01_TEMPERATURE_SETTINGS,
) )
from .coordinator import ShellyBlockCoordinator, ShellyRpcCoordinator, get_entry_data from .coordinator import ShellyBlockCoordinator, ShellyConfigEntry, ShellyRpcCoordinator
from .entity import ShellyRpcEntity from .entity import ShellyRpcEntity
from .utils import ( from .utils import (
async_remove_shelly_entity, async_remove_shelly_entity,
@ -54,14 +53,14 @@ from .utils import (
async def async_setup_entry( async def async_setup_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up climate device.""" """Set up climate device."""
if get_device_entry_gen(config_entry) in RPC_GENERATIONS: if get_device_entry_gen(config_entry) in RPC_GENERATIONS:
return async_setup_rpc_entry(hass, config_entry, async_add_entities) return async_setup_rpc_entry(hass, config_entry, async_add_entities)
coordinator = get_entry_data(hass)[config_entry.entry_id].block coordinator = config_entry.runtime_data.block
assert coordinator assert coordinator
if coordinator.device.initialized: if coordinator.device.initialized:
async_setup_climate_entities(async_add_entities, coordinator) async_setup_climate_entities(async_add_entities, coordinator)
@ -99,7 +98,7 @@ def async_setup_climate_entities(
@callback @callback
def async_restore_climate_entities( def async_restore_climate_entities(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
coordinator: ShellyBlockCoordinator, coordinator: ShellyBlockCoordinator,
) -> None: ) -> None:
@ -121,11 +120,11 @@ def async_restore_climate_entities(
@callback @callback
def async_setup_rpc_entry( def async_setup_rpc_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up entities for RPC device.""" """Set up entities for RPC device."""
coordinator = get_entry_data(hass)[config_entry.entry_id].rpc coordinator = config_entry.runtime_data.rpc
assert coordinator assert coordinator
climate_key_ids = get_rpc_key_ids(coordinator.device.status, "thermostat") climate_key_ids = get_rpc_key_ids(coordinator.device.status, "thermostat")

View file

@ -31,7 +31,6 @@ DOMAIN: Final = "shelly"
LOGGER: Logger = getLogger(__package__) LOGGER: Logger = getLogger(__package__)
DATA_CONFIG_ENTRY: Final = "config_entry"
CONF_COAP_PORT: Final = "coap_port" CONF_COAP_PORT: Final = "coap_port"
FIRMWARE_PATTERN: Final = re.compile(r"^(\d{8})") FIRMWARE_PATTERN: Final = re.compile(r"^(\d{8})")

View file

@ -39,7 +39,6 @@ from .const import (
BATTERY_DEVICES_WITH_PERMANENT_CONNECTION, BATTERY_DEVICES_WITH_PERMANENT_CONNECTION,
CONF_BLE_SCANNER_MODE, CONF_BLE_SCANNER_MODE,
CONF_SLEEP_PERIOD, CONF_SLEEP_PERIOD,
DATA_CONFIG_ENTRY,
DOMAIN, DOMAIN,
DUAL_MODE_LIGHT_MODELS, DUAL_MODE_LIGHT_MODELS,
ENTRY_RELOAD_COOLDOWN, ENTRY_RELOAD_COOLDOWN,
@ -85,9 +84,7 @@ class ShellyEntryData:
rpc_poll: ShellyRpcPollingCoordinator | None = None rpc_poll: ShellyRpcPollingCoordinator | None = None
def get_entry_data(hass: HomeAssistant) -> dict[str, ShellyEntryData]: ShellyConfigEntry = ConfigEntry[ShellyEntryData]
"""Return Shelly entry data for a given config entry."""
return cast(dict[str, ShellyEntryData], hass.data[DOMAIN][DATA_CONFIG_ENTRY])
class ShellyCoordinatorBase(DataUpdateCoordinator[None], Generic[_DeviceT]): class ShellyCoordinatorBase(DataUpdateCoordinator[None], Generic[_DeviceT]):
@ -96,7 +93,7 @@ class ShellyCoordinatorBase(DataUpdateCoordinator[None], Generic[_DeviceT]):
def __init__( def __init__(
self, self,
hass: HomeAssistant, hass: HomeAssistant,
entry: ConfigEntry, entry: ShellyConfigEntry,
device: _DeviceT, device: _DeviceT,
update_interval: float, update_interval: float,
) -> None: ) -> None:
@ -217,7 +214,7 @@ class ShellyBlockCoordinator(ShellyCoordinatorBase[BlockDevice]):
"""Coordinator for a Shelly block based device.""" """Coordinator for a Shelly block based device."""
def __init__( def __init__(
self, hass: HomeAssistant, entry: ConfigEntry, device: BlockDevice self, hass: HomeAssistant, entry: ShellyConfigEntry, device: BlockDevice
) -> None: ) -> None:
"""Initialize the Shelly block device coordinator.""" """Initialize the Shelly block device coordinator."""
self.entry = entry self.entry = entry
@ -424,7 +421,7 @@ class ShellyRestCoordinator(ShellyCoordinatorBase[BlockDevice]):
"""Coordinator for a Shelly REST device.""" """Coordinator for a Shelly REST device."""
def __init__( def __init__(
self, hass: HomeAssistant, device: BlockDevice, entry: ConfigEntry self, hass: HomeAssistant, device: BlockDevice, entry: ShellyConfigEntry
) -> None: ) -> None:
"""Initialize the Shelly REST device coordinator.""" """Initialize the Shelly REST device coordinator."""
update_interval = REST_SENSORS_UPDATE_INTERVAL update_interval = REST_SENSORS_UPDATE_INTERVAL
@ -458,7 +455,7 @@ class ShellyRpcCoordinator(ShellyCoordinatorBase[RpcDevice]):
"""Coordinator for a Shelly RPC based device.""" """Coordinator for a Shelly RPC based device."""
def __init__( def __init__(
self, hass: HomeAssistant, entry: ConfigEntry, device: RpcDevice self, hass: HomeAssistant, entry: ShellyConfigEntry, device: RpcDevice
) -> None: ) -> None:
"""Initialize the Shelly RPC device coordinator.""" """Initialize the Shelly RPC device coordinator."""
self.entry = entry self.entry = entry
@ -538,7 +535,7 @@ class ShellyRpcCoordinator(ShellyCoordinatorBase[RpcDevice]):
return _unsubscribe return _unsubscribe
async def _async_update_listener( async def _async_update_listener(
self, hass: HomeAssistant, entry: ConfigEntry self, hass: HomeAssistant, entry: ShellyConfigEntry
) -> None: ) -> None:
"""Reconfigure on update.""" """Reconfigure on update."""
async with self._connection_lock: async with self._connection_lock:
@ -721,7 +718,7 @@ class ShellyRpcPollingCoordinator(ShellyCoordinatorBase[RpcDevice]):
"""Polling coordinator for a Shelly RPC based device.""" """Polling coordinator for a Shelly RPC based device."""
def __init__( def __init__(
self, hass: HomeAssistant, entry: ConfigEntry, device: RpcDevice self, hass: HomeAssistant, entry: ShellyConfigEntry, device: RpcDevice
) -> None: ) -> None:
"""Initialize the RPC polling coordinator.""" """Initialize the RPC polling coordinator."""
super().__init__(hass, entry, device, RPC_SENSORS_POLLING_INTERVAL) super().__init__(hass, entry, device, RPC_SENSORS_POLLING_INTERVAL)
@ -747,10 +744,13 @@ def get_block_coordinator_by_device_id(
dev_reg = dr_async_get(hass) dev_reg = dr_async_get(hass)
if device := dev_reg.async_get(device_id): if device := dev_reg.async_get(device_id):
for config_entry in device.config_entries: for config_entry in device.config_entries:
if not (entry_data := get_entry_data(hass).get(config_entry)): entry = hass.config_entries.async_get_entry(config_entry)
continue if (
entry
if coordinator := entry_data.block: and entry.state == ConfigEntryState.LOADED
and isinstance(entry.runtime_data, ShellyEntryData)
and (coordinator := entry.runtime_data.block)
):
return coordinator return coordinator
return None return None
@ -763,23 +763,25 @@ def get_rpc_coordinator_by_device_id(
dev_reg = dr_async_get(hass) dev_reg = dr_async_get(hass)
if device := dev_reg.async_get(device_id): if device := dev_reg.async_get(device_id):
for config_entry in device.config_entries: for config_entry in device.config_entries:
if not (entry_data := get_entry_data(hass).get(config_entry)): entry = hass.config_entries.async_get_entry(config_entry)
continue if (
entry
if coordinator := entry_data.rpc: and entry.state == ConfigEntryState.LOADED
and isinstance(entry.runtime_data, ShellyEntryData)
and (coordinator := entry.runtime_data.rpc)
):
return coordinator return coordinator
return None return None
async def async_reconnect_soon(hass: HomeAssistant, entry: ConfigEntry) -> None: async def async_reconnect_soon(hass: HomeAssistant, entry: ShellyConfigEntry) -> None:
"""Try to reconnect soon.""" """Try to reconnect soon."""
if ( if (
not entry.data.get(CONF_SLEEP_PERIOD) not entry.data.get(CONF_SLEEP_PERIOD)
and not hass.is_stopping and not hass.is_stopping
and entry.state == ConfigEntryState.LOADED and entry.state == ConfigEntryState.LOADED
and (entry_data := get_entry_data(hass).get(entry.entry_id)) and (coordinator := entry.runtime_data.rpc)
and (coordinator := entry_data.rpc)
): ):
entry.async_create_background_task( entry.async_create_background_task(
hass, hass,

View file

@ -13,18 +13,17 @@ from homeassistant.components.cover import (
CoverEntity, CoverEntity,
CoverEntityFeature, CoverEntityFeature,
) )
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant, callback from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_platform import AddEntitiesCallback
from .coordinator import ShellyBlockCoordinator, ShellyRpcCoordinator, get_entry_data from .coordinator import ShellyBlockCoordinator, ShellyConfigEntry, ShellyRpcCoordinator
from .entity import ShellyBlockEntity, ShellyRpcEntity from .entity import ShellyBlockEntity, ShellyRpcEntity
from .utils import get_device_entry_gen, get_rpc_key_ids from .utils import get_device_entry_gen, get_rpc_key_ids
async def async_setup_entry( async def async_setup_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up covers for device.""" """Set up covers for device."""
@ -37,11 +36,11 @@ async def async_setup_entry(
@callback @callback
def async_setup_block_entry( def async_setup_block_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up cover for device.""" """Set up cover for device."""
coordinator = get_entry_data(hass)[config_entry.entry_id].block coordinator = config_entry.runtime_data.block
assert coordinator and coordinator.device.blocks assert coordinator and coordinator.device.blocks
blocks = [block for block in coordinator.device.blocks if block.type == "roller"] blocks = [block for block in coordinator.device.blocks if block.type == "roller"]
@ -54,11 +53,11 @@ def async_setup_block_entry(
@callback @callback
def async_setup_rpc_entry( def async_setup_rpc_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up entities for RPC device.""" """Set up entities for RPC device."""
coordinator = get_entry_data(hass)[config_entry.entry_id].rpc coordinator = config_entry.runtime_data.rpc
assert coordinator assert coordinator
cover_key_ids = get_rpc_key_ids(coordinator.device.status, "cover") cover_key_ids = get_rpc_key_ids(coordinator.device.status, "cover")

View file

@ -6,21 +6,20 @@ from typing import Any
from homeassistant.components.bluetooth import async_scanner_by_source from homeassistant.components.bluetooth import async_scanner_by_source
from homeassistant.components.diagnostics import async_redact_data from homeassistant.components.diagnostics import async_redact_data
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_PASSWORD, CONF_USERNAME from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from homeassistant.helpers.device_registry import format_mac from homeassistant.helpers.device_registry import format_mac
from .coordinator import get_entry_data from .coordinator import ShellyConfigEntry
TO_REDACT = {CONF_USERNAME, CONF_PASSWORD} TO_REDACT = {CONF_USERNAME, CONF_PASSWORD}
async def async_get_config_entry_diagnostics( async def async_get_config_entry_diagnostics(
hass: HomeAssistant, entry: ConfigEntry hass: HomeAssistant, entry: ShellyConfigEntry
) -> dict[str, Any]: ) -> dict[str, Any]:
"""Return diagnostics for a config entry.""" """Return diagnostics for a config entry."""
shelly_entry_data = get_entry_data(hass)[entry.entry_id] shelly_entry_data = entry.runtime_data
device_settings: str | dict = "not initialized" device_settings: str | dict = "not initialized"
device_status: str | dict = "not initialized" device_status: str | dict = "not initialized"

View file

@ -9,7 +9,6 @@ from typing import Any, cast
from aioshelly.block_device import Block from aioshelly.block_device import Block
from aioshelly.exceptions import DeviceConnectionError, InvalidAuthError, RpcCallError from aioshelly.exceptions import DeviceConnectionError, InvalidAuthError, RpcCallError
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant, State, callback from homeassistant.core import HomeAssistant, State, callback
from homeassistant.exceptions import HomeAssistantError from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC, DeviceInfo from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC, DeviceInfo
@ -24,7 +23,7 @@ from homeassistant.helpers.typing import StateType
from homeassistant.helpers.update_coordinator import CoordinatorEntity from homeassistant.helpers.update_coordinator import CoordinatorEntity
from .const import CONF_SLEEP_PERIOD, LOGGER from .const import CONF_SLEEP_PERIOD, LOGGER
from .coordinator import ShellyBlockCoordinator, ShellyRpcCoordinator, get_entry_data from .coordinator import ShellyBlockCoordinator, ShellyConfigEntry, ShellyRpcCoordinator
from .utils import ( from .utils import (
async_remove_shelly_entity, async_remove_shelly_entity,
get_block_entity_name, get_block_entity_name,
@ -36,13 +35,13 @@ from .utils import (
@callback @callback
def async_setup_entry_attribute_entities( def async_setup_entry_attribute_entities(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
sensors: Mapping[tuple[str, str], BlockEntityDescription], sensors: Mapping[tuple[str, str], BlockEntityDescription],
sensor_class: Callable, sensor_class: Callable,
) -> None: ) -> None:
"""Set up entities for attributes.""" """Set up entities for attributes."""
coordinator = get_entry_data(hass)[config_entry.entry_id].block coordinator = config_entry.runtime_data.block
assert coordinator assert coordinator
if coordinator.device.initialized: if coordinator.device.initialized:
async_setup_block_attribute_entities( async_setup_block_attribute_entities(
@ -104,7 +103,7 @@ def async_setup_block_attribute_entities(
@callback @callback
def async_restore_block_attribute_entities( def async_restore_block_attribute_entities(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
coordinator: ShellyBlockCoordinator, coordinator: ShellyBlockCoordinator,
sensors: Mapping[tuple[str, str], BlockEntityDescription], sensors: Mapping[tuple[str, str], BlockEntityDescription],
@ -139,13 +138,13 @@ def async_restore_block_attribute_entities(
@callback @callback
def async_setup_entry_rpc( def async_setup_entry_rpc(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
sensors: Mapping[str, RpcEntityDescription], sensors: Mapping[str, RpcEntityDescription],
sensor_class: Callable, sensor_class: Callable,
) -> None: ) -> None:
"""Set up entities for RPC sensors.""" """Set up entities for RPC sensors."""
coordinator = get_entry_data(hass)[config_entry.entry_id].rpc coordinator = config_entry.runtime_data.rpc
assert coordinator assert coordinator
if coordinator.device.initialized: if coordinator.device.initialized:
@ -161,18 +160,18 @@ def async_setup_entry_rpc(
@callback @callback
def async_setup_rpc_attribute_entities( def async_setup_rpc_attribute_entities(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
sensors: Mapping[str, RpcEntityDescription], sensors: Mapping[str, RpcEntityDescription],
sensor_class: Callable, sensor_class: Callable,
) -> None: ) -> None:
"""Set up entities for RPC attributes.""" """Set up entities for RPC attributes."""
coordinator = get_entry_data(hass)[config_entry.entry_id].rpc coordinator = config_entry.runtime_data.rpc
assert coordinator assert coordinator
polling_coordinator = None polling_coordinator = None
if not (sleep_period := config_entry.data[CONF_SLEEP_PERIOD]): if not (sleep_period := config_entry.data[CONF_SLEEP_PERIOD]):
polling_coordinator = get_entry_data(hass)[config_entry.entry_id].rpc_poll polling_coordinator = config_entry.runtime_data.rpc_poll
assert polling_coordinator assert polling_coordinator
entities = [] entities = []
@ -213,7 +212,7 @@ def async_setup_rpc_attribute_entities(
@callback @callback
def async_restore_rpc_attribute_entities( def async_restore_rpc_attribute_entities(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
coordinator: ShellyRpcCoordinator, coordinator: ShellyRpcCoordinator,
sensors: Mapping[str, RpcEntityDescription], sensors: Mapping[str, RpcEntityDescription],
@ -248,13 +247,13 @@ def async_restore_rpc_attribute_entities(
@callback @callback
def async_setup_entry_rest( def async_setup_entry_rest(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
sensors: Mapping[str, RestEntityDescription], sensors: Mapping[str, RestEntityDescription],
sensor_class: Callable, sensor_class: Callable,
) -> None: ) -> None:
"""Set up entities for REST sensors.""" """Set up entities for REST sensors."""
coordinator = get_entry_data(hass)[config_entry.entry_id].rest coordinator = config_entry.runtime_data.rest
assert coordinator assert coordinator
async_add_entities( async_add_entities(

View file

@ -15,7 +15,6 @@ from homeassistant.components.event import (
EventEntity, EventEntity,
EventEntityDescription, EventEntityDescription,
) )
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant, callback from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC, DeviceInfo from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC, DeviceInfo
from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_platform import AddEntitiesCallback
@ -26,7 +25,7 @@ from .const import (
RPC_INPUTS_EVENTS_TYPES, RPC_INPUTS_EVENTS_TYPES,
SHIX3_1_INPUTS_EVENTS_TYPES, SHIX3_1_INPUTS_EVENTS_TYPES,
) )
from .coordinator import ShellyBlockCoordinator, ShellyRpcCoordinator, get_entry_data from .coordinator import ShellyBlockCoordinator, ShellyConfigEntry, ShellyRpcCoordinator
from .entity import ShellyBlockEntity from .entity import ShellyBlockEntity
from .utils import ( from .utils import (
async_remove_shelly_entity, async_remove_shelly_entity,
@ -73,7 +72,7 @@ RPC_EVENT: Final = ShellyRpcEventDescription(
async def async_setup_entry( async def async_setup_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up sensors for device.""" """Set up sensors for device."""
@ -82,7 +81,7 @@ async def async_setup_entry(
coordinator: ShellyRpcCoordinator | ShellyBlockCoordinator | None = None coordinator: ShellyRpcCoordinator | ShellyBlockCoordinator | None = None
if get_device_entry_gen(config_entry) in RPC_GENERATIONS: if get_device_entry_gen(config_entry) in RPC_GENERATIONS:
coordinator = get_entry_data(hass)[config_entry.entry_id].rpc coordinator = config_entry.runtime_data.rpc
if TYPE_CHECKING: if TYPE_CHECKING:
assert coordinator assert coordinator
@ -97,7 +96,7 @@ async def async_setup_entry(
else: else:
entities.append(ShellyRpcEvent(coordinator, key, RPC_EVENT)) entities.append(ShellyRpcEvent(coordinator, key, RPC_EVENT))
else: else:
coordinator = get_entry_data(hass)[config_entry.entry_id].block coordinator = config_entry.runtime_data.block
if TYPE_CHECKING: if TYPE_CHECKING:
assert coordinator assert coordinator
assert coordinator.device.blocks assert coordinator.device.blocks

View file

@ -20,7 +20,6 @@ from homeassistant.components.light import (
LightEntityFeature, LightEntityFeature,
brightness_supported, brightness_supported,
) )
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant, callback from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_platform import AddEntitiesCallback
@ -38,7 +37,7 @@ from .const import (
SHELLY_PLUS_RGBW_CHANNELS, SHELLY_PLUS_RGBW_CHANNELS,
STANDARD_RGB_EFFECTS, STANDARD_RGB_EFFECTS,
) )
from .coordinator import ShellyBlockCoordinator, ShellyRpcCoordinator, get_entry_data from .coordinator import ShellyBlockCoordinator, ShellyConfigEntry, ShellyRpcCoordinator
from .entity import ShellyBlockEntity, ShellyRpcEntity from .entity import ShellyBlockEntity, ShellyRpcEntity
from .utils import ( from .utils import (
async_remove_shelly_entity, async_remove_shelly_entity,
@ -54,7 +53,7 @@ from .utils import (
async def async_setup_entry( async def async_setup_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up lights for device.""" """Set up lights for device."""
@ -67,11 +66,11 @@ async def async_setup_entry(
@callback @callback
def async_setup_block_entry( def async_setup_block_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up entities for block device.""" """Set up entities for block device."""
coordinator = get_entry_data(hass)[config_entry.entry_id].block coordinator = config_entry.runtime_data.block
assert coordinator assert coordinator
blocks = [] blocks = []
assert coordinator.device.blocks assert coordinator.device.blocks
@ -97,11 +96,11 @@ def async_setup_block_entry(
@callback @callback
def async_setup_rpc_entry( def async_setup_rpc_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up entities for RPC device.""" """Set up entities for RPC device."""
coordinator = get_entry_data(hass)[config_entry.entry_id].rpc coordinator = config_entry.runtime_data.rpc
assert coordinator assert coordinator
switch_key_ids = get_rpc_key_ids(coordinator.device.status, "switch") switch_key_ids = get_rpc_key_ids(coordinator.device.status, "switch")

View file

@ -14,7 +14,6 @@ from homeassistant.components.number import (
NumberMode, NumberMode,
RestoreNumber, RestoreNumber,
) )
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import PERCENTAGE, EntityCategory from homeassistant.const import PERCENTAGE, EntityCategory
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from homeassistant.exceptions import HomeAssistantError from homeassistant.exceptions import HomeAssistantError
@ -22,7 +21,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.entity_registry import RegistryEntry from homeassistant.helpers.entity_registry import RegistryEntry
from .const import CONF_SLEEP_PERIOD, LOGGER from .const import CONF_SLEEP_PERIOD, LOGGER
from .coordinator import ShellyBlockCoordinator from .coordinator import ShellyBlockCoordinator, ShellyConfigEntry
from .entity import ( from .entity import (
BlockEntityDescription, BlockEntityDescription,
ShellySleepingBlockAttributeEntity, ShellySleepingBlockAttributeEntity,
@ -58,7 +57,7 @@ NUMBERS: dict[tuple[str, str], BlockNumberDescription] = {
async def async_setup_entry( async def async_setup_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up numbers for device.""" """Set up numbers for device."""

View file

@ -16,7 +16,6 @@ from homeassistant.components.sensor import (
SensorExtraStoredData, SensorExtraStoredData,
SensorStateClass, SensorStateClass,
) )
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import ( from homeassistant.const import (
CONCENTRATION_PARTS_PER_MILLION, CONCENTRATION_PARTS_PER_MILLION,
DEGREE, DEGREE,
@ -38,7 +37,7 @@ from homeassistant.helpers.entity_registry import RegistryEntry
from homeassistant.helpers.typing import StateType from homeassistant.helpers.typing import StateType
from .const import CONF_SLEEP_PERIOD, SHAIR_MAX_WORK_HOURS from .const import CONF_SLEEP_PERIOD, SHAIR_MAX_WORK_HOURS
from .coordinator import ShellyBlockCoordinator, ShellyRpcCoordinator from .coordinator import ShellyBlockCoordinator, ShellyConfigEntry, ShellyRpcCoordinator
from .entity import ( from .entity import (
BlockEntityDescription, BlockEntityDescription,
RestEntityDescription, RestEntityDescription,
@ -995,7 +994,7 @@ RPC_SENSORS: Final = {
async def async_setup_entry( async def async_setup_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up sensors for device.""" """Set up sensors for device."""

View file

@ -22,13 +22,12 @@ from homeassistant.components.switch import (
SwitchEntityDescription, SwitchEntityDescription,
) )
from homeassistant.components.valve import DOMAIN as VALVE_DOMAIN from homeassistant.components.valve import DOMAIN as VALVE_DOMAIN
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant, callback from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.issue_registry import IssueSeverity, async_create_issue from homeassistant.helpers.issue_registry import IssueSeverity, async_create_issue
from .const import DOMAIN, GAS_VALVE_OPEN_STATES from .const import DOMAIN, GAS_VALVE_OPEN_STATES
from .coordinator import ShellyBlockCoordinator, ShellyRpcCoordinator, get_entry_data from .coordinator import ShellyBlockCoordinator, ShellyConfigEntry, ShellyRpcCoordinator
from .entity import ( from .entity import (
BlockEntityDescription, BlockEntityDescription,
ShellyBlockAttributeEntity, ShellyBlockAttributeEntity,
@ -64,7 +63,7 @@ GAS_VALVE_SWITCH = BlockSwitchDescription(
async def async_setup_entry( async def async_setup_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up switches for device.""" """Set up switches for device."""
@ -77,11 +76,11 @@ async def async_setup_entry(
@callback @callback
def async_setup_block_entry( def async_setup_block_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up entities for block device.""" """Set up entities for block device."""
coordinator = get_entry_data(hass)[config_entry.entry_id].block coordinator = config_entry.runtime_data.block
assert coordinator assert coordinator
# Add Shelly Gas Valve as a switch # Add Shelly Gas Valve as a switch
@ -127,11 +126,11 @@ def async_setup_block_entry(
@callback @callback
def async_setup_rpc_entry( def async_setup_rpc_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up entities for RPC device.""" """Set up entities for RPC device."""
coordinator = get_entry_data(hass)[config_entry.entry_id].rpc coordinator = config_entry.runtime_data.rpc
assert coordinator assert coordinator
switch_key_ids = get_rpc_key_ids(coordinator.device.status, "switch") switch_key_ids = get_rpc_key_ids(coordinator.device.status, "switch")

View file

@ -18,7 +18,6 @@ from homeassistant.components.update import (
UpdateEntityDescription, UpdateEntityDescription,
UpdateEntityFeature, UpdateEntityFeature,
) )
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import EntityCategory from homeassistant.const import EntityCategory
from homeassistant.core import HomeAssistant, callback from homeassistant.core import HomeAssistant, callback
from homeassistant.exceptions import HomeAssistantError from homeassistant.exceptions import HomeAssistantError
@ -26,7 +25,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.restore_state import RestoreEntity from homeassistant.helpers.restore_state import RestoreEntity
from .const import CONF_SLEEP_PERIOD, OTA_BEGIN, OTA_ERROR, OTA_PROGRESS, OTA_SUCCESS from .const import CONF_SLEEP_PERIOD, OTA_BEGIN, OTA_ERROR, OTA_PROGRESS, OTA_SUCCESS
from .coordinator import ShellyBlockCoordinator, ShellyRpcCoordinator from .coordinator import ShellyBlockCoordinator, ShellyConfigEntry, ShellyRpcCoordinator
from .entity import ( from .entity import (
RestEntityDescription, RestEntityDescription,
RpcEntityDescription, RpcEntityDescription,
@ -103,7 +102,7 @@ RPC_UPDATES: Final = {
async def async_setup_entry( async def async_setup_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up update entities for Shelly component.""" """Set up update entities for Shelly component."""

View file

@ -14,11 +14,10 @@ from homeassistant.components.valve import (
ValveEntityDescription, ValveEntityDescription,
ValveEntityFeature, ValveEntityFeature,
) )
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant, callback from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_platform import AddEntitiesCallback
from .coordinator import ShellyBlockCoordinator, get_entry_data from .coordinator import ShellyBlockCoordinator, ShellyConfigEntry
from .entity import ( from .entity import (
BlockEntityDescription, BlockEntityDescription,
ShellyBlockAttributeEntity, ShellyBlockAttributeEntity,
@ -42,7 +41,7 @@ GAS_VALVE = BlockValveDescription(
async def async_setup_entry( async def async_setup_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up valves for device.""" """Set up valves for device."""
@ -53,11 +52,11 @@ async def async_setup_entry(
@callback @callback
def async_setup_block_entry( def async_setup_block_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up valve for device.""" """Set up valve for device."""
coordinator = get_entry_data(hass)[config_entry.entry_id].block coordinator = config_entry.runtime_data.block
assert coordinator and coordinator.device.blocks assert coordinator and coordinator.device.blocks
if coordinator.model == MODEL_GAS: if coordinator.model == MODEL_GAS: