Refactor handling of Jinja2 templates and add filters

Author: ogenstad

infrahub_sdk/ctl/cli_commands.py

@@ -9,7 +9,6 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, Optional
 
-import jinja2
 import typer
 import ujson
 from rich.console import Console
@@ -18,7 +17,6 @@
 from rich.panel import Panel
 from rich.pretty import Pretty
 from rich.table import Table
-from rich.traceback import Traceback
 
 from .. import __version__ as sdk_version
 from ..async_typer import AsyncTyper
@@ -44,8 +42,14 @@
 )
 from ..ctl.validate import app as validate_app
 from ..exceptions import GraphQLError, ModuleImportError
-from ..jinja2 import identify_faulty_jinja_code
 from ..schema import MainSchemaTypesAll, SchemaRoot
+from ..template import Jinja2Template
+from ..template.exceptions import (
+    JinjaTemplateError,
+    JinjaTemplateNotFoundError,
+    JinjaTemplateSyntaxError,
+    JinjaTemplateUndefinedError,
+)
 from ..utils import get_branch, write_to_file
 from ..yaml import SchemaFile
 from .exporter import dump
@@ -174,37 +178,50 @@ async def run(
     await func(client=client, log=log, branch=branch, **variables_dict)
 
 
-def render_jinja2_template(template_path: Path, variables: dict[str, str], data: dict[str, Any]) -> str:
+async def render_jinja2_template(template_path: Path, variables: dict[str, Any], data: dict[str, Any]) -> str:
     if not template_path.is_file():
         console.print(f"[red]Unable to locate the template at {template_path}")
         raise typer.Exit(1)
 
-    templateLoader = jinja2.FileSystemLoader(searchpath=".")
-    templateEnv = jinja2.Environment(loader=templateLoader, trim_blocks=True, lstrip_blocks=True)
-    template = templateEnv.get_template(str(template_path))
-
+    variables["data"] = data
+    jinja_template = Jinja2Template(template_directory=Path())
     try:
-        rendered_tpl = template.render(**variables, data=data)  # type: ignore[arg-type]
-    except jinja2.TemplateSyntaxError as exc:
-        console.print("[red]Syntax Error detected on the template")
-        console.print(f"[yellow]  {exc}")
+        rendered_tpl = await jinja_template.render_from_file(template=template_path, variables=variables)
+    except JinjaTemplateNotFoundError as exc:
+        console.print("[red]An error occurred while rendering the jinja template")
+        console.print("")
+        if exc.base_template:
+            console.print(f"Base template: [yellow]{exc.base_template}")
+        console.print(f"Missing template: [yellow]{exc.filename}")
         raise typer.Exit(1) from exc
 
-    except jinja2.UndefinedError as exc:
+    except JinjaTemplateUndefinedError as exc:
+        console.print("[red]An error occurred while rendering the jinja template")
+        for error in exc.errors:
+            console.print(f"[yellow]{error.frame.filename} on line {error.frame.lineno}\n")
+            console.print(error.syntax)
+        console.print("")
+        console.print(exc.message)
+        raise typer.Exit(1) from exc
+    except JinjaTemplateSyntaxError as exc:
+        console.print("[red]A syntax error was encountered within the template")
+        console.print("")
+        if exc.filename:
+            console.print(f"Filename: [yellow]{exc.filename}")
+        console.print(f"Line number: [yellow]{exc.lineno}")
+        console.print()
+        console.print(exc.message)
+        raise typer.Exit(1) from exc
+    except JinjaTemplateError as exc:
         console.print("[red]An error occurred while rendering the jinja template")
-        traceback = Traceback(show_locals=False)
-        errors = identify_faulty_jinja_code(traceback=traceback)
-        for frame, syntax in errors:
-            console.print(f"[yellow]{frame.filename} on line {frame.lineno}\n")
-            console.print(syntax)
         console.print("")
-        console.print(traceback.trace.stacks[0].exc_value)
+        console.print(f"[yellow]{exc.message}")
         raise typer.Exit(1) from exc
 
     return rendered_tpl
 
 
-def _run_transform(
+async def _run_transform(
     query_name: str,
     variables: dict[str, Any],
     transform_func: Callable,
@@ -249,15 +266,15 @@ def _run_transform(
         raise typer.Abort()
 
     if asyncio.iscoroutinefunction(transform_func):
-        output = asyncio.run(transform_func(response))
+        output = await transform_func(response)
     else:
         output = transform_func(response)
     return output
 
 
 @app.command(name="render")
 @catch_exception(console=console)
-def render(
+async def render(
     transform_name: str = typer.Argument(default="", help="Name of the Python transformation", show_default=False),
     variables: Optional[list[str]] = typer.Argument(
         None, help="Variables to pass along with the query. Format key=value key=value."
@@ -289,7 +306,7 @@ def render(
     transform_func = functools.partial(render_jinja2_template, transform_config.template_path, variables_dict)
 
     # Query GQL and run the transform
-    result = _run_transform(
+    result = await _run_transform(
         query_name=transform_config.query,
         variables=variables_dict,
         transform_func=transform_func,

infrahub_sdk/pytest_plugin/items/jinja2_transform.py

@@ -1,22 +1,21 @@
 from __future__ import annotations
 
+import asyncio
 import difflib
+from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
 import jinja2
 import ujson
 from httpx import HTTPStatusError
-from rich.console import Console
-from rich.traceback import Traceback
 
-from ...jinja2 import identify_faulty_jinja_code
-from ..exceptions import Jinja2TransformError, Jinja2TransformUndefinedError, OutputMatchError
+from ...template import Jinja2Template
+from ...template.exceptions import JinjaTemplateError
+from ..exceptions import OutputMatchError
 from ..models import InfrahubInputOutputTest, InfrahubTestExpectedResult
 from .base import InfrahubItem
 
 if TYPE_CHECKING:
-    from pathlib import Path
-
     from pytest import ExceptionInfo
 
 
@@ -29,23 +28,15 @@ def get_jinja2_template(self) -> jinja2.Template:
         return self.get_jinja2_environment().get_template(str(self.resource_config.template_path))  # type: ignore[attr-defined]
 
     def render_jinja2_template(self, variables: dict[str, Any]) -> str | None:
+        jinja2_template = Jinja2Template(template_directory=Path(self.session.infrahub_config_path.parent))  # type: ignore[attr-defined]
+
         try:
-            return self.get_jinja2_template().render(**variables)
-        except jinja2.UndefinedError as exc:
-            traceback = Traceback(show_locals=False)
-            errors = identify_faulty_jinja_code(traceback=traceback)
-            console = Console()
-            with console.capture() as capture:
-                console.print(f"An error occurred while rendering Jinja2 transform:{self.name!r}\n", soft_wrap=True)
-                console.print(f"{exc.message}\n", soft_wrap=True)
-                for frame, syntax in errors:
-                    console.print(f"{frame.filename} on line {frame.lineno}\n", soft_wrap=True)
-                    console.print(syntax, soft_wrap=True)
-            str_output = capture.get()
+            return asyncio.run(
+                jinja2_template.render_from_file(template=Path(self.resource_config.template_path), variables=variables)  # type: ignore[attr-defined]
+            )
+        except JinjaTemplateError as exc:
             if self.test.expect == InfrahubTestExpectedResult.PASS:
-                raise Jinja2TransformUndefinedError(
-                    name=self.name, message=str_output, rtb=traceback, errors=errors
-                ) from exc
+                raise exc
             return None
 
     def get_result_differences(self, computed: Any) -> str | None:
@@ -99,7 +90,7 @@ def runtest(self) -> None:
             raise OutputMatchError(name=self.name, differences=differences)
 
     def repr_failure(self, excinfo: ExceptionInfo, style: str | None = None) -> str:
-        if isinstance(excinfo.value, (Jinja2TransformUndefinedError, Jinja2TransformError)):
+        if isinstance(excinfo.value, (JinjaTemplateError)):
             return excinfo.value.message
 
         return super().repr_failure(excinfo, style=style)

infrahub_sdk/template/__init__.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+import linecache
+from pathlib import Path
+from typing import Any, Callable, NoReturn
+
+import jinja2
+from jinja2.sandbox import SandboxedEnvironment
+from netutils.utils import jinja2_convenience_function
+from rich.syntax import Syntax
+from rich.traceback import Traceback
+
+from .constants import CURRATED_NETUTILS_FILTERS
+from .exceptions import (
+    JinjaTemplateError,
+    JinjaTemplateNotFoundError,
+    JinjaTemplateSyntaxError,
+    JinjaTemplateUndefinedError,
+)
+from .models import UndefinedJinja2Error
+
+netutils_filters = jinja2_convenience_function()
+
+
+class Jinja2Template:
+    FORBIDDEN_OPERATIONS: list[str] | None = None
+
+    def __init__(self, template_directory: Path | None = None) -> None:
+        self._template_directory = template_directory
+        self._filters: dict[str, Callable] = {}
+
+        self._filters.update(
+            {name: jinja_filter for name, jinja_filter in netutils_filters.items() if name in CURRATED_NETUTILS_FILTERS}
+        )
+
+    async def render_from_file(self, template: Path, variables: dict[str, Any]) -> str:
+        template_loader = jinja2.FileSystemLoader(searchpath=str(self._template_directory))
+        env = jinja2.Environment(loader=template_loader, trim_blocks=True, lstrip_blocks=True, enable_async=True)
+        env.filters.update(self._filters)
+
+        try:
+            jinja2_template = env.get_template(str(template))
+        except jinja2.TemplateSyntaxError as exc:
+            self._raise_template_syntax_error(error=exc)
+        except jinja2.TemplateNotFound as exc:
+            raise JinjaTemplateNotFoundError(message=exc.message, filename=str(exc.name))
+        return await self._render(template=jinja2_template, variables=variables)
+
+    async def render_from_string(self, template: str, variables: dict[str, Any]) -> str:
+        env = SandboxedEnvironment(enable_async=True, undefined=jinja2.StrictUndefined)
+        env.filters.update(self._filters)
+        try:
+            jinja2_template = env.from_string(template)
+        except jinja2.TemplateSyntaxError as exc:
+            self._raise_template_syntax_error(error=exc)
+
+        return await self._render(template=jinja2_template, variables=variables)
+
+    async def _render(self, template: jinja2.Template, variables: dict[str, Any]) -> str:
+        try:
+            output = await template.render_async(variables)
+        except jinja2.exceptions.TemplateNotFound as exc:
+            raise JinjaTemplateNotFoundError(message=exc.message, filename=str(exc.name), base_template=template.name)
+        except jinja2.TemplateSyntaxError as exc:
+            self._raise_template_syntax_error(error=exc)
+        except jinja2.UndefinedError as exc:
+            traceback = Traceback(show_locals=False)
+            errors = _identify_faulty_jinja_code(traceback=traceback)
+            raise JinjaTemplateUndefinedError(message=exc.message, errors=errors)
+        except Exception as exc:
+            if error_message := getattr(exc, "message", None):
+                message = error_message
+            else:
+                message = str(exc)
+            raise JinjaTemplateError(message=message or "Unknown template error")
+
+        return output
+
+    def _raise_template_syntax_error(self, error: jinja2.TemplateSyntaxError) -> NoReturn:
+        filename: str | None = None
+        if error.filename and self._template_directory:
+            filename = error.filename
+            if error.filename.startswith(str(self._template_directory)):
+                filename = error.filename[len(str(self._template_directory)) :]
+
+        raise JinjaTemplateSyntaxError(message=error.message, filename=filename, lineno=error.lineno)
+
+
+def _identify_faulty_jinja_code(traceback: Traceback, nbr_context_lines: int = 3) -> list[UndefinedJinja2Error]:
+    """This function identifies the faulty Jinja2 code and beautify it to provide meaningful information to the user.
+
+    We use the rich's Traceback to parse the complete stack trace and extract Frames for each exception found in the trace.
+    """
+    response = []
+
+    # Extract only the Jinja related exception
+    for frame in [frame for frame in traceback.trace.stacks[0].frames if not frame.filename.endswith(".py")]:
+        code = "".join(linecache.getlines(frame.filename))
+        if frame.filename == "<template>":
+            lexer_name = "text"
+        else:
+            lexer_name = Traceback._guess_lexer(frame.filename, code)
+        syntax = Syntax(
+            code,
+            lexer_name,
+            line_numbers=True,
+            line_range=(frame.lineno - nbr_context_lines, frame.lineno + nbr_context_lines),
+            highlight_lines={frame.lineno},
+            code_width=88,
+            theme=traceback.theme,
+            dedent=False,
+        )
+        response.append(UndefinedJinja2Error(frame=frame, syntax=syntax))
+
+    return response

infrahub_sdk/template/constants.py

@@ -0,0 +1,80 @@
+CURRATED_NETUTILS_FILTERS = [
+    "abbreviated_interface_name",
+    "abbreviated_interface_name_list",
+    "asn_to_int",
+    "bits_to_name",
+    "bytes_to_name",
+    "canonical_interface_name",
+    "canonical_interface_name_list",
+    "cidr_to_netmask",
+    "cidr_to_netmaskv6",
+    "clean_config",
+    "compare_version_loose",
+    "compare_version_strict",
+    "config_compliance",
+    "config_section_not_parsed",
+    "delimiter_change",
+    "diff_network_config",
+    "feature_compliance",
+    "find_unordered_cfg_lines",
+    "fqdn_to_ip",
+    "get_all_host",
+    "get_broadcast_address",
+    "get_first_usable",
+    "get_ips_sorted",
+    "get_nist_urls",
+    "get_nist_vendor_platform_urls",
+    "get_oui",
+    "get_peer_ip",
+    "get_range_ips",
+    "get_upgrade_path",
+    "get_usable_range",
+    "hash_data",
+    "int_to_asdot",
+    "interface_range_compress",
+    "interface_range_expansion",
+    "ip_addition",
+    "ip_subtract",
+    "ip_to_bin",
+    "ip_to_hex",
+    "ipaddress_address",
+    "ipaddress_interface",
+    "ipaddress_network",
+    "is_classful",
+    "is_fqdn_resolvable",
+    "is_ip",
+    "is_ip_range",
+    "is_ip_within",
+    "is_netmask",
+    "is_network",
+    "is_reversible_wildcardmask",
+    "is_valid_mac",
+    "longest_prefix_match",
+    "mac_normalize",
+    "mac_to_format",
+    "mac_to_int",
+    "mac_type",
+    "name_to_bits",
+    "name_to_bytes",
+    "name_to_name",
+    "netmask_to_cidr",
+    "netmask_to_wildcardmask",
+    "normalise_delimiter_caret_c",
+    "paloalto_panos_brace_to_set",
+    "paloalto_panos_clean_newlines",
+    "regex_findall",
+    "regex_match",
+    "regex_search",
+    "regex_split",
+    "regex_sub",
+    "sanitize_config",
+    "section_config",
+    "sort_interface_list",
+    "split_interface",
+    "uptime_seconds_to_string",
+    "uptime_string_to_seconds",
+    "version_metadata",
+    "vlanconfig_to_list",
+    "vlanlist_to_config",
+    "wildcardmask_to_netmask",
+]