269 lines
10 KiB
Python
269 lines
10 KiB
Python
"""Node-RED base entity class."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextlib import suppress
|
|
from typing import TYPE_CHECKING, Any, ClassVar
|
|
|
|
from homeassistant.const import (
|
|
CONF_DEVICE_CLASS,
|
|
CONF_ENTITY_CATEGORY,
|
|
CONF_ICON,
|
|
CONF_UNIT_OF_MEASUREMENT,
|
|
EntityCategory,
|
|
)
|
|
from homeassistant.core import HomeAssistant, callback
|
|
from homeassistant.helpers import device_registry as dr
|
|
from homeassistant.helpers.device_registry import DeviceInfo
|
|
from homeassistant.helpers.dispatcher import (
|
|
async_dispatcher_connect,
|
|
async_dispatcher_send,
|
|
)
|
|
from homeassistant.helpers.entity import Entity
|
|
from homeassistant.helpers.entity_registry import async_get
|
|
|
|
if TYPE_CHECKING:
|
|
from homeassistant.components.websocket_api.connection import ActiveConnection
|
|
|
|
from .const import (
|
|
CONF_ATTRIBUTES,
|
|
CONF_COMPONENT,
|
|
CONF_CONFIG,
|
|
CONF_DEVICE_INFO,
|
|
CONF_ENTITY_PICTURE,
|
|
CONF_NAME,
|
|
CONF_NODE_ID,
|
|
CONF_OPTIONS,
|
|
CONF_REMOVE,
|
|
CONF_SERVER_ID,
|
|
DOMAIN,
|
|
DOMAIN_DATA,
|
|
NODERED_CONFIG_UPDATE,
|
|
NODERED_DISCOVERY,
|
|
NODERED_DISCOVERY_UPDATED,
|
|
NODERED_ENTITY,
|
|
)
|
|
from .discovery import ALREADY_DISCOVERED, CHANGE_ENTITY_TYPE
|
|
|
|
|
|
class MissingConfigError(TypeError):
|
|
"""Raised when config is missing required values."""
|
|
|
|
|
|
class NodeRedEntity(Entity):
|
|
"""Base entity for Node-RED integration."""
|
|
|
|
component: ClassVar[str] = ""
|
|
_bidirectional = False
|
|
|
|
def __init__(self, hass: HomeAssistant, config: dict[str, Any]) -> None:
|
|
"""Initialize the entity."""
|
|
if not type(self).component:
|
|
msg = f"{type(self).__name__} must set class attribute `component`"
|
|
raise TypeError(msg)
|
|
self.hass = hass
|
|
self._server_id = config.get(CONF_SERVER_ID)
|
|
self._node_id = config.get(CONF_NODE_ID)
|
|
if self._server_id is None or self._node_id is None:
|
|
msg = f"{type(self).__name__} requires 'server_id' and 'node_id'"
|
|
raise MissingConfigError(msg)
|
|
self._attr_unique_id = f"{DOMAIN}-{self._server_id}-{self._node_id}"
|
|
self._attr_should_poll = False
|
|
|
|
device_info = config.get(CONF_DEVICE_INFO, {})
|
|
device_id = device_info.get("id")
|
|
if device_id:
|
|
self._attr_device_info = DeviceInfo(
|
|
identifiers=generate_device_identifiers(device_id),
|
|
hw_version=device_info.get("hw_version"),
|
|
manufacturer=device_info.get("manufacturer"),
|
|
model=device_info.get("model"),
|
|
name=device_info.get("name"),
|
|
sw_version=device_info.get("sw_version"),
|
|
)
|
|
else:
|
|
self._attr_device_info = None
|
|
|
|
self.update_discovery_config(config)
|
|
self.update_entity_state_attributes(config)
|
|
|
|
@callback
|
|
def handle_config_update(self, msg: dict[str, Any]) -> None:
|
|
"""Handle an incoming config update and write state."""
|
|
self.update_config(msg)
|
|
self.async_write_ha_state()
|
|
|
|
@callback
|
|
def handle_entity_update(self, msg: dict[str, Any]) -> None:
|
|
"""Update entity state attributes and write state."""
|
|
self.update_entity_state_attributes(msg)
|
|
self.async_write_ha_state()
|
|
|
|
def update_entity_state_attributes(self, msg: dict[str, Any]) -> None:
|
|
"""Set extra state attributes from incoming message."""
|
|
self._attr_extra_state_attributes = msg.get(CONF_ATTRIBUTES, {})
|
|
|
|
@callback
|
|
def handle_lost_connection(self) -> None:
|
|
"""Mark entity unavailable after losing connection."""
|
|
self._attr_available = False
|
|
self.async_write_ha_state()
|
|
|
|
@callback
|
|
def handle_discovery_update(
|
|
self, msg: dict[str, Any], connection: ActiveConnection
|
|
) -> None:
|
|
"""Update entity config/state based on discovery message."""
|
|
if CONF_REMOVE in msg:
|
|
if msg[CONF_REMOVE] == CHANGE_ENTITY_TYPE:
|
|
|
|
@callback
|
|
def recreate_entity() -> None:
|
|
del msg[CONF_REMOVE]
|
|
async_dispatcher_send(
|
|
self.hass,
|
|
NODERED_DISCOVERY.format(msg[CONF_COMPONENT]),
|
|
msg,
|
|
connection,
|
|
)
|
|
|
|
self.async_on_remove(recreate_entity)
|
|
# Schedule entity removal; recreation will happen when the on_remove
|
|
# callback runs and re-dispatches discovery for the new entity
|
|
self.hass.async_create_task(self.async_remove())
|
|
else:
|
|
|
|
@callback
|
|
def cleanup_discovery() -> None:
|
|
"""Remove discovery tracking for this entity, if present."""
|
|
with suppress(KeyError):
|
|
del self.hass.data[DOMAIN_DATA][ALREADY_DISCOVERED][
|
|
self.unique_id
|
|
]
|
|
|
|
self.async_on_remove(cleanup_discovery)
|
|
|
|
return
|
|
|
|
self.update_discovery_device_info(msg)
|
|
if CONF_CONFIG in msg:
|
|
self.update_discovery_config(msg)
|
|
|
|
# Bidirectional subscription handling
|
|
if self._bidirectional and "id" in msg:
|
|
self._attr_available = True
|
|
self._message_id = msg["id"]
|
|
self._connection = connection
|
|
self._connection.subscriptions[msg["id"]] = self.handle_lost_connection
|
|
|
|
self.async_write_ha_state()
|
|
|
|
def entity_category_mapper(self, category: str) -> None | EntityCategory:
|
|
"""Map Node-RED category strings to Home Assistant EntityCategory."""
|
|
if category == "config":
|
|
return EntityCategory.CONFIG
|
|
if category == "diagnostic":
|
|
return EntityCategory.DIAGNOSTIC
|
|
return None
|
|
|
|
def update_discovery_config(self, msg: dict[str, Any]) -> None:
|
|
"""Apply discovery config fields to the entity attributes."""
|
|
self._config = msg[CONF_CONFIG]
|
|
self._attr_icon = self._config.get(CONF_ICON)
|
|
self._attr_name = self._config.get(CONF_NAME, f"{DOMAIN} {self._node_id}")
|
|
self._attr_device_class = self._config.get(CONF_DEVICE_CLASS)
|
|
self._attr_entity_category = self.entity_category_mapper(
|
|
self._config.get(CONF_ENTITY_CATEGORY)
|
|
)
|
|
self._attr_entity_picture = self._config.get(CONF_ENTITY_PICTURE)
|
|
self._attr_unit_of_measurement = self._config.get(CONF_UNIT_OF_MEASUREMENT)
|
|
|
|
def update_config(self, msg: dict[str, Any]) -> None:
|
|
"""Apply runtime config updates to the entity."""
|
|
config = msg.get(CONF_CONFIG, {})
|
|
|
|
if config.get(CONF_NAME):
|
|
self._attr_name = config.get(CONF_NAME)
|
|
if config.get(CONF_ICON):
|
|
self._attr_icon = config.get(CONF_ICON)
|
|
if config.get(CONF_ENTITY_PICTURE):
|
|
self._attr_entity_picture = config.get(CONF_ENTITY_PICTURE)
|
|
if config.get(CONF_OPTIONS):
|
|
self._attr_options = config.get(CONF_OPTIONS)
|
|
|
|
def update_discovery_device_info(self, msg: dict[str, Any]) -> None:
|
|
"""Update or remove device info for this entity based on discovery msg."""
|
|
entity_registry = async_get(self.hass)
|
|
if self.unique_id is None:
|
|
raise RuntimeError("unique_id should be set during initialization")
|
|
entity_id = entity_registry.async_get_entity_id(
|
|
self.component,
|
|
DOMAIN,
|
|
self.unique_id,
|
|
)
|
|
device_info = msg.get(CONF_DEVICE_INFO)
|
|
|
|
# Exit early when there is no device info. If an entity existed, dissociate it from any device.
|
|
if device_info is None:
|
|
if entity_id is not None:
|
|
entity_registry.async_update_entity(entity_id, device_id=None)
|
|
return
|
|
|
|
device_registry = dr.async_get(self.hass)
|
|
identifiers = generate_device_identifiers(device_info["id"])
|
|
|
|
# Get or create the device
|
|
device = device_registry.async_get_or_create(
|
|
config_entry_id=self.hass.config_entries.async_entries(DOMAIN)[0].entry_id,
|
|
identifiers=identifiers,
|
|
name=device_info.get("name"),
|
|
manufacturer=device_info.get("manufacturer"),
|
|
model=device_info.get("model"),
|
|
)
|
|
|
|
# Update device properties
|
|
device_registry.async_update_device(
|
|
device_id=device.id,
|
|
hw_version=device_info.get("hw_version"),
|
|
name=device_info.get("name"),
|
|
manufacturer=device_info.get("manufacturer"),
|
|
model=device_info.get("model"),
|
|
sw_version=device_info.get("sw_version"),
|
|
)
|
|
|
|
# Associate entity with device
|
|
if entity_id is not None:
|
|
entity_registry.async_update_entity(entity_id, device_id=device.id)
|
|
|
|
async def async_added_to_hass(self) -> None:
|
|
"""Register dispatcher listeners when added to Home Assistant."""
|
|
self._remove_signal_entity_update = async_dispatcher_connect(
|
|
self.hass,
|
|
NODERED_ENTITY.format(self._server_id, self._node_id),
|
|
self.handle_entity_update,
|
|
)
|
|
self._remove_signal_discovery_update = async_dispatcher_connect(
|
|
self.hass,
|
|
NODERED_DISCOVERY_UPDATED.format(self.unique_id),
|
|
self.handle_discovery_update,
|
|
)
|
|
self._remove_signal_config_update = async_dispatcher_connect(
|
|
self.hass,
|
|
NODERED_CONFIG_UPDATE.format(self._server_id, self._node_id),
|
|
self.handle_config_update,
|
|
)
|
|
|
|
async def async_will_remove_from_hass(self) -> None:
|
|
"""Remove dispatcher listeners when the entity is removed from hass."""
|
|
if self._remove_signal_entity_update is not None:
|
|
self._remove_signal_entity_update()
|
|
if self._remove_signal_discovery_update is not None:
|
|
self._remove_signal_discovery_update()
|
|
if self._remove_signal_config_update is not None:
|
|
self._remove_signal_config_update()
|
|
|
|
|
|
def generate_device_identifiers(device_id: str) -> set[tuple[str, str]]:
|
|
"""Create identifiers set from device info."""
|
|
return {(DOMAIN, device_id)}
|