Use HassKey in notify (#126338)

This commit is contained in:
epenet 2024-09-21 13:17:01 +02:00 committed by GitHub
parent 37d527bd08
commit aa736b2de6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 21 additions and 20 deletions

View File

@ -20,6 +20,7 @@ from homeassistant.helpers.entity_component import EntityComponent
from homeassistant.helpers.restore_state import RestoreEntity
from homeassistant.helpers.typing import ConfigType
from homeassistant.util import dt as dt_util
from homeassistant.util.hass_dict import HassKey
from .const import ( # noqa: F401
ATTR_DATA,
@ -46,6 +47,7 @@ from .repairs import migrate_notify_issue # noqa: F401
# Platform specific data
ATTR_TITLE_DEFAULT = "Home Assistant"
DOMAIN_DATA: HassKey[EntityComponent[NotifyEntity]] = HassKey(DOMAIN)
ENTITY_ID_FORMAT = DOMAIN + ".{}"
MIN_TIME_BETWEEN_SCANS = timedelta(seconds=10)
@ -76,7 +78,9 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
# legacy platforms to finish setting up.
hass.async_create_task(setup, eager_start=True)
component = hass.data[DOMAIN] = EntityComponent[NotifyEntity](_LOGGER, DOMAIN, hass)
component = hass.data[DOMAIN_DATA] = EntityComponent[NotifyEntity](
_LOGGER, DOMAIN, hass
)
component.async_register_entity_service(
SERVICE_SEND_MESSAGE,
{
@ -113,14 +117,12 @@ class NotifyEntityDescription(EntityDescription, frozen_or_thawed=True):
async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Set up a config entry."""
component: EntityComponent[NotifyEntity] = hass.data[DOMAIN]
return await component.async_setup_entry(entry)
return await hass.data[DOMAIN_DATA].async_setup_entry(entry)
async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Unload a config entry."""
component: EntityComponent[NotifyEntity] = hass.data[DOMAIN]
return await component.async_unload_entry(entry)
return await hass.data[DOMAIN_DATA].async_unload_entry(entry)
class NotifyEntity(RestoreEntity):

View File

@ -3,13 +3,13 @@
from __future__ import annotations
import asyncio
from collections.abc import Callable, Coroutine, Mapping
from collections.abc import Coroutine, Mapping
from functools import partial
from typing import Any, Protocol, cast
from homeassistant.config import config_per_platform
from homeassistant.const import CONF_DESCRIPTION, CONF_NAME
from homeassistant.core import HomeAssistant, ServiceCall, callback
from homeassistant.core import CALLBACK_TYPE, HomeAssistant, ServiceCall, callback
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers import discovery
from homeassistant.helpers.service import async_set_service_schema
@ -21,6 +21,7 @@ from homeassistant.setup import (
async_start_setup,
)
from homeassistant.util import slugify
from homeassistant.util.hass_dict import HassKey
from homeassistant.util.yaml import load_yaml_dict
from .const import (
@ -35,8 +36,12 @@ from .const import (
)
CONF_FIELDS = "fields"
NOTIFY_SERVICES = "notify_services"
NOTIFY_DISCOVERY_DISPATCHER = "notify_discovery_dispatcher"
NOTIFY_SERVICES: HassKey[dict[str, list[BaseNotificationService]]] = HassKey(
f"{DOMAIN}_services"
)
NOTIFY_DISCOVERY_DISPATCHER: HassKey[CALLBACK_TYPE | None] = HassKey(
f"{DOMAIN}_discovery_dispatcher"
)
class LegacyNotifyPlatform(Protocol):
@ -160,11 +165,9 @@ async def async_reload(hass: HomeAssistant, integration_name: str) -> None:
if not _async_integration_has_notify_services(hass, integration_name):
return
notify_services: list[BaseNotificationService] = hass.data[NOTIFY_SERVICES][
integration_name
]
tasks = [
notify_service.async_register_services() for notify_service in notify_services
notify_service.async_register_services()
for notify_service in hass.data[NOTIFY_SERVICES][integration_name]
]
await asyncio.gather(*tasks)
@ -173,20 +176,16 @@ async def async_reload(hass: HomeAssistant, integration_name: str) -> None:
@bind_hass
async def async_reset_platform(hass: HomeAssistant, integration_name: str) -> None:
"""Unregister notify services for an integration."""
notify_discovery_dispatcher: Callable[[], None] | None = hass.data.get(
NOTIFY_DISCOVERY_DISPATCHER
)
notify_discovery_dispatcher = hass.data.get(NOTIFY_DISCOVERY_DISPATCHER)
if notify_discovery_dispatcher:
notify_discovery_dispatcher()
hass.data[NOTIFY_DISCOVERY_DISPATCHER] = None
if not _async_integration_has_notify_services(hass, integration_name):
return
notify_services: list[BaseNotificationService] = hass.data[NOTIFY_SERVICES][
integration_name
]
tasks = [
notify_service.async_unregister_services() for notify_service in notify_services
notify_service.async_unregister_services()
for notify_service in hass.data[NOTIFY_SERVICES][integration_name]
]
await asyncio.gather(*tasks)