entity.py

Go to the documentation of this file.

"""Base entity for the Switch as X integration."""
 
from __future__ import annotations
 
from typing import Any
 
from homeassistant.components.homeassistant import exposed_entities
from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
from homeassistant.const import (
    ATTR_ENTITY_ID,
    SERVICE_TURN_OFF,
    SERVICE_TURN_ON,
    STATE_ON,
    STATE_UNAVAILABLE,
)
from homeassistant.core import Event, EventStateChangedData, HomeAssistant, callback
from homeassistant.helpers import device_registry as dr, entity_registry as er
from homeassistant.helpers.device_registry import DeviceInfo
from homeassistant.helpers.entity import Entity, ToggleEntity
from homeassistant.helpers.event import async_track_state_change_event
 
from .const import DOMAIN as SWITCH_AS_X_DOMAIN
 
 
class BaseEntity(Entity):
    """Represents a Switch as an X."""
 
    _attr_should_poll = False
    _is_new_entity: bool
 
    def __init__(
        self,
        hass: HomeAssistant,
        config_entry_title: str,
        domain: str,
        switch_entity_id: str,
        unique_id: str,
    ) -> None:
        """Initialize Switch as an X."""
        registry = er.async_get(hass)
        device_registry = dr.async_get(hass)
        wrapped_switch = registry.async_get(switch_entity_id)
        device_id = wrapped_switch.device_id if wrapped_switch else None
        entity_category = wrapped_switch.entity_category if wrapped_switch else None
        has_entity_name = wrapped_switch.has_entity_name if wrapped_switch else False
 
        name: str | None = config_entry_title
        if wrapped_switch:
            name = wrapped_switch.original_name
 
        self._device_id_device_id = device_id
        if device_id and (device := device_registry.async_get(device_id)):
            self._attr_device_info_attr_device_info = DeviceInfo(
                connections=device.connections,
                identifiers=device.identifiers,
            )
        self._attr_entity_category_attr_entity_category = entity_category
        self._attr_has_entity_name_attr_has_entity_name = has_entity_name
        self._attr_name_attr_name = name
        self._attr_unique_id_attr_unique_id = unique_id
        self._switch_entity_id_switch_entity_id = switch_entity_id
 
        self._is_new_entity_is_new_entity = (
            registry.async_get_entity_id(domain, SWITCH_AS_X_DOMAIN, unique_id) is None
        )
 
    @callback
    def async_state_changed_listener(
        self, event: Event[EventStateChangedData] | None = None
    ) -> None:
        """Handle child updates."""
        if (
            state := self.hasshass.states.get(self._switch_entity_id_switch_entity_id)
        ) is None or state.state == STATE_UNAVAILABLE:
            self._attr_available_attr_available = False
            return
 
        self._attr_available_attr_available = True
 
    async def async_added_to_hass(self) -> None:
        """Register callbacks and copy the wrapped entity's custom name if set."""
 
        @callback
        def _async_state_changed_listener(
            event: Event[EventStateChangedData] | None = None,
        ) -> None:
            """Handle child updates."""
            self.async_state_changed_listenerasync_state_changed_listener(event)
            self.async_write_ha_stateasync_write_ha_state()
 
        self.async_on_removeasync_on_remove(
            async_track_state_change_event(
                self.hasshass, [self._switch_entity_id_switch_entity_id], _async_state_changed_listener
            )
        )
 
        # Call once on adding
        _async_state_changed_listener()
 
        # Update entity options
        registry = er.async_get(self.hasshass)
        if registry.async_get(self.entity_identity_id) is not None:
            registry.async_update_entity_options(
                self.entity_identity_id,
                SWITCH_AS_X_DOMAIN,
                self.async_generate_entity_optionsasync_generate_entity_options(),
            )
 
        if not self._is_new_entity_is_new_entity or not (
            wrapped_switch := registry.async_get(self._switch_entity_id_switch_entity_id)
        ):
            return
 
        def copy_custom_name(wrapped_switch: er.RegistryEntry) -> None:
            """Copy the name set by user from the wrapped entity."""
            if wrapped_switch.name is None:
                return
            registry.async_update_entity(self.entity_identity_id, name=wrapped_switch.name)
 
        def copy_expose_settings() -> None:
            """Copy assistant expose settings from the wrapped entity.
 
            Also unexpose the wrapped entity if exposed.
            """
            expose_settings = exposed_entities.async_get_entity_settings(
                self.hasshass, self._switch_entity_id_switch_entity_id
            )
            for assistant, settings in expose_settings.items():
                if (should_expose := settings.get("should_expose")) is None:
                    continue
                exposed_entities.async_expose_entity(
                    self.hasshass, assistant, self.entity_identity_id, should_expose
                )
                exposed_entities.async_expose_entity(
                    self.hasshass, assistant, self._switch_entity_id_switch_entity_id, False
                )
 
        copy_custom_name(wrapped_switch)
        copy_expose_settings()
 
    @callback
    def async_generate_entity_options(self) -> dict[str, Any]:
        """Generate entity options."""
        return {"entity_id": self._switch_entity_id_switch_entity_id, "invert": False}
 
 
class BaseToggleEntity(BaseEntity, ToggleEntity):
    """Represents a Switch as a ToggleEntity."""
 
    async def async_turn_on(self, **kwargs: Any) -> None:
        """Forward the turn_on command to the switch in this light switch."""
        await self.hasshass.services.async_call(
            SWITCH_DOMAIN,
            SERVICE_TURN_ON,
            {ATTR_ENTITY_ID: self._switch_entity_id_switch_entity_id},
            blocking=True,
            context=self._context_context,
        )
 
    async def async_turn_off(self, **kwargs: Any) -> None:
        """Forward the turn_off command to the switch in this light switch."""
        await self.hasshass.services.async_call(
            SWITCH_DOMAIN,
            SERVICE_TURN_OFF,
            {ATTR_ENTITY_ID: self._switch_entity_id_switch_entity_id},
            blocking=True,
            context=self._context_context,
        )
 
    @callback
    def async_state_changed_listener(
        self, event: Event[EventStateChangedData] | None = None
    ) -> None:
        """Handle child updates."""
        super().async_state_changed_listener(event)
        if (
            not self.availableavailable
            or (state := self.hasshass.states.get(self._switch_entity_id_switch_entity_id)) is None
        ):
            return
 
        self._attr_is_on_attr_is_on = state.state == STATE_ON
 
 
class BaseInvertableEntity(BaseEntity):
    """Represents a Switch as an X."""
 
    def __init__(
        self,
        hass: HomeAssistant,
        config_entry_title: str,
        domain: str,
        invert: bool,
        switch_entity_id: str,
        unique_id: str,
    ) -> None:
        """Initialize Switch as an X."""
        super().__init__(hass, config_entry_title, domain, switch_entity_id, unique_id)
        self._invert_state_invert_state = invert
 
    @callback
    def async_generate_entity_options(self) -> dict[str, Any]:
        """Generate entity options."""
        return super().async_generate_entity_options() | {"invert": self._invert_state_invert_state}