summaryrefslogtreecommitdiffhomepage
path: root/scripts/jinja2/idtracking.py
diff options
context:
space:
mode:
Diffstat (limited to 'scripts/jinja2/idtracking.py')
-rw-r--r--scripts/jinja2/idtracking.py318
1 files changed, 318 insertions, 0 deletions
diff --git a/scripts/jinja2/idtracking.py b/scripts/jinja2/idtracking.py
new file mode 100644
index 0000000..38c525e
--- /dev/null
+++ b/scripts/jinja2/idtracking.py
@@ -0,0 +1,318 @@
+import typing as t
+
+from . import nodes
+from .visitor import NodeVisitor
+
+VAR_LOAD_PARAMETER = "param"
+VAR_LOAD_RESOLVE = "resolve"
+VAR_LOAD_ALIAS = "alias"
+VAR_LOAD_UNDEFINED = "undefined"
+
+
+def find_symbols(
+ nodes: t.Iterable[nodes.Node], parent_symbols: t.Optional["Symbols"] = None
+) -> "Symbols":
+ sym = Symbols(parent=parent_symbols)
+ visitor = FrameSymbolVisitor(sym)
+ for node in nodes:
+ visitor.visit(node)
+ return sym
+
+
+def symbols_for_node(
+ node: nodes.Node, parent_symbols: t.Optional["Symbols"] = None
+) -> "Symbols":
+ sym = Symbols(parent=parent_symbols)
+ sym.analyze_node(node)
+ return sym
+
+
+class Symbols:
+ def __init__(
+ self, parent: t.Optional["Symbols"] = None, level: t.Optional[int] = None
+ ) -> None:
+ if level is None:
+ if parent is None:
+ level = 0
+ else:
+ level = parent.level + 1
+
+ self.level: int = level
+ self.parent = parent
+ self.refs: t.Dict[str, str] = {}
+ self.loads: t.Dict[str, t.Any] = {}
+ self.stores: t.Set[str] = set()
+
+ def analyze_node(self, node: nodes.Node, **kwargs: t.Any) -> None:
+ visitor = RootVisitor(self)
+ visitor.visit(node, **kwargs)
+
+ def _define_ref(
+ self, name: str, load: t.Optional[t.Tuple[str, t.Optional[str]]] = None
+ ) -> str:
+ ident = f"l_{self.level}_{name}"
+ self.refs[name] = ident
+ if load is not None:
+ self.loads[ident] = load
+ return ident
+
+ def find_load(self, target: str) -> t.Optional[t.Any]:
+ if target in self.loads:
+ return self.loads[target]
+
+ if self.parent is not None:
+ return self.parent.find_load(target)
+
+ return None
+
+ def find_ref(self, name: str) -> t.Optional[str]:
+ if name in self.refs:
+ return self.refs[name]
+
+ if self.parent is not None:
+ return self.parent.find_ref(name)
+
+ return None
+
+ def ref(self, name: str) -> str:
+ rv = self.find_ref(name)
+ if rv is None:
+ raise AssertionError(
+ "Tried to resolve a name to a reference that was"
+ f" unknown to the frame ({name!r})"
+ )
+ return rv
+
+ def copy(self) -> "Symbols":
+ rv = t.cast(Symbols, object.__new__(self.__class__))
+ rv.__dict__.update(self.__dict__)
+ rv.refs = self.refs.copy()
+ rv.loads = self.loads.copy()
+ rv.stores = self.stores.copy()
+ return rv
+
+ def store(self, name: str) -> None:
+ self.stores.add(name)
+
+ # If we have not see the name referenced yet, we need to figure
+ # out what to set it to.
+ if name not in self.refs:
+ # If there is a parent scope we check if the name has a
+ # reference there. If it does it means we might have to alias
+ # to a variable there.
+ if self.parent is not None:
+ outer_ref = self.parent.find_ref(name)
+ if outer_ref is not None:
+ self._define_ref(name, load=(VAR_LOAD_ALIAS, outer_ref))
+ return
+
+ # Otherwise we can just set it to undefined.
+ self._define_ref(name, load=(VAR_LOAD_UNDEFINED, None))
+
+ def declare_parameter(self, name: str) -> str:
+ self.stores.add(name)
+ return self._define_ref(name, load=(VAR_LOAD_PARAMETER, None))
+
+ def load(self, name: str) -> None:
+ if self.find_ref(name) is None:
+ self._define_ref(name, load=(VAR_LOAD_RESOLVE, name))
+
+ def branch_update(self, branch_symbols: t.Sequence["Symbols"]) -> None:
+ stores: t.Dict[str, int] = {}
+ for branch in branch_symbols:
+ for target in branch.stores:
+ if target in self.stores:
+ continue
+ stores[target] = stores.get(target, 0) + 1
+
+ for sym in branch_symbols:
+ self.refs.update(sym.refs)
+ self.loads.update(sym.loads)
+ self.stores.update(sym.stores)
+
+ for name, branch_count in stores.items():
+ if branch_count == len(branch_symbols):
+ continue
+
+ target = self.find_ref(name) # type: ignore
+ assert target is not None, "should not happen"
+
+ if self.parent is not None:
+ outer_target = self.parent.find_ref(name)
+ if outer_target is not None:
+ self.loads[target] = (VAR_LOAD_ALIAS, outer_target)
+ continue
+ self.loads[target] = (VAR_LOAD_RESOLVE, name)
+
+ def dump_stores(self) -> t.Dict[str, str]:
+ rv: t.Dict[str, str] = {}
+ node: t.Optional["Symbols"] = self
+
+ while node is not None:
+ for name in sorted(node.stores):
+ if name not in rv:
+ rv[name] = self.find_ref(name) # type: ignore
+
+ node = node.parent
+
+ return rv
+
+ def dump_param_targets(self) -> t.Set[str]:
+ rv = set()
+ node: t.Optional["Symbols"] = self
+
+ while node is not None:
+ for target, (instr, _) in self.loads.items():
+ if instr == VAR_LOAD_PARAMETER:
+ rv.add(target)
+
+ node = node.parent
+
+ return rv
+
+
+class RootVisitor(NodeVisitor):
+ def __init__(self, symbols: "Symbols") -> None:
+ self.sym_visitor = FrameSymbolVisitor(symbols)
+
+ def _simple_visit(self, node: nodes.Node, **kwargs: t.Any) -> None:
+ for child in node.iter_child_nodes():
+ self.sym_visitor.visit(child)
+
+ visit_Template = _simple_visit
+ visit_Block = _simple_visit
+ visit_Macro = _simple_visit
+ visit_FilterBlock = _simple_visit
+ visit_Scope = _simple_visit
+ visit_If = _simple_visit
+ visit_ScopedEvalContextModifier = _simple_visit
+
+ def visit_AssignBlock(self, node: nodes.AssignBlock, **kwargs: t.Any) -> None:
+ for child in node.body:
+ self.sym_visitor.visit(child)
+
+ def visit_CallBlock(self, node: nodes.CallBlock, **kwargs: t.Any) -> None:
+ for child in node.iter_child_nodes(exclude=("call",)):
+ self.sym_visitor.visit(child)
+
+ def visit_OverlayScope(self, node: nodes.OverlayScope, **kwargs: t.Any) -> None:
+ for child in node.body:
+ self.sym_visitor.visit(child)
+
+ def visit_For(
+ self, node: nodes.For, for_branch: str = "body", **kwargs: t.Any
+ ) -> None:
+ if for_branch == "body":
+ self.sym_visitor.visit(node.target, store_as_param=True)
+ branch = node.body
+ elif for_branch == "else":
+ branch = node.else_
+ elif for_branch == "test":
+ self.sym_visitor.visit(node.target, store_as_param=True)
+ if node.test is not None:
+ self.sym_visitor.visit(node.test)
+ return
+ else:
+ raise RuntimeError("Unknown for branch")
+
+ if branch:
+ for item in branch:
+ self.sym_visitor.visit(item)
+
+ def visit_With(self, node: nodes.With, **kwargs: t.Any) -> None:
+ for target in node.targets:
+ self.sym_visitor.visit(target)
+ for child in node.body:
+ self.sym_visitor.visit(child)
+
+ def generic_visit(self, node: nodes.Node, *args: t.Any, **kwargs: t.Any) -> None:
+ raise NotImplementedError(f"Cannot find symbols for {type(node).__name__!r}")
+
+
+class FrameSymbolVisitor(NodeVisitor):
+ """A visitor for `Frame.inspect`."""
+
+ def __init__(self, symbols: "Symbols") -> None:
+ self.symbols = symbols
+
+ def visit_Name(
+ self, node: nodes.Name, store_as_param: bool = False, **kwargs: t.Any
+ ) -> None:
+ """All assignments to names go through this function."""
+ if store_as_param or node.ctx == "param":
+ self.symbols.declare_parameter(node.name)
+ elif node.ctx == "store":
+ self.symbols.store(node.name)
+ elif node.ctx == "load":
+ self.symbols.load(node.name)
+
+ def visit_NSRef(self, node: nodes.NSRef, **kwargs: t.Any) -> None:
+ self.symbols.load(node.name)
+
+ def visit_If(self, node: nodes.If, **kwargs: t.Any) -> None:
+ self.visit(node.test, **kwargs)
+ original_symbols = self.symbols
+
+ def inner_visit(nodes: t.Iterable[nodes.Node]) -> "Symbols":
+ self.symbols = rv = original_symbols.copy()
+
+ for subnode in nodes:
+ self.visit(subnode, **kwargs)
+
+ self.symbols = original_symbols
+ return rv
+
+ body_symbols = inner_visit(node.body)
+ elif_symbols = inner_visit(node.elif_)
+ else_symbols = inner_visit(node.else_ or ())
+ self.symbols.branch_update([body_symbols, elif_symbols, else_symbols])
+
+ def visit_Macro(self, node: nodes.Macro, **kwargs: t.Any) -> None:
+ self.symbols.store(node.name)
+
+ def visit_Import(self, node: nodes.Import, **kwargs: t.Any) -> None:
+ self.generic_visit(node, **kwargs)
+ self.symbols.store(node.target)
+
+ def visit_FromImport(self, node: nodes.FromImport, **kwargs: t.Any) -> None:
+ self.generic_visit(node, **kwargs)
+
+ for name in node.names:
+ if isinstance(name, tuple):
+ self.symbols.store(name[1])
+ else:
+ self.symbols.store(name)
+
+ def visit_Assign(self, node: nodes.Assign, **kwargs: t.Any) -> None:
+ """Visit assignments in the correct order."""
+ self.visit(node.node, **kwargs)
+ self.visit(node.target, **kwargs)
+
+ def visit_For(self, node: nodes.For, **kwargs: t.Any) -> None:
+ """Visiting stops at for blocks. However the block sequence
+ is visited as part of the outer scope.
+ """
+ self.visit(node.iter, **kwargs)
+
+ def visit_CallBlock(self, node: nodes.CallBlock, **kwargs: t.Any) -> None:
+ self.visit(node.call, **kwargs)
+
+ def visit_FilterBlock(self, node: nodes.FilterBlock, **kwargs: t.Any) -> None:
+ self.visit(node.filter, **kwargs)
+
+ def visit_With(self, node: nodes.With, **kwargs: t.Any) -> None:
+ for target in node.values:
+ self.visit(target)
+
+ def visit_AssignBlock(self, node: nodes.AssignBlock, **kwargs: t.Any) -> None:
+ """Stop visiting at block assigns."""
+ self.visit(node.target, **kwargs)
+
+ def visit_Scope(self, node: nodes.Scope, **kwargs: t.Any) -> None:
+ """Stop visiting at scopes."""
+
+ def visit_Block(self, node: nodes.Block, **kwargs: t.Any) -> None:
+ """Stop visiting at blocks."""
+
+ def visit_OverlayScope(self, node: nodes.OverlayScope, **kwargs: t.Any) -> None:
+ """Do not visit into overlay scopes."""