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>
This commit is contained in:
David Poll 2023-03-13 03:00:05 -07:00 committed by GitHub
parent 5c4f93fa36
commit 7284af6a3e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 178 additions and 6 deletions

View file

@ -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),
)

View file

@ -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)

View file

@ -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.

View file

@ -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

View file

@ -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

View file

@ -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 (

View file

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

View file

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