Get all variables accessed in a jinja2 template

jinja2_TrackingCodeGenerator.py

from jinja2 import Environment, compiler, nodes
from jinja2.compiler import Frame
import re

env = Environment()


class Template:
    def __init__(self, template_str: str) -> None:
        # escape block statements
        template_str = re.sub(r"(\{\%|\%\})", "{{\"\\1\"}}", template_str)

        self.template_str = template_str
        self.ast = None

    def validate(self):
        self.ast = env.parse(self.template_str)

    def used_variables(self):
        if self.ast is None:
            self.validate()
            
        # this code is inspired by:
        # the already build-in function jinja2.meta.find_undeclared_variables(...) which does not output nested keys
        # https://stackoverflow.com/a/71664198
        # https://stackoverflow.com/a/8284419
        code_gen = TrackingCodeGenerator(self.ast.environment)
        code_gen.visit(self.ast)
        return code_gen.variables

    def compile(self):
        return env.from_string(self.template_str)


class TrackingCodeGenerator(compiler.CodeGenerator):
    """We abuse the code generator for introspection."""

    def __init__(self, environment: "Environment") -> None:
        super().__init__(environment, "<introspection>", "<introspection>")
        self.variables = list[str]()
        self.stack = list[nodes.Node]()

    def write(self, x: str) -> None:
        """Don't write."""

    def capture_node(self, node: nodes.Node):
        # if the current node's parent differs from the last element in stack,
        # we process a new variable, so we clean up the stack first and parse the variable
        if len(self.stack) > 0 and (not hasattr(node, "node") or node.node != self.stack[-1]):
            if parsed_variable := self.parse_jinja_variable(self.stack[-1]):
                self.variables.append(parsed_variable)
            self.stack = []

        self.stack.append(node)

    def parse_jinja_variable(self, variable: nodes.Node, suffix=""):
        if type(variable) is nodes.Name:
            return self.join_keys(variable.name, suffix)
        elif type(variable) is nodes.Getattr:
            return self.parse_jinja_variable(variable.node, self.join_keys(variable.attr, suffix))
        elif type(variable) is nodes.Getitem:
            return self.parse_jinja_variable(variable.node, self.join_keys(str(variable.arg.value), suffix))
        return variable

    @staticmethod
    def join_keys(*keys: list[str]):
        return ".".join(k for k in keys if k)

    # --- track visiting of names, getattr, getitem and cleanup on frame leave
    def visit_Name(self, node: nodes.Name, frame: Frame):
        super().visit_Name(node, frame)
        self.capture_node(node)

    def visit_Getattr(self, node: nodes.Getattr, frame: Frame):
        super().visit_Getattr(node, frame)
        self.capture_node(node)

    def visit_Getitem(self, node: nodes.Getitem, frame: Frame):
        super().visit_Getitem(node, frame)
        self.capture_node(node)

    def leave_frame(self, frame: Frame, with_python_scope: bool = False):
        super().leave_frame(frame, with_python_scope)

        # clean up stack before leaving frame so that the last variable is added to the variables list
        if len(self.stack) > 0:
            if parsed_variable := self.parse_jinja_variable(self.stack[-1]):
                self.variables.append(parsed_variable)
            self.stack = []