Extract device template functions into a devices Jinja2 extension (home-assistant#156619)

Author: frenck

homeassistant/helpers/template/__init__.py

@@ -1166,13 +1166,6 @@ def expand(hass: HomeAssistant, *args: Any) -> Iterable[State]:
     return list(found.values())
 
 
-def device_entities(hass: HomeAssistant, _device_id: str) -> Iterable[str]:
-    """Get entity ids for entities tied to a device."""
-    entity_reg = er.async_get(hass)
-    entries = er.async_entries_for_device(entity_reg, _device_id)
-    return [entry.entity_id for entry in entries]
-
-
 def integration_entities(hass: HomeAssistant, entry_name: str) -> Iterable[str]:
     """Get entity ids for entities tied to an integration/domain.
 
@@ -1214,65 +1207,6 @@ def config_entry_id(hass: HomeAssistant, entity_id: str) -> str | None:
     return None
 
 
-def device_id(hass: HomeAssistant, entity_id_or_device_name: str) -> str | None:
-    """Get a device ID from an entity ID or device name."""
-    entity_reg = er.async_get(hass)
-    entity = entity_reg.async_get(entity_id_or_device_name)
-    if entity is not None:
-        return entity.device_id
-
-    dev_reg = dr.async_get(hass)
-    return next(
-        (
-            device_id
-            for device_id, device in dev_reg.devices.items()
-            if (name := device.name_by_user or device.name)
-            and (str(entity_id_or_device_name) == name)
-        ),
-        None,
-    )
-
-
-def device_name(hass: HomeAssistant, lookup_value: str) -> str | None:
-    """Get the device name from an device id, or entity id."""
-    device_reg = dr.async_get(hass)
-    if device := device_reg.async_get(lookup_value):
-        return device.name_by_user or device.name
-
-    ent_reg = er.async_get(hass)
-    # Import here, not at top-level to avoid circular import
-    from homeassistant.helpers import config_validation as cv  # noqa: PLC0415
-
-    try:
-        cv.entity_id(lookup_value)
-    except vol.Invalid:
-        pass
-    else:
-        if entity := ent_reg.async_get(lookup_value):
-            if entity.device_id and (device := device_reg.async_get(entity.device_id)):
-                return device.name_by_user or device.name
-
-    return None
-
-
-def device_attr(hass: HomeAssistant, device_or_entity_id: str, attr_name: str) -> Any:
-    """Get the device specific attribute."""
-    device_reg = dr.async_get(hass)
-    if not isinstance(device_or_entity_id, str):
-        raise TemplateError("Must provide a device or entity ID")
-    device = None
-    if (
-        "." in device_or_entity_id
-        and (_device_id := device_id(hass, device_or_entity_id)) is not None
-    ):
-        device = device_reg.async_get(_device_id)
-    elif "." not in device_or_entity_id:
-        device = device_reg.async_get(device_or_entity_id)
-    if device is None or not hasattr(device, attr_name):
-        return None
-    return getattr(device, attr_name)
-
-
 def config_entry_attr(
     hass: HomeAssistant, config_entry_id_: str, attr_name: str
 ) -> Any:
@@ -1291,13 +1225,6 @@ def config_entry_attr(
     return getattr(config_entry, attr_name)
 
 
-def is_device_attr(
-    hass: HomeAssistant, device_or_entity_id: str, attr_name: str, attr_value: Any
-) -> bool:
-    """Test if a device's attribute is a specific value."""
-    return bool(device_attr(hass, device_or_entity_id, attr_name) == attr_value)
-
-
 def issues(hass: HomeAssistant) -> dict[tuple[str, str], dict[str, Any]]:
     """Return all open issues."""
     current_issues = ir.async_get(hass).issues
@@ -2260,6 +2187,7 @@ def __init__(
             "homeassistant.helpers.template.extensions.CollectionExtension"
         )
         self.add_extension("homeassistant.helpers.template.extensions.CryptoExtension")
+        self.add_extension("homeassistant.helpers.template.extensions.DeviceExtension")
         self.add_extension("homeassistant.helpers.template.extensions.FloorExtension")
         self.add_extension("homeassistant.helpers.template.extensions.LabelExtension")
         self.add_extension("homeassistant.helpers.template.extensions.MathExtension")
@@ -2377,23 +2305,6 @@ def wrapper(_: Any, *args: _P.args, **kwargs: _P.kwargs) -> _R:
         self.globals["config_entry_id"] = hassfunction(config_entry_id)
         self.filters["config_entry_id"] = self.globals["config_entry_id"]
 
-        # Device extensions
-
-        self.globals["device_name"] = hassfunction(device_name)
-        self.filters["device_name"] = self.globals["device_name"]
-
-        self.globals["device_attr"] = hassfunction(device_attr)
-        self.filters["device_attr"] = self.globals["device_attr"]
-
-        self.globals["device_entities"] = hassfunction(device_entities)
-        self.filters["device_entities"] = self.globals["device_entities"]
-
-        self.globals["is_device_attr"] = hassfunction(is_device_attr)
-        self.tests["is_device_attr"] = hassfunction(is_device_attr, pass_eval_context)
-
-        self.globals["device_id"] = hassfunction(device_id)
-        self.filters["device_id"] = self.globals["device_id"]
-
         # Issue extensions
 
         self.globals["issues"] = hassfunction(issues)
@@ -2415,12 +2326,9 @@ def warn_unsupported(*args: Any, **kwargs: Any) -> NoReturn:
                 "area_id",
                 "area_name",
                 "closest",
-                "device_attr",
-                "device_id",
                 "distance",
                 "expand",
                 "has_value",
-                "is_device_attr",
                 "is_hidden_entity",
                 "is_state_attr",
                 "is_state",
@@ -2438,7 +2346,6 @@ def warn_unsupported(*args: Any, **kwargs: Any) -> NoReturn:
                 "area_id",
                 "area_name",
                 "closest",
-                "device_id",
                 "expand",
                 "has_value",
             ]

