From e389ef9920e4fd2d157b22f6133acae2184d4daa Mon Sep 17 00:00:00 2001 From: Jan-Philipp Benecke Date: Wed, 20 Nov 2024 20:20:12 +0100 Subject: [PATCH] Use runtime_data in sabnzbd (#131069) --- homeassistant/components/sabnzbd/__init__.py | 20 +++++++++++--------- homeassistant/components/sabnzbd/button.py | 8 ++++---- homeassistant/components/sabnzbd/number.py | 6 +++--- homeassistant/components/sabnzbd/sensor.py | 10 +++------- 4 files changed, 21 insertions(+), 23 deletions(-) diff --git a/homeassistant/components/sabnzbd/__init__.py b/homeassistant/components/sabnzbd/__init__.py index abc9b1b1356..74920ba6465 100644 --- a/homeassistant/components/sabnzbd/__init__.py +++ b/homeassistant/components/sabnzbd/__init__.py @@ -48,20 +48,24 @@ SERVICE_SPEED_SCHEMA = SERVICE_BASE_SCHEMA.extend( } ) +type SabnzbdConfigEntry = ConfigEntry[SabnzbdUpdateCoordinator] + @callback -def async_get_entry_id_for_service_call(hass: HomeAssistant, call: ServiceCall) -> str: +def async_get_entry_for_service_call( + hass: HomeAssistant, call: ServiceCall +) -> SabnzbdConfigEntry: """Get the entry ID related to a service call (by device ID).""" call_data_api_key = call.data[ATTR_API_KEY] for entry in hass.config_entries.async_entries(DOMAIN): if entry.data[ATTR_API_KEY] == call_data_api_key: - return entry.entry_id + return entry raise ValueError(f"No api for API key: {call_data_api_key}") -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, entry: SabnzbdConfigEntry) -> bool: """Set up the SabNzbd Component.""" sab_api = await get_client(hass, entry.data) @@ -70,7 +74,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: coordinator = SabnzbdUpdateCoordinator(hass, entry, sab_api) await coordinator.async_config_entry_first_refresh() - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator + entry.runtime_data = coordinator @callback def extract_api( @@ -82,8 +86,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def wrapper(call: ServiceCall) -> None: """Wrap the service function.""" - entry_id = async_get_entry_id_for_service_call(hass, call) - coordinator: SabnzbdUpdateCoordinator = hass.data[DOMAIN][entry_id] + config_entry = async_get_entry_for_service_call(hass, call) + coordinator = config_entry.runtime_data try: await func(call, coordinator) @@ -155,11 +159,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: return True -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_unload_entry(hass: HomeAssistant, entry: SabnzbdConfigEntry) -> bool: """Unload a Sabnzbd config entry.""" unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS) - if unload_ok: - hass.data[DOMAIN].pop(entry.entry_id) loaded_entries = [ entry diff --git a/homeassistant/components/sabnzbd/button.py b/homeassistant/components/sabnzbd/button.py index 4efecba18a7..79038e84775 100644 --- a/homeassistant/components/sabnzbd/button.py +++ b/homeassistant/components/sabnzbd/button.py @@ -7,13 +7,13 @@ from typing import Any from pysabnzbd import SabnzbdApiException from homeassistant.components.button import ButtonEntity, ButtonEntityDescription -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import SabnzbdUpdateCoordinator +from . import SabnzbdConfigEntry from .const import DOMAIN +from .coordinator import SabnzbdUpdateCoordinator from .entity import SabnzbdEntity @@ -40,11 +40,11 @@ BUTTON_DESCRIPTIONS: tuple[SabnzbdButtonEntityDescription, ...] = ( async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: SabnzbdConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up buttons from a config entry.""" - coordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data async_add_entities( SabnzbdButton(coordinator, description) for description in BUTTON_DESCRIPTIONS diff --git a/homeassistant/components/sabnzbd/number.py b/homeassistant/components/sabnzbd/number.py index 31faca3f78b..d8536cb6b37 100644 --- a/homeassistant/components/sabnzbd/number.py +++ b/homeassistant/components/sabnzbd/number.py @@ -12,12 +12,12 @@ from homeassistant.components.number import ( NumberEntityDescription, NumberMode, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import PERCENTAGE from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import AddEntitiesCallback +from . import SabnzbdConfigEntry from .const import DOMAIN from .coordinator import SabnzbdUpdateCoordinator from .entity import SabnzbdEntity @@ -48,11 +48,11 @@ NUMBER_DESCRIPTIONS: tuple[SabnzbdNumberEntityDescription, ...] = ( async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: SabnzbdConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up the SABnzbd number entity.""" - coordinator = hass.data[DOMAIN][config_entry.entry_id] + coordinator = config_entry.runtime_data async_add_entities( SabnzbdNumber(coordinator, description) for description in NUMBER_DESCRIPTIONS diff --git a/homeassistant/components/sabnzbd/sensor.py b/homeassistant/components/sabnzbd/sensor.py index 1c6a1279263..115b9de3793 100644 --- a/homeassistant/components/sabnzbd/sensor.py +++ b/homeassistant/components/sabnzbd/sensor.py @@ -10,14 +10,12 @@ from homeassistant.components.sensor import ( SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import UnitOfDataRate, UnitOfInformation from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.typing import StateType -from .const import DOMAIN -from .coordinator import SabnzbdUpdateCoordinator +from . import SabnzbdConfigEntry from .entity import SabnzbdEntity @@ -116,13 +114,11 @@ SENSOR_TYPES: tuple[SabnzbdSensorEntityDescription, ...] = ( async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: SabnzbdConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up a Sabnzbd sensor entry.""" - - entry_id = config_entry.entry_id - coordinator: SabnzbdUpdateCoordinator = hass.data[DOMAIN][entry_id] + coordinator = config_entry.runtime_data async_add_entities([SabnzbdSensor(coordinator, sensor) for sensor in SENSOR_TYPES])