Move TriggerBaseEntity into helpers (home-assistant#91945)

* Move TriggerBaseEntity * mypy
konnected-io · Apr 25, 2023 · 1fa82fa · 1fa82fa
1 parent b56f0ad
commit 1fa82fa
Show file tree

Hide file tree

Showing 2 changed files with 160 additions and 166 deletions.
diff --git a/homeassistant/components/template/trigger_entity.py b/homeassistant/components/template/trigger_entity.py
@@ -1,174 +1,11 @@
 """Trigger entity."""
 from __future__ import annotations
 
-import logging
-from typing import Any
-
-from homeassistant.const import (
-    ATTR_ENTITY_PICTURE,
-    ATTR_FRIENDLY_NAME,
-    ATTR_ICON,
-    CONF_DEVICE_CLASS,
-    CONF_ICON,
-    CONF_NAME,
-    CONF_UNIQUE_ID,
-)
-from homeassistant.core import HomeAssistant, State, callback
-from homeassistant.exceptions import TemplateError
-from homeassistant.helpers import template
-from homeassistant.helpers.entity import Entity
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers.template_entity import TriggerBaseEntity
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from . import TriggerUpdateCoordinator
-from .const import CONF_ATTRIBUTES, CONF_AVAILABILITY, CONF_PICTURE
-
-CONF_TO_ATTRIBUTE = {
-    CONF_ICON: ATTR_ICON,
-    CONF_NAME: ATTR_FRIENDLY_NAME,
-    CONF_PICTURE: ATTR_ENTITY_PICTURE,
-}
-
-
-class TriggerBaseEntity(Entity):
-    """Template Base entity based on trigger data."""
-
-    domain: str
-    extra_template_keys: tuple | None = None
-    extra_template_keys_complex: tuple | None = None
-    _unique_id: str | None
-
-    def __init__(
-        self,
-        hass: HomeAssistant,
-        config: dict,
-    ) -> None:
-        """Initialize the entity."""
-        self.hass = hass
-
-        self._set_unique_id(config.get(CONF_UNIQUE_ID))
-
-        self._config = config
-
-        self._static_rendered = {}
-        self._to_render_simple = []
-        self._to_render_complex: list[str] = []
-
-        for itm in (
-            CONF_AVAILABILITY,
-            CONF_ICON,
-            CONF_NAME,
-            CONF_PICTURE,
-        ):
-            if itm not in config:
-                continue
-
-            if config[itm].is_static:
-                self._static_rendered[itm] = config[itm].template
-            else:
-                self._to_render_simple.append(itm)
-
-        if self.extra_template_keys is not None:
-            self._to_render_simple.extend(self.extra_template_keys)
-
-        if self.extra_template_keys_complex is not None:
-            self._to_render_complex.extend(self.extra_template_keys_complex)
-
-        # We make a copy so our initial render is 'unknown' and not 'unavailable'
-        self._rendered = dict(self._static_rendered)
-        self._parse_result = {CONF_AVAILABILITY}
-
-    @property
-    def name(self):
-        """Name of the entity."""
-        return self._rendered.get(CONF_NAME)
-
-    @property
-    def unique_id(self):
-        """Return unique ID of the entity."""
-        return self._unique_id
-
-    @property
-    def device_class(self):
-        """Return device class of the entity."""
-        return self._config.get(CONF_DEVICE_CLASS)
-
-    @property
-    def icon(self) -> str | None:
-        """Return icon."""
-        return self._rendered.get(CONF_ICON)
-
-    @property
-    def entity_picture(self) -> str | None:
-        """Return entity picture."""
-        return self._rendered.get(CONF_PICTURE)
-
-    @property
-    def available(self):
-        """Return availability of the entity."""
-        return (
-            self._rendered is not self._static_rendered
-            and
-            # Check against False so `None` is ok
-            self._rendered.get(CONF_AVAILABILITY) is not False
-        )
-
-    @property
-    def extra_state_attributes(self) -> dict[str, Any] | None:
-        """Return extra attributes."""
-        return self._rendered.get(CONF_ATTRIBUTES)
-
-    async def async_added_to_hass(self) -> None:
-        """Handle being added to Home Assistant."""
-        template.attach(self.hass, self._config)
-
-    def _set_unique_id(self, unique_id: str | None) -> None:
-        """Set unique id."""
-        self._unique_id = unique_id
-
-    def restore_attributes(self, last_state: State) -> None:
-        """Restore attributes."""
-        for conf_key, attr in CONF_TO_ATTRIBUTE.items():
-            if conf_key not in self._config or attr not in last_state.attributes:
-                continue
-            self._rendered[conf_key] = last_state.attributes[attr]
-
-        if CONF_ATTRIBUTES in self._config:
-            extra_state_attributes = {}
-            for attr in self._config[CONF_ATTRIBUTES]:
-                if attr not in last_state.attributes:
-                    continue
-                extra_state_attributes[attr] = last_state.attributes[attr]
-            self._rendered[CONF_ATTRIBUTES] = extra_state_attributes
-
-    def _render_templates(self, variables: dict[str, Any]) -> None:
-        """Render templates."""
-        try:
-            rendered = dict(self._static_rendered)
-
-            for key in self._to_render_simple:
-                rendered[key] = self._config[key].async_render(
-                    variables,
-                    parse_result=key in self._parse_result,
-                )
-
-            for key in self._to_render_complex:
-                rendered[key] = template.render_complex(
-                    self._config[key],
-                    variables,
-                )
-
-            if CONF_ATTRIBUTES in self._config:
-                rendered[CONF_ATTRIBUTES] = template.render_complex(
-                    self._config[CONF_ATTRIBUTES],
-                    variables,
-                )
-
-            self._rendered = rendered
-        except TemplateError as err:
-            logging.getLogger(f"{__package__}.{self.entity_id.split('.')[0]}").error(
-                "Error rendering %s template for %s: %s", key, self.entity_id, err
-            )
-            self._rendered = self._static_rendered
 
 
 class TriggerEntity(TriggerBaseEntity, CoordinatorEntity[TriggerUpdateCoordinator]):

diff --git a/homeassistant/helpers/template_entity.py b/homeassistant/helpers/template_entity.py
@@ -16,6 +16,9 @@
 )
 from homeassistant.const import (
     ATTR_ENTITY_ID,
+    ATTR_ENTITY_PICTURE,
+    ATTR_FRIENDLY_NAME,
+    ATTR_ICON,
     CONF_DEVICE_CLASS,
     CONF_ICON,
     CONF_NAME,
@@ -31,7 +34,13 @@
 from .entity import Entity
 from .event import TrackTemplate, TrackTemplateResult, async_track_template_result
 from .script import Script, _VarsType
-from .template import Template, TemplateStateFromEntityId, result_as_boolean
+from .template import (
+    Template,
+    TemplateStateFromEntityId,
+    attach as template_attach,
+    render_complex,
+    result_as_boolean,
+)
 from .typing import ConfigType
 
 _LOGGER = logging.getLogger(__name__)
@@ -40,6 +49,12 @@
 CONF_ATTRIBUTES = "attributes"
 CONF_PICTURE = "picture"
 
+CONF_TO_ATTRIBUTE = {
+    CONF_ICON: ATTR_ICON,
+    CONF_NAME: ATTR_FRIENDLY_NAME,
+    CONF_PICTURE: ATTR_ENTITY_PICTURE,
+}
+
 TEMPLATE_ENTITY_BASE_SCHEMA = vol.Schema(
     {
         vol.Optional(CONF_ICON): cv.template,
@@ -440,3 +455,145 @@ def __init__(
         self._attr_native_unit_of_measurement = config.get(CONF_UNIT_OF_MEASUREMENT)
         self._attr_device_class = config.get(CONF_DEVICE_CLASS)
         self._attr_state_class = config.get(CONF_STATE_CLASS)
+
+
+class TriggerBaseEntity(Entity):
+    """Template Base entity based on trigger data."""
+
+    domain: str
+    extra_template_keys: tuple | None = None
+    extra_template_keys_complex: tuple | None = None
+    _unique_id: str | None
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        config: dict,
+    ) -> None:
+        """Initialize the entity."""
+        self.hass = hass
+
+        self._set_unique_id(config.get(CONF_UNIQUE_ID))
+
+        self._config = config
+
+        self._static_rendered = {}
+        self._to_render_simple = []
+        self._to_render_complex: list[str] = []
+
+        for itm in (
+            CONF_AVAILABILITY,
+            CONF_ICON,
+            CONF_NAME,
+            CONF_PICTURE,
+        ):
+            if itm not in config:
+                continue
+
+            if config[itm].is_static:
+                self._static_rendered[itm] = config[itm].template
+            else:
+                self._to_render_simple.append(itm)
+
+        if self.extra_template_keys is not None:
+            self._to_render_simple.extend(self.extra_template_keys)
+
+        if self.extra_template_keys_complex is not None:
+            self._to_render_complex.extend(self.extra_template_keys_complex)
+
+        # We make a copy so our initial render is 'unknown' and not 'unavailable'
+        self._rendered = dict(self._static_rendered)
+        self._parse_result = {CONF_AVAILABILITY}
+
+    @property
+    def name(self) -> str | None:
+        """Name of the entity."""
+        return self._rendered.get(CONF_NAME)
+
+    @property
+    def unique_id(self) -> str | None:
+        """Return unique ID of the entity."""
+        return self._unique_id
+
+    @property
+    def device_class(self):  # type: ignore[no-untyped-def]
+        """Return device class of the entity."""
+        return self._config.get(CONF_DEVICE_CLASS)
+
+    @property
+    def icon(self) -> str | None:
+        """Return icon."""
+        return self._rendered.get(CONF_ICON)
+
+    @property
+    def entity_picture(self) -> str | None:
+        """Return entity picture."""
+        return self._rendered.get(CONF_PICTURE)
+
+    @property
+    def available(self) -> bool:
+        """Return availability of the entity."""
+        return (
+            self._rendered is not self._static_rendered
+            and
+            # Check against False so `None` is ok
+            self._rendered.get(CONF_AVAILABILITY) is not False
+        )
+
+    @property
+    def extra_state_attributes(self) -> dict[str, Any] | None:
+        """Return extra attributes."""
+        return self._rendered.get(CONF_ATTRIBUTES)
+
+    async def async_added_to_hass(self) -> None:
+        """Handle being added to Home Assistant."""
+        template_attach(self.hass, self._config)
+
+    def _set_unique_id(self, unique_id: str | None) -> None:
+        """Set unique id."""
+        self._unique_id = unique_id
+
+    def restore_attributes(self, last_state: State) -> None:
+        """Restore attributes."""
+        for conf_key, attr in CONF_TO_ATTRIBUTE.items():
+            if conf_key not in self._config or attr not in last_state.attributes:
+                continue
+            self._rendered[conf_key] = last_state.attributes[attr]
+
+        if CONF_ATTRIBUTES in self._config:
+            extra_state_attributes = {}
+            for attr in self._config[CONF_ATTRIBUTES]:
+                if attr not in last_state.attributes:
+                    continue
+                extra_state_attributes[attr] = last_state.attributes[attr]
+            self._rendered[CONF_ATTRIBUTES] = extra_state_attributes
+
+    def _render_templates(self, variables: dict[str, Any]) -> None:
+        """Render templates."""
+        try:
+            rendered = dict(self._static_rendered)
+
+            for key in self._to_render_simple:
+                rendered[key] = self._config[key].async_render(
+                    variables,
+                    parse_result=key in self._parse_result,
+                )
+
+            for key in self._to_render_complex:
+                rendered[key] = render_complex(
+                    self._config[key],
+                    variables,
+                )
+
+            if CONF_ATTRIBUTES in self._config:
+                rendered[CONF_ATTRIBUTES] = render_complex(
+                    self._config[CONF_ATTRIBUTES],
+                    variables,
+                )
+
+            self._rendered = rendered
+        except TemplateError as err:
+            logging.getLogger(f"{__package__}.{self.entity_id.split('.')[0]}").error(
+                "Error rendering %s template for %s: %s", key, self.entity_id, err
+            )
+            self._rendered = self._static_rendered