diff --git a/homeassistant/components/pushbullet/__init__.py b/homeassistant/components/pushbullet/__init__.py index e5892afc926..4adfbcad4f9 100644 --- a/homeassistant/components/pushbullet/__init__.py +++ b/homeassistant/components/pushbullet/__init__.py @@ -21,6 +21,8 @@ from homeassistant.helpers.typing import ConfigType from .api import PushBulletNotificationProvider from .const import DATA_HASS_CONFIG, DOMAIN +type PushbulletConfigEntry = ConfigEntry[PushBulletNotificationProvider] + PLATFORMS = [Platform.SENSOR] _LOGGER = logging.getLogger(__name__) @@ -35,7 +37,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: return True -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, entry: PushbulletConfigEntry) -> bool: """Set up pushbullet from a config entry.""" try: @@ -49,7 +51,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: raise ConfigEntryNotReady from err pb_provider = PushBulletNotificationProvider(hass, pushbullet) - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = pb_provider + entry.runtime_data = pb_provider def start_listener(event: Event) -> None: """Start the listener thread.""" @@ -72,11 +74,8 @@ 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: PushbulletConfigEntry) -> bool: """Unload a config entry.""" if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS): - pb_provider: PushBulletNotificationProvider = hass.data[DOMAIN].pop( - entry.entry_id - ) - await hass.async_add_executor_job(pb_provider.close) + await hass.async_add_executor_job(entry.runtime_data.close) return unload_ok diff --git a/homeassistant/components/pushbullet/notify.py b/homeassistant/components/pushbullet/notify.py index f2e70695b27..26ecc859ad2 100644 --- a/homeassistant/components/pushbullet/notify.py +++ b/homeassistant/components/pushbullet/notify.py @@ -22,8 +22,7 @@ from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType -from .api import PushBulletNotificationProvider -from .const import ATTR_FILE, ATTR_FILE_URL, ATTR_URL, DOMAIN +from .const import ATTR_FILE, ATTR_FILE_URL, ATTR_URL _LOGGER = logging.getLogger(__name__) @@ -36,10 +35,10 @@ async def async_get_service( """Get the Pushbullet notification service.""" if TYPE_CHECKING: assert discovery_info is not None - pb_provider: PushBulletNotificationProvider = hass.data[DOMAIN][ - discovery_info["entry_id"] - ] - return PushBulletNotificationService(hass, pb_provider.pushbullet) + entry = hass.config_entries.async_get_entry(discovery_info["entry_id"]) + if TYPE_CHECKING: + assert entry is not None + return PushBulletNotificationService(hass, entry.runtime_data.pushbullet) class PushBulletNotificationService(BaseNotificationService): diff --git a/homeassistant/components/pushbullet/sensor.py b/homeassistant/components/pushbullet/sensor.py index 3ab55ecf072..ade6f9362ed 100644 --- a/homeassistant/components/pushbullet/sensor.py +++ b/homeassistant/components/pushbullet/sensor.py @@ -3,13 +3,13 @@ from __future__ import annotations from homeassistant.components.sensor import SensorEntity, SensorEntityDescription -from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_NAME, MAX_LENGTH_STATE_STATE from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback +from . import PushbulletConfigEntry from .api import PushBulletNotificationProvider from .const import DATA_UPDATED, DOMAIN @@ -69,12 +69,12 @@ SENSOR_KEYS: list[str] = [desc.key for desc in SENSOR_TYPES] async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: PushbulletConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up the Pushbullet sensors from config entry.""" - pb_provider: PushBulletNotificationProvider = hass.data[DOMAIN][entry.entry_id] + pb_provider = entry.runtime_data entities = [ PushBulletNotificationSensor(entry.data[CONF_NAME], pb_provider, description)