Skip to content

Instantly share code, notes, and snippets.

Last active July 28, 2023 21:27
Show Gist options
  • Save wolflu05/3919730ac2e305a3ba80c615c0e15bed to your computer and use it in GitHub Desktop.
Save wolflu05/3919730ac2e305a3ba80c615c0e15bed to your computer and use it in GitHub Desktop.
Get all variables accessed in a jinja2 template
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:
# this code is inspired by:
# the already build-in function jinja2.meta.find_undeclared_variables(...) which does not output nested keys
code_gen = TrackingCodeGenerator(self.ast.environment)
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.stack = []
def parse_jinja_variable(self, variable: nodes.Node, suffix=""):
if type(variable) is nodes.Name:
return self.join_keys(, 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
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)
def visit_Getattr(self, node: nodes.Getattr, frame: Frame):
super().visit_Getattr(node, frame)
def visit_Getitem(self, node: nodes.Getitem, frame: Frame):
super().visit_Getitem(node, frame)
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.stack = []
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment