From 7284af6a3ecdf1f92d92238b2ab4747c4fa6c312 Mon Sep 17 00:00:00 2001 From: David Poll Date: Mon, 13 Mar 2023 03:00:05 -0700 Subject: [PATCH] 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 --------- Co-authored-by: Erik Montnemery --- homeassistant/bootstrap.py | 2 + .../components/homeassistant/__init__.py | 25 ++++++-- .../components/homeassistant/services.yaml | 6 ++ homeassistant/helpers/template.py | 62 +++++++++++++++++++ tests/components/homeassistant/test_init.py | 18 ++++++ tests/helpers/test_template.py | 61 ++++++++++++++++++ .../custom_jinja/inner/inner_test.jinja | 5 ++ tests/testing_config/custom_jinja/test.jinja | 5 ++ 8 files changed, 178 insertions(+), 6 deletions(-) create mode 100644 tests/testing_config/custom_jinja/inner/inner_test.jinja create mode 100644 tests/testing_config/custom_jinja/test.jinja diff --git a/homeassistant/bootstrap.py b/homeassistant/bootstrap.py index 29772e865af..9ba4e99a082 100644 --- a/homeassistant/bootstrap.py +++ b/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), ) diff --git a/homeassistant/components/homeassistant/__init__.py b/homeassistant/components/homeassistant/__init__.py index 5602fd6b59a..4b033fd7119 100644 --- a/homeassistant/components/homeassistant/__init__.py +++ b/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 - `frontend.reload_themes` service, as that reloads the themes. + service, as that reloads the core YAML configuration, the + `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 { - ha.DOMAIN: SERVICE_RELOAD_CORE_CONFIG, - "frontend": "reload_themes", - }.items() + for domain, service in ( + (ha.DOMAIN, SERVICE_RELOAD_CORE_CONFIG), + ("frontend", "reload_themes"), + (ha.DOMAIN, SERVICE_RELOAD_CUSTOM_JINJA), + ) ] await asyncio.gather(*tasks) diff --git a/homeassistant/components/homeassistant/services.yaml b/homeassistant/components/homeassistant/services.yaml index da52ff50d2f..20f23402a73 100644 --- a/homeassistant/components/homeassistant/services.yaml +++ b/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. diff --git a/homeassistant/helpers/template.py b/homeassistant/helpers/template.py index 5205d51273f..1c5d15801f8 100644 --- a/homeassistant/helpers/template.py +++ b/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 diff --git a/tests/components/homeassistant/test_init.py b/tests/components/homeassistant/test_init.py index 8b982dc1c31..4a042416965 100644 --- a/tests/components/homeassistant/test_init.py +++ b/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 diff --git a/tests/helpers/test_template.py b/tests/helpers/test_template.py index 5122a4238ea..740040835c6 100644 --- a/tests/helpers/test_template.py +++ b/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 ( diff --git a/tests/testing_config/custom_jinja/inner/inner_test.jinja b/tests/testing_config/custom_jinja/inner/inner_test.jinja new file mode 100644 index 00000000000..e0cf20be762 --- /dev/null +++ b/tests/testing_config/custom_jinja/inner/inner_test.jinja @@ -0,0 +1,5 @@ +{% macro test_macro() -%} +inner macro +{%- endmacro %} + +{% set test_variable = "inner variable" %} \ No newline at end of file diff --git a/tests/testing_config/custom_jinja/test.jinja b/tests/testing_config/custom_jinja/test.jinja new file mode 100644 index 00000000000..44be0f3a57a --- /dev/null +++ b/tests/testing_config/custom_jinja/test.jinja @@ -0,0 +1,5 @@ +{% macro test_macro() -%} +macro +{%- endmacro %} + +{% set test_variable = "variable" %} \ No newline at end of file