Add an in-memory-preloading loader for Jinja imports (#88850)

* Adds a loader to enable jinja imports.

* Switch to in-memory

* Move loading custom_jinja off of the event loop

* Raise TemplateNotFound if template doesn't exist

* Fix docstring

* Adds a service to reload custom jinja

* Remove IO from test setup

* Improve coverage and small refactor

* Incorporate feedback and use .jinja extension

* Check the loaded sources in test.

* Incorporate PR feedback.

* Update homeassistant/helpers/template.py

Co-authored-by: Erik Montnemery <erik@montnemery.com>

---------

Co-authored-by: Erik Montnemery <erik@montnemery.com>

homeassistant/bootstrap.py

@@ -31,6 +31,7 @@ from .helpers import (
     entity_registry,
     issue_registry,
     recorder,
+    template,
 )
 from .helpers.dispatcher import async_dispatcher_send
 from .helpers.typing import ConfigType
@@ -244,6 +245,7 @@ async def load_registries(hass: core.HomeAssistant) -> None:
         entity_registry.async_load(hass),
         issue_registry.async_load(hass),
         hass.async_add_executor_job(_cache_uname_processor),
+        template.async_load_custom_jinja(hass),
     )
 
 

homeassistant/components/homeassistant/__init__.py

@@ -30,6 +30,7 @@ from homeassistant.helpers.service import (
     async_extract_referenced_entity_ids,
     async_register_admin_service,
 )
+from homeassistant.helpers.template import async_load_custom_jinja
 from homeassistant.helpers.typing import ConfigType
 
 ATTR_ENTRY_ID = "entry_id"
@@ -38,6 +39,7 @@ _LOGGER = logging.getLogger(__name__)
 DOMAIN = ha.DOMAIN
 SERVICE_RELOAD_CORE_CONFIG = "reload_core_config"
 SERVICE_RELOAD_CONFIG_ENTRY = "reload_config_entry"
+SERVICE_RELOAD_CUSTOM_JINJA = "reload_custom_jinja"
 SERVICE_CHECK_CONFIG = "check_config"
 SERVICE_UPDATE_ENTITY = "update_entity"
 SERVICE_SET_LOCATION = "set_location"
@@ -258,6 +260,14 @@ async def async_setup(hass: ha.HomeAssistant, config: ConfigType) -> bool:  # no
         vol.Schema({ATTR_LATITUDE: cv.latitude, ATTR_LONGITUDE: cv.longitude}),
     )
 
+    async def async_handle_reload_jinja(call: ha.ServiceCall) -> None:
+        """Service handler to reload custom Jinja."""
+        await async_load_custom_jinja(hass)
+
+    async_register_admin_service(
+        hass, ha.DOMAIN, SERVICE_RELOAD_CUSTOM_JINJA, async_handle_reload_jinja
+    )
+
     async def async_handle_reload_config_entry(call: ha.ServiceCall) -> None:
         """Service handler for reloading a config entry."""
         reload_entries = set()
@@ -288,8 +298,10 @@ async def async_setup(hass: ha.HomeAssistant, config: ConfigType) -> bool:  # no
         reload of YAML configurations for the domain that support it.
 
         Additionally, it also calls the `homeasssitant.reload_core_config`