homeassistant/helpers/template/extensions/__init__.py

@@ -3,6 +3,7 @@
 from .base64 import Base64Extension
 from .collection import CollectionExtension
 from .crypto import CryptoExtension
+from .devices import DeviceExtension
 from .floors import FloorExtension
 from .labels import LabelExtension
 from .math import MathExtension
@@ -13,6 +14,7 @@
     "Base64Extension",
     "CollectionExtension",
     "CryptoExtension",
+    "DeviceExtension",
     "FloorExtension",
     "LabelExtension",
     "MathExtension",

homeassistant/helpers/template/extensions/devices.py

@@ -0,0 +1,139 @@
+"""Device functions for Home Assistant templates."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, Any
+
+import voluptuous as vol
+
+from homeassistant.exceptions import TemplateError
+from homeassistant.helpers import (
+    config_validation as cv,
+    device_registry as dr,
+    entity_registry as er,
+)
+
+from .base import BaseTemplateExtension, TemplateFunction
+
+if TYPE_CHECKING:
+    from homeassistant.helpers.template import TemplateEnvironment
+
+
+class DeviceExtension(BaseTemplateExtension):
+    """Extension for device-related template functions."""
+
+    def __init__(self, environment: TemplateEnvironment) -> None:
+        """Initialize the device extension."""
+        super().__init__(
+            environment,
+            functions=[
+                TemplateFunction(
+                    "device_entities",
+                    self.device_entities,
+                    as_global=True,
+                    as_filter=True,
+                    requires_hass=True,
+                ),
+                TemplateFunction(
+                    "device_id",
+                    self.device_id,
+                    as_global=True,
+                    as_filter=True,
+                    requires_hass=True,
+                    limited_ok=False,
+                ),
+                TemplateFunction(
+                    "device_name",
+                    self.device_name,
+                    as_global=True,
+                    as_filter=True,
+                    requires_hass=True,
+                    limited_ok=False,
+                ),
+                TemplateFunction(
+                    "device_attr",
+                    self.device_attr,
+                    as_global=True,
+                    as_filter=True,
+                    requires_hass=True,
+                    limited_ok=False,
+                ),
+                TemplateFunction(
+                    "is_device_attr",
+                    self.is_device_attr,
+                    as_global=True,
+                    as_test=True,
+                    requires_hass=True,
+                    limited_ok=False,
+                ),
+            ],
+        )
+
+    def device_entities(self, _device_id: str) -> Iterable[str]:
+        """Get entity ids for entities tied to a device."""
+        entity_reg = er.async_get(self.hass)
+        entries = er.async_entries_for_device(entity_reg, _device_id)
+        return [entry.entity_id for entry in entries]
+
+    def device_id(self, entity_id_or_device_name: str) -> str | None:
+        """Get a device ID from an entity ID or device name."""
+        entity_reg = er.async_get(self.hass)
+        entity = entity_reg.async_get(entity_id_or_device_name)
+        if entity is not None:
+            return entity.device_id
+
+        dev_reg = dr.async_get(self.hass)
+        return next(
+            (
+                device_id
+                for device_id, device in dev_reg.devices.items()
+                if (name := device.name_by_user or device.name)
+                and (str(entity_id_or_device_name) == name)
+            ),
+            None,
+        )
+
+    def device_name(self, lookup_value: str) -> str | None:
+        """Get the device name from an device id, or entity id."""
+        device_reg = dr.async_get(self.hass)
+        if device := device_reg.async_get(lookup_value):
+            return device.name_by_user or device.name
+
+        ent_reg = er.async_get(self.hass)
+
+        try:
+            cv.entity_id(lookup_value)
+        except vol.Invalid:
+            pass
+        else:
+            if entity := ent_reg.async_get(lookup_value):
+                if entity.device_id and (
+                    device := device_reg.async_get(entity.device_id)
+                ):
+                    return device.name_by_user or device.name
+
+        return None
+
+    def device_attr(self, device_or_entity_id: str, attr_name: str) -> Any:
+        """Get the device specific attribute."""
+        device_reg = dr.async_get(self.hass)
+        if not isinstance(device_or_entity_id, str):
+            raise TemplateError("Must provide a device or entity ID")
+        device = None
+        if (
+            "." in device_or_entity_id
+            and (_device_id := self.device_id(device_or_entity_id)) is not None
+        ):
+            device = device_reg.async_get(_device_id)
+        elif "." not in device_or_entity_id:
+            device = device_reg.async_get(device_or_entity_id)
+        if device is None or not hasattr(device, attr_name):
+            return None
+        return getattr(device, attr_name)
+
+    def is_device_attr(
+        self, device_or_entity_id: str, attr_name: str, attr_value: Any
+    ) -> bool:
+        """Test if a device's attribute is a specific value."""
+        return bool(self.device_attr(device_or_entity_id, attr_name) == attr_value)