1import typing as t
2
3from . import nodes
4from .visitor import NodeVisitor
5
6if t.TYPE_CHECKING:
7 import typing_extensions as te
8
9VAR_LOAD_PARAMETER = "param"
10VAR_LOAD_RESOLVE = "resolve"
11VAR_LOAD_ALIAS = "alias"
12VAR_LOAD_UNDEFINED = "undefined"
13
14
15def find_symbols(
16 nodes: t.Iterable[nodes.Node], parent_symbols: t.Optional["Symbols"] = None
17) -> "Symbols":
18 sym = Symbols(parent=parent_symbols)
19 visitor = FrameSymbolVisitor(sym)
20 for node in nodes:
21 visitor.visit(node)
22 return sym
23
24
25def symbols_for_node(
26 node: nodes.Node, parent_symbols: t.Optional["Symbols"] = None
27) -> "Symbols":
28 sym = Symbols(parent=parent_symbols)
29 sym.analyze_node(node)
30 return sym
31
32
33class Symbols:
34 def __init__(
35 self, parent: t.Optional["Symbols"] = None, level: t.Optional[int] = None
36 ) -> None:
37 if level is None:
38 if parent is None:
39 level = 0
40 else:
41 level = parent.level + 1
42
43 self.level: int = level
44 self.parent = parent
45 self.refs: t.Dict[str, str] = {}
46 self.loads: t.Dict[str, t.Any] = {}
47 self.stores: t.Set[str] = set()
48
49 def analyze_node(self, node: nodes.Node, **kwargs: t.Any) -> None:
50 visitor = RootVisitor(self)
51 visitor.visit(node, **kwargs)
52
53 def _define_ref(
54 self, name: str, load: t.Optional[t.Tuple[str, t.Optional[str]]] = None
55 ) -> str:
56 ident = f"l_{self.level}_{name}"
57 self.refs[name] = ident
58 if load is not None:
59 self.loads[ident] = load
60 return ident
61
62 def find_load(self, target: str) -> t.Optional[t.Any]:
63 if target in self.loads:
64 return self.loads[target]
65
66 if self.parent is not None:
67 return self.parent.find_load(target)
68
69 return None
70
71 def find_ref(self, name: str) -> t.Optional[str]:
72 if name in self.refs:
73 return self.refs[name]
74
75 if self.parent is not None:
76 return self.parent.find_ref(name)
77
78 return None
79
80 def ref(self, name: str) -> str:
81 rv = self.find_ref(name)
82 if rv is None:
83 raise AssertionError(
84 "Tried to resolve a name to a reference that was"
85 f" unknown to the frame ({name!r})"
86 )
87 return rv
88
89 def copy(self) -> "te.Self":
90 rv = object.__new__(self.__class__)
91 rv.__dict__.update(self.__dict__)
92 rv.refs = self.refs.copy()
93 rv.loads = self.loads.copy()
94 rv.stores = self.stores.copy()
95 return rv
96
97 def store(self, name: str) -> None:
98 self.stores.add(name)
99
100 # If we have not see the name referenced yet, we need to figure
101 # out what to set it to.
102 if name not in self.refs:
103 # If there is a parent scope we check if the name has a
104 # reference there. If it does it means we might have to alias
105 # to a variable there.
106 if self.parent is not None:
107 outer_ref = self.parent.find_ref(name)
108 if outer_ref is not None:
109 self._define_ref(name, load=(VAR_LOAD_ALIAS, outer_ref))
110 return
111
112 # Otherwise we can just set it to undefined.
113 self._define_ref(name, load=(VAR_LOAD_UNDEFINED, None))
114
115 def declare_parameter(self, name: str) -> str:
116 self.stores.add(name)
117 return self._define_ref(name, load=(VAR_LOAD_PARAMETER, None))
118
119 def load(self, name: str) -> None:
120 if self.find_ref(name) is None:
121 self._define_ref(name, load=(VAR_LOAD_RESOLVE, name))
122
123 def branch_update(self, branch_symbols: t.Sequence["Symbols"]) -> None:
124 stores: t.Set[str] = set()
125
126 for branch in branch_symbols:
127 stores.update(branch.stores)
128
129 stores.difference_update(self.stores)
130
131 for sym in branch_symbols:
132 self.refs.update(sym.refs)
133 self.loads.update(sym.loads)
134 self.stores.update(sym.stores)
135
136 for name in stores:
137 target = self.find_ref(name)
138 assert target is not None, "should not happen"
139
140 if self.parent is not None:
141 outer_target = self.parent.find_ref(name)
142 if outer_target is not None:
143 self.loads[target] = (VAR_LOAD_ALIAS, outer_target)
144 continue
145 self.loads[target] = (VAR_LOAD_RESOLVE, name)
146
147 def dump_stores(self) -> t.Dict[str, str]:
148 rv: t.Dict[str, str] = {}
149 node: t.Optional[Symbols] = self
150
151 while node is not None:
152 for name in sorted(node.stores):
153 if name not in rv:
154 rv[name] = self.find_ref(name) # type: ignore
155
156 node = node.parent
157
158 return rv
159
160 def dump_param_targets(self) -> t.Set[str]:
161 rv = set()
162 node: t.Optional[Symbols] = self
163
164 while node is not None:
165 for target, (instr, _) in self.loads.items():
166 if instr == VAR_LOAD_PARAMETER:
167 rv.add(target)
168
169 node = node.parent
170
171 return rv
172
173
174class RootVisitor(NodeVisitor):
175 def __init__(self, symbols: "Symbols") -> None:
176 self.sym_visitor = FrameSymbolVisitor(symbols)
177
178 def _simple_visit(self, node: nodes.Node, **kwargs: t.Any) -> None:
179 for child in node.iter_child_nodes():
180 self.sym_visitor.visit(child)
181
182 visit_Template = _simple_visit
183 visit_Block = _simple_visit
184 visit_Macro = _simple_visit
185 visit_FilterBlock = _simple_visit
186 visit_Scope = _simple_visit
187 visit_If = _simple_visit
188 visit_ScopedEvalContextModifier = _simple_visit
189
190 def visit_AssignBlock(self, node: nodes.AssignBlock, **kwargs: t.Any) -> None:
191 for child in node.body:
192 self.sym_visitor.visit(child)
193
194 def visit_CallBlock(self, node: nodes.CallBlock, **kwargs: t.Any) -> None:
195 for child in node.iter_child_nodes(exclude=("call",)):
196 self.sym_visitor.visit(child)
197
198 def visit_OverlayScope(self, node: nodes.OverlayScope, **kwargs: t.Any) -> None:
199 for child in node.body:
200 self.sym_visitor.visit(child)
201
202 def visit_For(
203 self, node: nodes.For, for_branch: str = "body", **kwargs: t.Any
204 ) -> None:
205 if for_branch == "body":
206 self.sym_visitor.visit(node.target, store_as_param=True)
207 branch = node.body
208 elif for_branch == "else":
209 branch = node.else_
210 elif for_branch == "test":
211 self.sym_visitor.visit(node.target, store_as_param=True)
212 if node.test is not None:
213 self.sym_visitor.visit(node.test)
214 return
215 else:
216 raise RuntimeError("Unknown for branch")
217
218 if branch:
219 for item in branch:
220 self.sym_visitor.visit(item)
221
222 def visit_With(self, node: nodes.With, **kwargs: t.Any) -> None:
223 for target in node.targets:
224 self.sym_visitor.visit(target)
225 for child in node.body:
226 self.sym_visitor.visit(child)
227
228 def generic_visit(self, node: nodes.Node, *args: t.Any, **kwargs: t.Any) -> None:
229 raise NotImplementedError(f"Cannot find symbols for {type(node).__name__!r}")
230
231
232class FrameSymbolVisitor(NodeVisitor):
233 """A visitor for `Frame.inspect`."""
234
235 def __init__(self, symbols: "Symbols") -> None:
236 self.symbols = symbols
237
238 def visit_Name(
239 self, node: nodes.Name, store_as_param: bool = False, **kwargs: t.Any
240 ) -> None:
241 """All assignments to names go through this function."""
242 if store_as_param or node.ctx == "param":
243 self.symbols.declare_parameter(node.name)
244 elif node.ctx == "store":
245 self.symbols.store(node.name)
246 elif node.ctx == "load":
247 self.symbols.load(node.name)
248
249 def visit_NSRef(self, node: nodes.NSRef, **kwargs: t.Any) -> None:
250 self.symbols.load(node.name)
251
252 def visit_If(self, node: nodes.If, **kwargs: t.Any) -> None:
253 self.visit(node.test, **kwargs)
254 original_symbols = self.symbols
255
256 def inner_visit(nodes: t.Iterable[nodes.Node]) -> "Symbols":
257 self.symbols = rv = original_symbols.copy()
258
259 for subnode in nodes:
260 self.visit(subnode, **kwargs)
261
262 self.symbols = original_symbols
263 return rv
264
265 body_symbols = inner_visit(node.body)
266 elif_symbols = inner_visit(node.elif_)
267 else_symbols = inner_visit(node.else_ or ())
268 self.symbols.branch_update([body_symbols, elif_symbols, else_symbols])
269
270 def visit_Macro(self, node: nodes.Macro, **kwargs: t.Any) -> None:
271 self.symbols.store(node.name)
272
273 def visit_Import(self, node: nodes.Import, **kwargs: t.Any) -> None:
274 self.generic_visit(node, **kwargs)
275 self.symbols.store(node.target)
276
277 def visit_FromImport(self, node: nodes.FromImport, **kwargs: t.Any) -> None:
278 self.generic_visit(node, **kwargs)
279
280 for name in node.names:
281 if isinstance(name, tuple):
282 self.symbols.store(name[1])
283 else:
284 self.symbols.store(name)
285
286 def visit_Assign(self, node: nodes.Assign, **kwargs: t.Any) -> None:
287 """Visit assignments in the correct order."""
288 self.visit(node.node, **kwargs)
289 self.visit(node.target, **kwargs)
290
291 def visit_For(self, node: nodes.For, **kwargs: t.Any) -> None:
292 """Visiting stops at for blocks. However the block sequence
293 is visited as part of the outer scope.
294 """
295 self.visit(node.iter, **kwargs)
296
297 def visit_CallBlock(self, node: nodes.CallBlock, **kwargs: t.Any) -> None:
298 self.visit(node.call, **kwargs)
299
300 def visit_FilterBlock(self, node: nodes.FilterBlock, **kwargs: t.Any) -> None:
301 self.visit(node.filter, **kwargs)
302
303 def visit_With(self, node: nodes.With, **kwargs: t.Any) -> None:
304 for target in node.values:
305 self.visit(target)
306
307 def visit_AssignBlock(self, node: nodes.AssignBlock, **kwargs: t.Any) -> None:
308 """Stop visiting at block assigns."""
309 self.visit(node.target, **kwargs)
310
311 def visit_Scope(self, node: nodes.Scope, **kwargs: t.Any) -> None:
312 """Stop visiting at scopes."""
313
314 def visit_Block(self, node: nodes.Block, **kwargs: t.Any) -> None:
315 """Stop visiting at blocks."""
316
317 def visit_OverlayScope(self, node: nodes.OverlayScope, **kwargs: t.Any) -> None:
318 """Do not visit into overlay scopes."""