idtracking.py 9.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286
  1. from jinja2.visitor import NodeVisitor
  2. from jinja2._compat import iteritems
  3. VAR_LOAD_PARAMETER = 'param'
  4. VAR_LOAD_RESOLVE = 'resolve'
  5. VAR_LOAD_ALIAS = 'alias'
  6. VAR_LOAD_UNDEFINED = 'undefined'
  7. def find_symbols(nodes, parent_symbols=None):
  8. sym = Symbols(parent=parent_symbols)
  9. visitor = FrameSymbolVisitor(sym)
  10. for node in nodes:
  11. visitor.visit(node)
  12. return sym
  13. def symbols_for_node(node, parent_symbols=None):
  14. sym = Symbols(parent=parent_symbols)
  15. sym.analyze_node(node)
  16. return sym
  17. class Symbols(object):
  18. def __init__(self, parent=None, level=None):
  19. if level is None:
  20. if parent is None:
  21. level = 0
  22. else:
  23. level = parent.level + 1
  24. self.level = level
  25. self.parent = parent
  26. self.refs = {}
  27. self.loads = {}
  28. self.stores = set()
  29. def analyze_node(self, node, **kwargs):
  30. visitor = RootVisitor(self)
  31. visitor.visit(node, **kwargs)
  32. def _define_ref(self, name, load=None):
  33. ident = 'l_%d_%s' % (self.level, name)
  34. self.refs[name] = ident
  35. if load is not None:
  36. self.loads[ident] = load
  37. return ident
  38. def find_load(self, target):
  39. if target in self.loads:
  40. return self.loads[target]
  41. if self.parent is not None:
  42. return self.parent.find_load(target)
  43. def find_ref(self, name):
  44. if name in self.refs:
  45. return self.refs[name]
  46. if self.parent is not None:
  47. return self.parent.find_ref(name)
  48. def ref(self, name):
  49. rv = self.find_ref(name)
  50. if rv is None:
  51. raise AssertionError('Tried to resolve a name to a reference that '
  52. 'was unknown to the frame (%r)' % name)
  53. return rv
  54. def copy(self):
  55. rv = object.__new__(self.__class__)
  56. rv.__dict__.update(self.__dict__)
  57. rv.refs = self.refs.copy()
  58. rv.loads = self.loads.copy()
  59. rv.stores = self.stores.copy()
  60. return rv
  61. def store(self, name):
  62. self.stores.add(name)
  63. # If we have not see the name referenced yet, we need to figure
  64. # out what to set it to.
  65. if name not in self.refs:
  66. # If there is a parent scope we check if the name has a
  67. # reference there. If it does it means we might have to alias
  68. # to a variable there.
  69. if self.parent is not None:
  70. outer_ref = self.parent.find_ref(name)
  71. if outer_ref is not None:
  72. self._define_ref(name, load=(VAR_LOAD_ALIAS, outer_ref))
  73. return
  74. # Otherwise we can just set it to undefined.
  75. self._define_ref(name, load=(VAR_LOAD_UNDEFINED, None))
  76. def declare_parameter(self, name):
  77. self.stores.add(name)
  78. return self._define_ref(name, load=(VAR_LOAD_PARAMETER, None))
  79. def load(self, name):
  80. target = self.find_ref(name)
  81. if target is None:
  82. self._define_ref(name, load=(VAR_LOAD_RESOLVE, name))
  83. def branch_update(self, branch_symbols):
  84. stores = {}
  85. for branch in branch_symbols:
  86. for target in branch.stores:
  87. if target in self.stores:
  88. continue
  89. stores[target] = stores.get(target, 0) + 1
  90. for sym in branch_symbols:
  91. self.refs.update(sym.refs)
  92. self.loads.update(sym.loads)
  93. self.stores.update(sym.stores)
  94. for name, branch_count in iteritems(stores):
  95. if branch_count == len(branch_symbols):
  96. continue
  97. target = self.find_ref(name)
  98. assert target is not None, 'should not happen'
  99. if self.parent is not None:
  100. outer_target = self.parent.find_ref(name)
  101. if outer_target is not None:
  102. self.loads[target] = (VAR_LOAD_ALIAS, outer_target)
  103. continue
  104. self.loads[target] = (VAR_LOAD_RESOLVE, name)
  105. def dump_stores(self):
  106. rv = {}
  107. node = self
  108. while node is not None:
  109. for name in node.stores:
  110. if name not in rv:
  111. rv[name] = self.find_ref(name)
  112. node = node.parent
  113. return rv
  114. def dump_param_targets(self):
  115. rv = set()
  116. node = self
  117. while node is not None:
  118. for target, (instr, _) in iteritems(self.loads):
  119. if instr == VAR_LOAD_PARAMETER:
  120. rv.add(target)
  121. node = node.parent
  122. return rv
  123. class RootVisitor(NodeVisitor):
  124. def __init__(self, symbols):
  125. self.sym_visitor = FrameSymbolVisitor(symbols)
  126. def _simple_visit(self, node, **kwargs):
  127. for child in node.iter_child_nodes():
  128. self.sym_visitor.visit(child)
  129. visit_Template = visit_Block = visit_Macro = visit_FilterBlock = \
  130. visit_Scope = visit_If = visit_ScopedEvalContextModifier = \
  131. _simple_visit
  132. def visit_AssignBlock(self, node, **kwargs):
  133. for child in node.body:
  134. self.sym_visitor.visit(child)
  135. def visit_CallBlock(self, node, **kwargs):
  136. for child in node.iter_child_nodes(exclude=('call',)):
  137. self.sym_visitor.visit(child)
  138. def visit_OverlayScope(self, node, **kwargs):
  139. for child in node.body:
  140. self.sym_visitor.visit(child)
  141. def visit_For(self, node, for_branch='body', **kwargs):
  142. if for_branch == 'body':
  143. self.sym_visitor.visit(node.target, store_as_param=True)
  144. branch = node.body
  145. elif for_branch == 'else':
  146. branch = node.else_
  147. elif for_branch == 'test':
  148. self.sym_visitor.visit(node.target, store_as_param=True)
  149. if node.test is not None:
  150. self.sym_visitor.visit(node.test)
  151. return
  152. else:
  153. raise RuntimeError('Unknown for branch')
  154. for item in branch or ():
  155. self.sym_visitor.visit(item)
  156. def visit_With(self, node, **kwargs):
  157. for target in node.targets:
  158. self.sym_visitor.visit(target)
  159. for child in node.body:
  160. self.sym_visitor.visit(child)
  161. def generic_visit(self, node, *args, **kwargs):
  162. raise NotImplementedError('Cannot find symbols for %r' %
  163. node.__class__.__name__)
  164. class FrameSymbolVisitor(NodeVisitor):
  165. """A visitor for `Frame.inspect`."""
  166. def __init__(self, symbols):
  167. self.symbols = symbols
  168. def visit_Name(self, node, store_as_param=False, **kwargs):
  169. """All assignments to names go through this function."""
  170. if store_as_param or node.ctx == 'param':
  171. self.symbols.declare_parameter(node.name)
  172. elif node.ctx == 'store':
  173. self.symbols.store(node.name)
  174. elif node.ctx == 'load':
  175. self.symbols.load(node.name)
  176. def visit_NSRef(self, node, **kwargs):
  177. self.symbols.load(node.name)
  178. def visit_If(self, node, **kwargs):
  179. self.visit(node.test, **kwargs)
  180. original_symbols = self.symbols
  181. def inner_visit(nodes):
  182. self.symbols = rv = original_symbols.copy()
  183. for subnode in nodes:
  184. self.visit(subnode, **kwargs)
  185. self.symbols = original_symbols
  186. return rv
  187. body_symbols = inner_visit(node.body)
  188. elif_symbols = inner_visit(node.elif_)
  189. else_symbols = inner_visit(node.else_ or ())
  190. self.symbols.branch_update([body_symbols, elif_symbols, else_symbols])
  191. def visit_Macro(self, node, **kwargs):
  192. self.symbols.store(node.name)
  193. def visit_Import(self, node, **kwargs):
  194. self.generic_visit(node, **kwargs)
  195. self.symbols.store(node.target)
  196. def visit_FromImport(self, node, **kwargs):
  197. self.generic_visit(node, **kwargs)
  198. for name in node.names:
  199. if isinstance(name, tuple):
  200. self.symbols.store(name[1])
  201. else:
  202. self.symbols.store(name)
  203. def visit_Assign(self, node, **kwargs):
  204. """Visit assignments in the correct order."""
  205. self.visit(node.node, **kwargs)
  206. self.visit(node.target, **kwargs)
  207. def visit_For(self, node, **kwargs):
  208. """Visiting stops at for blocks. However the block sequence
  209. is visited as part of the outer scope.
  210. """
  211. self.visit(node.iter, **kwargs)
  212. def visit_CallBlock(self, node, **kwargs):
  213. self.visit(node.call, **kwargs)
  214. def visit_FilterBlock(self, node, **kwargs):
  215. self.visit(node.filter, **kwargs)
  216. def visit_With(self, node, **kwargs):
  217. for target in node.values:
  218. self.visit(target)
  219. def visit_AssignBlock(self, node, **kwargs):
  220. """Stop visiting at block assigns."""
  221. self.visit(node.target, **kwargs)
  222. def visit_Scope(self, node, **kwargs):
  223. """Stop visiting at scopes."""
  224. def visit_Block(self, node, **kwargs):
  225. """Stop visiting at blocks."""
  226. def visit_OverlayScope(self, node, **kwargs):
  227. """Do not visit into overlay scopes."""