Extract label template functions into a label Jinja2 extension (home-assistant#156439)

Co-authored-by: Copilot <[email protected]>

Author: frenck

homeassistant/helpers/template/__init__.py

@@ -61,7 +61,6 @@
     entity_registry as er,
     floor_registry as fr,
     issue_registry as ir,
-    label_registry as lr,
     location as loc_helper,
 )
 from homeassistant.helpers.singleton import singleton
@@ -1514,100 +1513,6 @@ def area_devices(hass: HomeAssistant, area_id_or_name: str) -> Iterable[str]:
     return [entry.id for entry in entries]
 
 
-def labels(hass: HomeAssistant, lookup_value: Any = None) -> Iterable[str | None]:
-    """Return all labels, or those from a area ID, device ID, or entity ID."""
-    label_reg = lr.async_get(hass)
-    if lookup_value is None:
-        return list(label_reg.labels)
-
-    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
-
-    lookup_value = str(lookup_value)
-
-    try:
-        cv.entity_id(lookup_value)
-    except vol.Invalid:
-        pass
-    else:
-        if entity := ent_reg.async_get(lookup_value):
-            return list(entity.labels)
-
-    # Check if this could be a device ID
-    dev_reg = dr.async_get(hass)
-    if device := dev_reg.async_get(lookup_value):
-        return list(device.labels)
-
-    # Check if this could be a area ID
-    area_reg = ar.async_get(hass)
-    if area := area_reg.async_get_area(lookup_value):
-        return list(area.labels)
-
-    return []
-
-
-def label_id(hass: HomeAssistant, lookup_value: Any) -> str | None:
-    """Get the label ID from a label name."""
-    label_reg = lr.async_get(hass)
-    if label := label_reg.async_get_label_by_name(str(lookup_value)):
-        return label.label_id
-    return None
-
-
-def label_name(hass: HomeAssistant, lookup_value: str) -> str | None:
-    """Get the label name from a label ID."""
-    label_reg = lr.async_get(hass)
-    if label := label_reg.async_get_label(lookup_value):
-        return label.name
-    return None
-
-
-def label_description(hass: HomeAssistant, lookup_value: str) -> str | None:
-    """Get the label description from a label ID."""
-    label_reg = lr.async_get(hass)
-    if label := label_reg.async_get_label(lookup_value):
-        return label.description
-    return None
-
-
-def _label_id_or_name(hass: HomeAssistant, label_id_or_name: str) -> str | None:
-    """Get the label ID from a label name or ID."""
-    # If label_name returns a value, we know the input was an ID, otherwise we
-    # assume it's a name, and if it's neither, we return early.
-    if label_name(hass, label_id_or_name) is not None:
-        return label_id_or_name
-    return label_id(hass, label_id_or_name)
-
-
-def label_areas(hass: HomeAssistant, label_id_or_name: str) -> Iterable[str]:
-    """Return areas for a given label ID or name."""
-    if (_label_id := _label_id_or_name(hass, label_id_or_name)) is None:
-        return []
-    area_reg = ar.async_get(hass)
-    entries = ar.async_entries_for_label(area_reg, _label_id)
-    return [entry.id for entry in entries]
-
-
-def label_devices(hass: HomeAssistant, label_id_or_name: str) -> Iterable[str]:
-    """Return device IDs for a given label ID or name."""
-    if (_label_id := _label_id_or_name(hass, label_id_or_name)) is None:
-        return []
-    dev_reg = dr.async_get(hass)
-    entries = dr.async_entries_for_label(dev_reg, _label_id)
-    return [entry.id for entry in entries]
-
-
-def label_entities(hass: HomeAssistant, label_id_or_name: str) -> Iterable[str]:
-    """Return entities for a given label ID or name."""
-    if (_label_id := _label_id_or_name(hass, label_id_or_name)) is None:
-        return []
-    ent_reg = er.async_get(hass)
-    entries = er.async_entries_for_label(ent_reg, _label_id)
-    return [entry.entity_id for entry in entries]
-
-
 def closest(hass: HomeAssistant, *args: Any) -> State | None:
     """Find closest entity.
 
@@ -2454,6 +2359,7 @@ def __init__(
             "homeassistant.helpers.template.extensions.CollectionExtension"
         )
         self.add_extension("homeassistant.helpers.template.extensions.CryptoExtension")
+        self.add_extension("homeassistant.helpers.template.extensions.LabelExtension")
         self.add_extension("homeassistant.helpers.template.extensions.MathExtension")
         self.add_extension("homeassistant.helpers.template.extensions.RegexExtension")
         self.add_extension("homeassistant.helpers.template.extensions.StringExtension")
@@ -2603,29 +2509,6 @@ def wrapper(_: Any, *args: _P.args, **kwargs: _P.kwargs) -> _R:
         self.globals["device_id"] = hassfunction(device_id)
         self.filters["device_id"] = self.globals["device_id"]
 
-        # Label extensions
-
-        self.globals["labels"] = hassfunction(labels)
-        self.filters["labels"] = self.globals["labels"]
-
-        self.globals["label_id"] = hassfunction(label_id)
-        self.filters["label_id"] = self.globals["label_id"]
-
-        self.globals["label_name"] = hassfunction(label_name)
-        self.filters["label_name"] = self.globals["label_name"]
-
-        self.globals["label_description"] = hassfunction(label_description)
-        self.filters["label_description"] = self.globals["label_description"]
-
-        self.globals["label_areas"] = hassfunction(label_areas)
-        self.filters["label_areas"] = self.globals["label_areas"]
-
-        self.globals["label_devices"] = hassfunction(label_devices)
-        self.filters["label_devices"] = self.globals["label_devices"]
-
-        self.globals["label_entities"] = hassfunction(label_entities)
-        self.filters["label_entities"] = self.globals["label_entities"]
-
         # Issue extensions
 
         self.globals["issues"] = hassfunction(issues)
@@ -2658,8 +2541,6 @@ def warn_unsupported(*args: Any, **kwargs: Any) -> NoReturn:
                 "is_hidden_entity",
                 "is_state_attr",
                 "is_state",
-                "label_id",
-                "label_name",
                 "now",
                 "relative_time",
                 "state_attr",
@@ -2679,8 +2560,6 @@ def warn_unsupported(*args: Any, **kwargs: Any) -> NoReturn:
                 "floor_id",
                 "floor_name",
                 "has_value",
-                "label_id",
-                "label_name",
             ]
             hass_tests = [
                 "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 .labels import LabelExtension
 from .math import MathExtension
 from .regex import RegexExtension
 from .string import StringExtension
@@ -11,6 +12,7 @@
     "Base64Extension",
     "CollectionExtension",
     "CryptoExtension",
+    "LabelExtension",
     "MathExtension",
     "RegexExtension",
     "StringExtension",

homeassistant/helpers/template/extensions/base.py

@@ -4,8 +4,10 @@
 
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, NoReturn
+from functools import wraps
+from typing import TYPE_CHECKING, Any, Concatenate, NoReturn
 
+from jinja2 import pass_context
 from jinja2.ext import Extension
 from jinja2.nodes import Node
 from jinja2.parser import Parser
@@ -32,6 +34,28 @@ class TemplateFunction:
     requires_hass: bool = False  # Whether this function requires hass to be available
 
 
+def _pass_context[**_P, _R](
+    func: Callable[Concatenate[Any, _P], _R],
+    jinja_context: Callable[
+        [Callable[Concatenate[Any, _P], _R]],
+        Callable[Concatenate[Any, _P], _R],
+    ] = pass_context,
+) -> Callable[Concatenate[Any, _P], _R]:
+    """Wrap function to pass context.
+
+    We mark these as a context functions to ensure they get
+    evaluated fresh with every execution, rather than executed
+    at compile time and the value stored. The context itself
+    can be discarded.
+    """
+
+    @wraps(func)
+    def wrapper(_: Any, *args: _P.args, **kwargs: _P.kwargs) -> _R:
+        return func(*args, **kwargs)
+
+    return jinja_context(wrapper)
+
+
 class BaseTemplateExtension(Extension):
     """Base class for Home Assistant template extensions."""
 
@@ -65,12 +89,20 @@ def __init__(
                         environment.tests[template_func.name] = unsupported_func
                     continue
 
+                func = template_func.func
+
+                if template_func.requires_hass:
+                    # We wrap these as a context functions to ensure they get
+                    # evaluated fresh with every execution, rather than executed
+                    # at compile time and the value stored.
+                    func = _pass_context(func)
+
                 if template_func.as_global:
-                    environment.globals[template_func.name] = template_func.func
+                    environment.globals[template_func.name] = func
                 if template_func.as_filter:
-                    environment.filters[template_func.name] = template_func.func
+                    environment.filters[template_func.name] = func
                 if template_func.as_test:
-                    environment.tests[template_func.name] = template_func.func
+                    environment.tests[template_func.name] = func
 
     @staticmethod
     def _create_unsupported_function(name: str) -> Callable[[], NoReturn]: