From 4a7e4955e873b74009fa19703756d79b0aca91b8 Mon Sep 17 00:00:00 2001 From: Philipp Rasch Date: Wed, 15 Oct 2025 18:46:59 +0200 Subject: [PATCH] ops.files.template: allow custom jinja2 template loaders --- src/pyinfra/api/util.py | 11 +++++++---- src/pyinfra/operations/files.py | 11 +++++++---- 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/src/pyinfra/api/util.py b/src/pyinfra/api/util.py index 46d53ae2d..9ed19cc1f 100644 --- a/src/pyinfra/api/util.py +++ b/src/pyinfra/api/util.py @@ -10,7 +10,7 @@ from typing import IO, TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type, Union import click -from jinja2 import Environment, FileSystemLoader, StrictUndefined +from jinja2 import Environment, FileSystemLoader, StrictUndefined, Template from paramiko import SSHException from typeguard import TypeCheckError, check_type @@ -26,7 +26,7 @@ BLOCKSIZE = 65536 # Caches -TEMPLATES: Dict[Any, Any] = {} +TEMPLATES: Dict[str, Template] = {} FILE_SHAS: Dict[Any, Any] = {} PYINFRA_INSTALL_DIR = path.normpath(path.join(path.dirname(__file__), "..")) @@ -139,7 +139,9 @@ def get_operation_order_from_stack(state: "State"): return line_numbers -def get_template(filename_or_io: str | IO, jinja_env_kwargs: dict[str, Any] | None = None): +def get_template( + filename_or_io: str | IO, jinja_env_kwargs: dict[str, Any] | None = None +) -> Template: """ Gets a jinja2 ``Template`` object for the input filename or string, with caching based on the filename of the template, or the SHA1 of the input string. @@ -155,10 +157,11 @@ def get_template(filename_or_io: str | IO, jinja_env_kwargs: dict[str, Any] | No with file_data as file_io: template_string = file_io.read() + default_loader = FileSystemLoader(getcwd()) template = Environment( undefined=StrictUndefined, keep_trailing_newline=True, - loader=FileSystemLoader(getcwd()), + loader=jinja_env_kwargs.pop("loader", default_loader), **jinja_env_kwargs, ).from_string(template_string) diff --git a/src/pyinfra/operations/files.py b/src/pyinfra/operations/files.py index ff9df2adf..08f9ebe84 100644 --- a/src/pyinfra/operations/files.py +++ b/src/pyinfra/operations/files.py @@ -1156,11 +1156,14 @@ def template( if not set explicitly. Notes: - Common convention is to store templates in a "templates" directory and - have a filename suffix with '.j2' (for jinja2). + Common convention is to store templates in a "templates" directory and + have a filename suffix with '.j2' (for jinja2). - For information on the template syntax, see - `the jinja2 docs `_. + The default template lookup directory (used with jinjas ``extends``, ``import`` and + ``include`` statements) is the current working directory. + + For information on the template syntax, see + `the jinja2 docs `_. **Examples:**