Extend base jinja2 extension with limited template errors (home-assistant#156431)

Author: frenck

homeassistant/helpers/template/extensions/base.py

@@ -4,12 +4,14 @@
 
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, NoReturn
 
 from jinja2.ext import Extension
 from jinja2.nodes import Node
 from jinja2.parser import Parser
 
+from homeassistant.exceptions import TemplateError
+
 if TYPE_CHECKING:
     from homeassistant.core import HomeAssistant
     from homeassistant.helpers.template import TemplateEnvironment
@@ -50,8 +52,17 @@ def __init__(
                 if template_func.requires_hass and self.environment.hass is None:
                     continue
 
-                # Skip functions not allowed in limited environments
+                # Register unsupported stub for functions not allowed in limited environments
                 if self.environment.limited and not template_func.limited_ok:
+                    unsupported_func = self._create_unsupported_function(
+                        template_func.name
+                    )
+                    if template_func.as_global:
+                        environment.globals[template_func.name] = unsupported_func
+                    if template_func.as_filter:
+                        environment.filters[template_func.name] = unsupported_func
+                    if template_func.as_test:
+                        environment.tests[template_func.name] = unsupported_func
                     continue
 
                 if template_func.as_global:
@@ -61,6 +72,17 @@ def __init__(
                 if template_func.as_test:
                     environment.tests[template_func.name] = template_func.func
 
+    @staticmethod
+    def _create_unsupported_function(name: str) -> Callable[[], NoReturn]:
+        """Create a function that raises an error for unsupported functions in limited templates."""
+
+        def unsupported(*args: Any, **kwargs: Any) -> NoReturn:
+            raise TemplateError(
+                f"Use of '{name}' is not supported in limited templates"
+            )
+
+        return unsupported
+
     @property
     def hass(self) -> HomeAssistant:
         """Return the Home Assistant instance.

tests/helpers/template/extensions/test_base.py

@@ -4,6 +4,7 @@
 
 import pytest
 
+from homeassistant.exceptions import TemplateError
 from homeassistant.helpers.template import TemplateEnvironment
 from homeassistant.helpers.template.extensions.base import (
     BaseTemplateExtension,
@@ -86,7 +87,7 @@ def test_func() -> str:
 
 
 def test_limited_ok_functions_not_registered_in_limited_env() -> None:
-    """Test that functions with limited_ok=False are not registered in limited env."""
+    """Test that functions with limited_ok=False raise error in limited env."""
     # Create a limited environment without hass
     env = TemplateEnvironment(None, limited=True)
 
@@ -116,9 +117,18 @@ def restricted_func() -> str:
         ],
     )
 
-    # Only the allowed function should be registered
+    # The allowed function should be registered and work
     assert "allowed_func" in env.globals
-    assert "restricted_func" not in env.globals
+    assert env.globals["allowed_func"]() == "allowed"
+
+    # The restricted function should be registered but raise TemplateError
+    assert "restricted_func" in env.globals
+    with pytest.raises(
+        TemplateError,
+        match="Use of 'restricted_func' is not supported in limited templates",
+    ):
+        env.globals["restricted_func"]()
+
     assert extension is not None