314 lines
10 KiB
Python
314 lines
10 KiB
Python
"""Support for esphome entities."""
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable
|
|
import functools
|
|
import math
|
|
from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast
|
|
|
|
from aioesphomeapi import (
|
|
EntityCategory as EsphomeEntityCategory,
|
|
EntityInfo,
|
|
EntityState,
|
|
build_unique_id,
|
|
)
|
|
import voluptuous as vol
|
|
|
|
from homeassistant.config_entries import ConfigEntry
|
|
from homeassistant.const import EntityCategory
|
|
from homeassistant.core import HomeAssistant, callback
|
|
from homeassistant.helpers import entity_platform
|
|
import homeassistant.helpers.config_validation as cv
|
|
import homeassistant.helpers.device_registry as dr
|
|
from homeassistant.helpers.device_registry import DeviceInfo
|
|
from homeassistant.helpers.dispatcher import async_dispatcher_connect
|
|
from homeassistant.helpers.entity import Entity
|
|
from homeassistant.helpers.entity_platform import AddEntitiesCallback
|
|
|
|
from .domain_data import DomainData
|
|
|
|
# Import config flow so that it's added to the registry
|
|
from .entry_data import RuntimeEntryData
|
|
from .enum_mapper import EsphomeEnumMapper
|
|
|
|
_R = TypeVar("_R")
|
|
_InfoT = TypeVar("_InfoT", bound=EntityInfo)
|
|
_EntityT = TypeVar("_EntityT", bound="EsphomeEntity[Any,Any]")
|
|
_StateT = TypeVar("_StateT", bound=EntityState)
|
|
|
|
|
|
async def platform_async_setup_entry(
|
|
hass: HomeAssistant,
|
|
entry: ConfigEntry,
|
|
async_add_entities: AddEntitiesCallback,
|
|
*,
|
|
info_type: type[_InfoT],
|
|
entity_type: type[_EntityT],
|
|
state_type: type[_StateT],
|
|
) -> None:
|
|
"""Set up an esphome platform.
|
|
|
|
This method is in charge of receiving, distributing and storing
|
|
info and state updates.
|
|
"""
|
|
entry_data: RuntimeEntryData = DomainData.get(hass).get_entry_data(entry)
|
|
entry_data.info[info_type] = {}
|
|
entry_data.state.setdefault(state_type, {})
|
|
platform = entity_platform.async_get_current_platform()
|
|
|
|
@callback
|
|
def async_list_entities(infos: list[EntityInfo]) -> None:
|
|
"""Update entities of this platform when entities are listed."""
|
|
current_infos = entry_data.info[info_type]
|
|
new_infos: dict[int, EntityInfo] = {}
|
|
add_entities: list[_EntityT] = []
|
|
|
|
for info in infos:
|
|
if not current_infos.pop(info.key, None):
|
|
# Create new entity
|
|
entity = entity_type(entry_data, platform.domain, info, state_type)
|
|
add_entities.append(entity)
|
|
new_infos[info.key] = info
|
|
|
|
# Anything still in current_infos is now gone
|
|
if current_infos:
|
|
hass.async_create_task(
|
|
entry_data.async_remove_entities(current_infos.values())
|
|
)
|
|
|
|
# Then update the actual info
|
|
entry_data.info[info_type] = new_infos
|
|
|
|
if new_infos:
|
|
entry_data.async_update_entity_infos(new_infos.values())
|
|
|
|
if add_entities:
|
|
# Add entities to Home Assistant
|
|
async_add_entities(add_entities)
|
|
|
|
entry_data.cleanup_callbacks.append(
|
|
entry_data.async_register_static_info_callback(info_type, async_list_entities)
|
|
)
|
|
|
|
|
|
def esphome_state_property(
|
|
func: Callable[[_EntityT], _R],
|
|
) -> Callable[[_EntityT], _R | None]:
|
|
"""Wrap a state property of an esphome entity.
|
|
|
|
This checks if the state object in the entity is set, and
|
|
prevents writing NAN values to the Home Assistant state machine.
|
|
"""
|
|
|
|
@functools.wraps(func)
|
|
def _wrapper(self: _EntityT) -> _R | None:
|
|
# pylint: disable-next=protected-access
|
|
if not self._has_state:
|
|
return None
|
|
val = func(self)
|
|
if isinstance(val, float) and not math.isfinite(val):
|
|
# Home Assistant doesn't use NaN or inf values in state machine
|
|
# (not JSON serializable)
|
|
return None
|
|
return val
|
|
|
|
return _wrapper
|
|
|
|
|
|
ICON_SCHEMA = vol.Schema(cv.icon)
|
|
|
|
|
|
ENTITY_CATEGORIES: EsphomeEnumMapper[
|
|
EsphomeEntityCategory, EntityCategory | None
|
|
] = EsphomeEnumMapper(
|
|
{
|
|
EsphomeEntityCategory.NONE: None,
|
|
EsphomeEntityCategory.CONFIG: EntityCategory.CONFIG,
|
|
EsphomeEntityCategory.DIAGNOSTIC: EntityCategory.DIAGNOSTIC,
|
|
}
|
|
)
|
|
|
|
|
|
class EsphomeEntity(Entity, Generic[_InfoT, _StateT]):
|
|
"""Define a base esphome entity."""
|
|
|
|
_attr_should_poll = False
|
|
_static_info: _InfoT
|
|
_state: _StateT
|
|
_has_state: bool
|
|
|
|
def __init__(
|
|
self,
|
|
entry_data: RuntimeEntryData,
|
|
domain: str,
|
|
entity_info: EntityInfo,
|
|
state_type: type[_StateT],
|
|
) -> None:
|
|
"""Initialize."""
|
|
|
|
self._entry_data = entry_data
|
|
self._on_entry_data_changed()
|
|
self._key = entity_info.key
|
|
self._state_type = state_type
|
|
self._on_static_info_update(entity_info)
|
|
assert entry_data.device_info is not None
|
|
device_info = entry_data.device_info
|
|
self._device_info = device_info
|
|
self._attr_device_info = DeviceInfo(
|
|
connections={(dr.CONNECTION_NETWORK_MAC, device_info.mac_address)}
|
|
)
|
|
self._entry_id = entry_data.entry_id
|
|
#
|
|
# If `friendly_name` is set, we use the Friendly naming rules, if
|
|
# `friendly_name` is not set we make an exception to the naming rules for
|
|
# backwards compatibility and use the Legacy naming rules.
|
|
#
|
|
# Friendly naming
|
|
# - Friendly name is prepended to entity names
|
|
# - Device Name is prepended to entity ids
|
|
# - Entity id is constructed from device name and object id
|
|
#
|
|
# Legacy naming
|
|
# - Device name is not prepended to entity names
|
|
# - Device name is not prepended to entity ids
|
|
# - Entity id is constructed from entity name
|
|
#
|
|
if not device_info.friendly_name:
|
|
return
|
|
self._attr_has_entity_name = True
|
|
self.entity_id = f"{domain}.{device_info.name}_{entity_info.object_id}"
|
|
|
|
async def async_added_to_hass(self) -> None:
|
|
"""Register callbacks."""
|
|
entry_data = self._entry_data
|
|
hass = self.hass
|
|
key = self._key
|
|
|
|
self.async_on_remove(
|
|
entry_data.async_register_key_static_info_remove_callback(
|
|
self._static_info,
|
|
functools.partial(self.async_remove, force_remove=True),
|
|
)
|
|
)
|
|
self.async_on_remove(
|
|
async_dispatcher_connect(
|
|
hass,
|
|
entry_data.signal_device_updated,
|
|
self._on_device_update,
|
|
)
|
|
)
|
|
self.async_on_remove(
|
|
entry_data.async_subscribe_state_update(
|
|
self._state_type, key, self._on_state_update
|
|
)
|
|
)
|
|
self.async_on_remove(
|
|
entry_data.async_register_key_static_info_updated_callback(
|
|
self._static_info, self._on_static_info_update
|
|
)
|
|
)
|
|
self._update_state_from_entry_data()
|
|
|
|
@callback
|
|
def _on_static_info_update(self, static_info: EntityInfo) -> None:
|
|
"""Save the static info for this entity when it changes.
|
|
|
|
This method can be overridden in child classes to know
|
|
when the static info changes.
|
|
"""
|
|
device_info = self._entry_data.device_info
|
|
if TYPE_CHECKING:
|
|
static_info = cast(_InfoT, static_info)
|
|
assert device_info
|
|
self._static_info = static_info
|
|
self._attr_unique_id = build_unique_id(device_info.mac_address, static_info)
|
|
self._attr_entity_registry_enabled_default = not static_info.disabled_by_default
|
|
self._attr_name = static_info.name
|
|
if entity_category := static_info.entity_category:
|
|
self._attr_entity_category = ENTITY_CATEGORIES.from_esphome(entity_category)
|
|
else:
|
|
self._attr_entity_category = None
|
|
if icon := static_info.icon:
|
|
self._attr_icon = cast(str, ICON_SCHEMA(icon))
|
|
else:
|
|
self._attr_icon = None
|
|
|
|
@callback
|
|
def _update_state_from_entry_data(self) -> None:
|
|
"""Update state from entry data."""
|
|
|
|
state = self._entry_data.state
|
|
key = self._key
|
|
state_type = self._state_type
|
|
has_state = key in state[state_type]
|
|
if has_state:
|
|
self._state = cast(_StateT, state[state_type][key])
|
|
self._has_state = has_state
|
|
|
|
@callback
|
|
def _on_state_update(self) -> None:
|
|
"""Call when state changed.
|
|
|
|
Behavior can be changed in child classes
|
|
"""
|
|
self._update_state_from_entry_data()
|
|
self.async_write_ha_state()
|
|
|
|
@callback
|
|
def _on_entry_data_changed(self) -> None:
|
|
entry_data = self._entry_data
|
|
self._api_version = entry_data.api_version
|
|
self._client = entry_data.client
|
|
|
|
@callback
|
|
def _on_device_update(self) -> None:
|
|
"""Call when device updates or entry data changes."""
|
|
self._on_entry_data_changed()
|
|
if not self._entry_data.available:
|
|
# Only write state if the device has gone unavailable
|
|
# since _on_state_update will be called if the device
|
|
# is available when the full state arrives
|
|
# through the next entity state packet.
|
|
self.async_write_ha_state()
|
|
|
|
@property
|
|
def available(self) -> bool:
|
|
"""Return if the entity is available."""
|
|
if self._device_info.has_deep_sleep:
|
|
# During deep sleep the ESP will not be connectable (by design)
|
|
# For these cases, show it as available
|
|
return self._entry_data.expected_disconnect
|
|
|
|
return self._entry_data.available
|
|
|
|
|
|
class EsphomeAssistEntity(Entity):
|
|
"""Define a base entity for Assist Pipeline entities."""
|
|
|
|
_attr_has_entity_name = True
|
|
_attr_should_poll = False
|
|
|
|
def __init__(self, entry_data: RuntimeEntryData) -> None:
|
|
"""Initialize the binary sensor."""
|
|
self._entry_data: RuntimeEntryData = entry_data
|
|
assert entry_data.device_info is not None
|
|
device_info = entry_data.device_info
|
|
self._device_info = device_info
|
|
self._attr_unique_id = (
|
|
f"{device_info.mac_address}-{self.entity_description.key}"
|
|
)
|
|
self._attr_device_info = DeviceInfo(
|
|
connections={(dr.CONNECTION_NETWORK_MAC, device_info.mac_address)}
|
|
)
|
|
|
|
@callback
|
|
def _update(self) -> None:
|
|
self.async_write_ha_state()
|
|
|
|
async def async_added_to_hass(self) -> None:
|
|
"""Register update callback."""
|
|
await super().async_added_to_hass()
|
|
self.async_on_remove(
|
|
self._entry_data.async_subscribe_assist_pipeline_update(self._update)
|
|
)
|