idtracking.py 9.0 KB

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