"""TrueNAS HA shared entity model.""" from __future__ import annotations from collections.abc import Mapping from logging import getLogger from typing import Any, Callable from homeassistant.config_entries import ConfigEntry from homeassistant.const import ATTR_ATTRIBUTION, CONF_HOST, CONF_NAME from homeassistant.core import HomeAssistant, callback from homeassistant.helpers import ( entity_platform as ep, entity_registry as er, ) from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity import DeviceInfo, Entity from homeassistant.helpers.update_coordinator import CoordinatorEntity from homeassistant.util import slugify from .const import ( ATTRIBUTION, DOMAIN, ) from .coordinator import TrueNASCoordinator from .helper import format_attribute _LOGGER = getLogger(__name__) # --------------------------- # async_add_entities # --------------------------- async def async_add_entities( hass: HomeAssistant, config_entry: ConfigEntry, dispatcher: dict[str, Callable] ): """Add entities.""" platform = ep.async_get_current_platform() services = platform.platform.SENSOR_SERVICES descriptions = platform.platform.SENSOR_TYPES for service in services: platform.async_register_entity_service(service[0], service[1], service[2]) @callback async def async_update_controller(coordinator): """Update the values of the controller.""" async def async_check_exist(obj, coordinator, uid: None) -> None: """Check entity exists.""" entity_registry = er.async_get(hass) if uid: unique_id = f"{obj._inst.lower()}-{obj.entity_description.key}-{slugify(str(obj._data[obj.entity_description.data_reference]).lower())}" else: unique_id = f"{obj._inst.lower()}-{obj.entity_description.key}" entity_id = entity_registry.async_get_entity_id( platform.domain, DOMAIN, unique_id ) entity = entity_registry.async_get(entity_id) if entity is None or ( (entity_id not in platform.entities) and (entity.disabled is False) ): _LOGGER.debug("Add entity %s", entity_id) await platform.async_add_entities([obj]) for entity_description in descriptions: data = coordinator.data[entity_description.data_path] if not entity_description.data_reference: if data.get(entity_description.data_attribute) is None: continue obj = dispatcher[entity_description.func]( coordinator, entity_description ) await async_check_exist(obj, coordinator, None) else: for uid in data: obj = dispatcher[entity_description.func]( coordinator, entity_description, uid ) await async_check_exist(obj, coordinator, uid) await async_update_controller(hass.data[DOMAIN][config_entry.entry_id]) unsub = async_dispatcher_connect(hass, "update_sensors", async_update_controller) config_entry.async_on_unload(unsub) # --------------------------- # TrueNASEntity # --------------------------- class TrueNASEntity(CoordinatorEntity[TrueNASCoordinator], Entity): """Define entity.""" _attr_has_entity_name = True def __init__( self, coordinator: TrueNASCoordinator, entity_description, uid: str | None = None, ): """Initialize entity.""" super().__init__(coordinator) self.entity_description = entity_description self._inst = coordinator.config_entry.data[CONF_NAME] self._config_entry = self.coordinator.config_entry self._attr_extra_state_attributes = {ATTR_ATTRIBUTION: ATTRIBUTION} self._uid = uid if self._uid: self._data = coordinator.data[self.entity_description.data_path][self._uid] else: self._data = coordinator.data[self.entity_description.data_path] platform = ep.async_get_current_platform() dev_group = self.entity_description.ha_group if self.entity_description.ha_group.startswith("data__"): dev_group = self.entity_description.ha_group[6:] if dev_group in self._data: dev_group = self._data[dev_group] self.entity_id = f"{platform.domain}.{self._inst.lower()}_{slugify(str(dev_group).lower())}_{slugify(str(self.name).lower())}" @callback def _handle_coordinator_update(self) -> None: self._data = self.coordinator.data[self.entity_description.data_path] if self._uid: self._data = self.coordinator.data[self.entity_description.data_path][ self._uid ] super()._handle_coordinator_update() @property def name(self) -> str: """Return the name for this entity.""" if not self._uid: return f"{self.entity_description.name}" if self.entity_description.name: return f"{self._data[self.entity_description.data_name]} {self.entity_description.name}" return f"{self._data[self.entity_description.data_name]}" @property def unique_id(self) -> str: """Return a unique id for this entity.""" if self._uid: return f"{self._inst.lower()}-{self.entity_description.key}-{slugify(str(self._data[self.entity_description.data_reference]).lower())}" else: return f"{self._inst.lower()}-{self.entity_description.key}" @property def device_info(self) -> DeviceInfo: """Return a description for device registry.""" dev_connection = DOMAIN dev_connection_value = f"{self._inst}_{self.entity_description.ha_group}" dev_group = self.entity_description.ha_group if self.entity_description.ha_group == "System": dev_connection_value = ( f"{self._inst}_{self.coordinator.data['system_info']['hostname']}" ) if self.entity_description.ha_group.startswith("data__"): dev_group = self.entity_description.ha_group[6:] if dev_group in self._data: dev_group = self._data[dev_group] dev_connection_value = dev_group if self.entity_description.ha_connection: dev_connection = self.entity_description.ha_connection if self.entity_description.ha_connection_value: dev_connection_value = self.entity_description.ha_connection_value if dev_connection_value.startswith("data__"): dev_connection_value = f"{self._inst}_{dev_connection_value[6:]}" dev_connection_value = ( f"{self._inst}_{self._data[dev_connection_value]}" ) if self.entity_description.ha_group == "System": return DeviceInfo( connections={(dev_connection, f"{dev_connection_value}")}, identifiers={(dev_connection, f"{dev_connection_value}")}, name=dev_group, model=f"{self.coordinator.data['system_info']['system_product']}", manufacturer=f"{self.coordinator.data['system_info']['system_manufacturer']}", sw_version=f"{self.coordinator.data['system_info']['version']}", configuration_url=f"http://{self.coordinator.config_entry.data[CONF_HOST]}", ) else: return DeviceInfo( connections={(dev_connection, f"{dev_connection_value}")}, default_name=dev_group, default_model=f"{self.coordinator.data['system_info']['system_product']}", default_manufacturer=f"{self.coordinator.data['system_info']['system_manufacturer']}", via_device=( DOMAIN, f"{self._inst}_{self.coordinator.data['system_info']['hostname']}", ), ) @property def extra_state_attributes(self) -> Mapping[str, Any]: """Return the state attributes.""" attributes = super().extra_state_attributes for variable in self.entity_description.data_attributes_list: if variable in self._data: attributes[format_attribute(variable)] = self._data[variable] return attributes async def start(self): """Run function.""" raise NotImplementedError() async def stop(self): """Stop function.""" raise NotImplementedError() async def restart(self): """Restart function.""" raise NotImplementedError() async def reload(self): """Reload function.""" raise NotImplementedError() async def snapshot(self): """Snapshot function.""" raise NotImplementedError()