-        service, as that reloads the core YAML configuration, and the
+        service, as that reloads the core YAML configuration, the
-        `frontend.reload_themes` service, as that reloads the themes.
+        `frontend.reload_themes` service that reloads the themes, and the
+        `homeassistant.reload_custom_jinja` service that reloads any custom
+        jinja into memory.
 
         We only do so, if there are no configuration errors.
         """
@@ -315,10 +327,11 @@ async def async_setup(hass: ha.HomeAssistant, config: ConfigType) -> bool:  # no
             hass.services.async_call(
                 domain, service, context=call.context, blocking=True
             )
-            for domain, service in {
+            for domain, service in (
-                ha.DOMAIN: SERVICE_RELOAD_CORE_CONFIG,
+                (ha.DOMAIN, SERVICE_RELOAD_CORE_CONFIG),
-                "frontend": "reload_themes",
+                ("frontend", "reload_themes"),
-            }.items()
+                (ha.DOMAIN, SERVICE_RELOAD_CUSTOM_JINJA),
+            )
         ]
 
         await asyncio.gather(*tasks)

homeassistant/components/homeassistant/services.yaml

@@ -59,6 +59,12 @@ update_entity:
   target:
     entity: {}
 
+reload_custom_jinja:
+  name: Reload custom Jinja2 templates
+  description: >-
+    Reload Jinja2 templates found in the custom_jinja folder in your config.
+    New values will be applied on the next render of the template.
+
 reload_config_entry:
   name: Reload config entry
   description: Reload a config entry that matches a target.

homeassistant/helpers/template.py

@@ -14,6 +14,7 @@ import json
 import logging
 import math
 from operator import attrgetter, contains
+import pathlib
 import random
 import re
 import statistics
@@ -73,6 +74,7 @@ from homeassistant.util.read_only_dict import ReadOnlyDict
 from homeassistant.util.thread import ThreadWithException
 
 from . import area_registry, device_registry, entity_registry, location as loc_helper
+from .singleton import singleton
 from .typing import TemplateVarsType
 
 # mypy: allow-untyped-defs, no-check-untyped-defs
@@ -85,6 +87,7 @@ _RENDER_INFO = "template.render_info"
 _ENVIRONMENT = "template.environment"
 _ENVIRONMENT_LIMITED = "template.environment_limited"
 _ENVIRONMENT_STRICT = "template.environment_strict"
+_HASS_LOADER = "template.hass_loader"
 
 _RE_JINJA_DELIMITERS = re.compile(r"\{%|\{\{|\{#")
 # Match "simple" ints and floats. -1.0, 1, +5, 5.0
@@ -120,6 +123,8 @@ template_cv: ContextVar[tuple[str, str] | None] = ContextVar(
 CACHED_TEMPLATE_STATES = 512
 EVAL_CACHE_SIZE = 512
 
+MAX_CUSTOM_JINJA_SIZE = 5 * 1024 * 1024
+
 
 @bind_hass
 def attach(hass: HomeAssistant, obj: Any) -> None:
@@ -2056,6 +2061,60 @@ class LoggingUndefined(jinja2.Undefined):
         return super().__bool__()
 
 
+async def async_load_custom_jinja(hass: HomeAssistant) -> None:
+    """Load all custom jinja files under 5MiB into memory."""
+    return await hass.async_add_executor_job(_load_custom_jinja, hass)
+
+
+def _load_custom_jinja(hass: HomeAssistant) -> None:
+    result = {}
+    jinja_path = hass.config.path("custom_jinja")
+    all_files = [
+        item
+        for item in pathlib.Path(jinja_path).rglob("*.jinja")
+        if item.is_file() and item.stat().st_size <= MAX_CUSTOM_JINJA_SIZE
+    ]
+    for file in all_files:
+        content = file.read_text()
+        path = str(file.relative_to(jinja_path))
+        result[path] = content
+
+    _get_hass_loader(hass).sources = result
+
+
+@singleton(_HASS_LOADER)
+def _get_hass_loader(hass: HomeAssistant) -> HassLoader:
+    return HassLoader({})
+
+
+class HassLoader(jinja2.BaseLoader):
+    """An in-memory jinja loader that keeps track of templates that need to be reloaded."""
+
+    def __init__(self, sources: dict[str, str]) -> None:
+        """Initialize an empty HassLoader."""
+        self._sources = sources
+        self._reload = 0
+
+    @property
+    def sources(self) -> dict[str, str]:
+        """Map filename to jinja source."""
+        return self._sources
+
+    @sources.setter
+    def sources(self, value: dict[str, str]) -> None:
+        self._sources = value
+        self._reload += 1
+
+    def get_source(
+        self, environment: jinja2.Environment, template: str
+    ) -> tuple[str, str | None, Callable[[], bool] | None]:
+        """Get in-memory sources."""
+        if template not in self._sources:
+            raise jinja2.TemplateNotFound(template)
+        cur_reload = self._reload
+        return self._sources[template], template, lambda: cur_reload == self._reload
+
+
 class TemplateEnvironment(ImmutableSandboxedEnvironment):
     """The Home Assistant template environment."""
 
@@ -2159,6 +2218,9 @@ class TemplateEnvironment(ImmutableSandboxedEnvironment):
         if hass is None:
             return
 
+        # This environment has access to hass, attach its loader to enable imports.
+        self.loader = _get_hass_loader(hass)
+
         # 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

tests/components/homeassistant/test_init.py

@@ -14,6 +14,7 @@ from homeassistant.components.homeassistant import (
     SERVICE_CHECK_CONFIG,
     SERVICE_RELOAD_ALL,
     SERVICE_RELOAD_CORE_CONFIG,
+    SERVICE_RELOAD_CUSTOM_JINJA,
     SERVICE_SET_LOCATION,
 )
 from homeassistant.const import (
@@ -575,6 +576,21 @@ async def test_save_persistent_states(hass: HomeAssistant) -> None:
         assert mock_save.called
 
 
+async def test_reload_custom_jinja(hass: HomeAssistant) -> None:
+    """Test we can call reload_custom_jinja."""
+    await async_setup_component(hass, "homeassistant", {})
+    with patch(
+        "homeassistant.components.homeassistant.async_load_custom_jinja",
+        return_value=None,
+    ) as mock_load_custom_jinja:
+        await hass.services.async_call(
+            "homeassistant",
+            SERVICE_RELOAD_CUSTOM_JINJA,
+            blocking=True,
+        )
+        assert mock_load_custom_jinja.called
+
+
 async def test_reload_all(
     hass: HomeAssistant, caplog: pytest.LogCaptureFixture
 ) -> None:
@@ -586,6 +602,7 @@ async def test_reload_all(
     notify = async_mock_service(hass, "notify", "reload")
     core_config = async_mock_service(hass, "homeassistant", "reload_core_config")
     themes = async_mock_service(hass, "frontend", "reload_themes")
+    jinja = async_mock_service(hass, "homeassistant", "reload_custom_jinja")
 
     with patch(
         "homeassistant.config.async_check_ha_config_file",
@@ -632,3 +649,4 @@ async def test_reload_all(
     assert len(test2) == 1
     assert len(core_config) == 1
     assert len(themes) == 1
+    assert len(jinja) == 1

tests/helpers/test_template.py

@@ -243,6 +243,67 @@ def test_iterating_domain_states(hass: HomeAssistant) -> None:
     )
 
 
+async def test_import(hass: HomeAssistant) -> None:
+    """Test that imports work from the config/custom_jinja folder."""
+    await template.async_load_custom_jinja(hass)
+    assert "test.jinja" in template._get_hass_loader(hass).sources
+    assert "inner/inner_test.jinja" in template._get_hass_loader(hass).sources
+    assert (
+        template.Template(
+            """
+            {% import 'test.jinja' as t %}
+            {{ t.test_macro() }} {{ t.test_variable }}
+            """,
+            hass,
+        ).async_render()
+        == "macro variable"
+    )
+
+    assert (
+        template.Template(
+            """
+            {% import 'inner/inner_test.jinja' as t %}
+            {{ t.test_macro() }} {{ t.test_variable }}
+            """,
+            hass,
+        ).async_render()
+        == "inner macro inner variable"
+    )
+
+    with pytest.raises(TemplateError):
+        template.Template(
+            """
+            {% import 'notfound.jinja' as t %}
+            {{ t.test_macro() }} {{ t.test_variable }}
+            """,
+            hass,
+        ).async_render()
+
+
+async def test_import_change(hass: HomeAssistant) -> None:
+    """Test that a change in HassLoader results in updated imports."""
+    await template.async_load_custom_jinja(hass)
+    to_test = template.Template(
+        """
+        {% import 'test.jinja' as t %}
+        {{ t.test_macro() }} {{ t.test_variable }}
+        """,
+        hass,
+    )
+    assert to_test.async_render() == "macro variable"
+
+    template._get_hass_loader(hass).sources = {
+        "test.jinja": """
+            {% macro test_macro() -%}
+            macro2
+            {%- endmacro %}
+
+            {% set test_variable = "variable2" %}
+            """
+    }
+    assert to_test.async_render() == "macro2 variable2"
+
+
 def test_loop_controls(hass: HomeAssistant) -> None:
     """Test that loop controls are enabled."""
     assert (

tests/testing_config/custom_jinja/inner/inner_test.jinja

@@ -0,0 +1,5 @@
+{% macro test_macro() -%}
+inner macro
+{%- endmacro %}
+
+{% set test_variable = "inner variable" %}

tests/testing_config/custom_jinja/test.jinja

@@ -0,0 +1,5 @@
+{% macro test_macro() -%}
+macro
+{%- endmacro %}
+
+{% set test_variable = "variable" %}