tree_walk.py 5.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179
  1. # -*- coding: utf-8 -*-
  2. """
  3. Part of the astor library for Python AST manipulation.
  4. License: 3-clause BSD
  5. Copyright 2012 (c) Patrick Maupin
  6. Copyright 2013 (c) Berker Peksag
  7. This file contains a TreeWalk class that views a node tree
  8. as a unified whole and allows several modes of traversal.
  9. """
  10. from .node_util import iter_node
  11. class MetaFlatten(type):
  12. """This metaclass is used to flatten classes to remove
  13. class hierarchy.
  14. This makes it easier to manipulate classes (find
  15. attributes in a single dict, etc.)
  16. """
  17. def __new__(clstype, name, bases, clsdict):
  18. newbases = (object,)
  19. newdict = {}
  20. for base in reversed(bases):
  21. if base not in newbases:
  22. newdict.update(vars(base))
  23. newdict.update(clsdict)
  24. # These are class-bound, we should let Python recreate them.
  25. newdict.pop('__dict__', None)
  26. newdict.pop('__weakref__', None)
  27. # Delegate the real work to type
  28. return type.__new__(clstype, name, newbases, newdict)
  29. MetaFlatten = MetaFlatten('MetaFlatten', (object,), {})
  30. class TreeWalk(MetaFlatten):
  31. """The TreeWalk class can be used as a superclass in order
  32. to walk an AST or similar tree.
  33. Unlike other treewalkers, this class can walk a tree either
  34. recursively or non-recursively. Subclasses can define
  35. methods with the following signatures::
  36. def pre_xxx(self):
  37. pass
  38. def post_xxx(self):
  39. pass
  40. def init_xxx(self):
  41. pass
  42. Where 'xxx' is one of:
  43. - A class name
  44. - An attribute member name concatenated with '_name'
  45. For example, 'pre_targets_name' will process nodes
  46. that are referenced by the name 'targets' in their
  47. parent's node.
  48. - An attribute member name concatenated with '_item'
  49. For example, 'pre_targets_item' will process nodes
  50. that are in a list that is the targets attribute
  51. of some node.
  52. pre_xxx will process a node before processing any of its subnodes.
  53. if the return value from pre_xxx evalates to true, then walk
  54. will not process any of the subnodes. Those can be manually
  55. processed, if desired, by calling self.walk(node) on the subnodes
  56. before returning True.
  57. post_xxx will process a node after processing all its subnodes.
  58. init_xxx methods can decorate the class instance with subclass-specific
  59. information. A single init_whatever method could be written, but to
  60. make it easy to keep initialization with use, any number of init_xxx
  61. methods can be written. They will be called in alphabetical order.
  62. """
  63. def __init__(self, node=None):
  64. self.nodestack = []
  65. self.setup()
  66. if node is not None:
  67. self.walk(node)
  68. def setup(self):
  69. """All the node-specific handlers are setup at
  70. object initialization time.
  71. """
  72. self.pre_handlers = pre_handlers = {}
  73. self.post_handlers = post_handlers = {}
  74. for name in sorted(vars(type(self))):
  75. if name.startswith('init_'):
  76. getattr(self, name)()
  77. elif name.startswith('pre_'):
  78. pre_handlers[name[4:]] = getattr(self, name)
  79. elif name.startswith('post_'):
  80. post_handlers[name[5:]] = getattr(self, name)
  81. def walk(self, node, name='', list=list, len=len, type=type):
  82. """Walk the tree starting at a given node.
  83. Maintain a stack of nodes.
  84. """
  85. pre_handlers = self.pre_handlers.get
  86. post_handlers = self.post_handlers.get
  87. nodestack = self.nodestack
  88. emptystack = len(nodestack)
  89. append, pop = nodestack.append, nodestack.pop
  90. append([node, name, list(iter_node(node, name + '_item')), -1])
  91. while len(nodestack) > emptystack:
  92. node, name, subnodes, index = nodestack[-1]
  93. if index >= len(subnodes):
  94. handler = (post_handlers(type(node).__name__) or
  95. post_handlers(name + '_name'))
  96. if handler is None:
  97. pop()
  98. continue
  99. self.cur_node = node
  100. self.cur_name = name
  101. handler()
  102. current = nodestack and nodestack[-1]
  103. popstack = current and current[0] is node
  104. if popstack and current[-1] >= len(current[-2]):
  105. pop()
  106. continue
  107. nodestack[-1][-1] = index + 1
  108. if index < 0:
  109. handler = (pre_handlers(type(node).__name__) or
  110. pre_handlers(name + '_name'))
  111. if handler is not None:
  112. self.cur_node = node
  113. self.cur_name = name
  114. if handler():
  115. pop()
  116. else:
  117. node, name = subnodes[index]
  118. append([node, name, list(iter_node(node, name + '_item')), -1])
  119. @property
  120. def parent(self):
  121. """Return the parent node of the current node."""
  122. nodestack = self.nodestack
  123. if len(nodestack) < 2:
  124. return None
  125. return nodestack[-2][0]
  126. @property
  127. def parent_name(self):
  128. """Return the parent node and name."""
  129. nodestack = self.nodestack
  130. if len(nodestack) < 2:
  131. return None
  132. return nodestack[-2][:2]
  133. def replace(self, new_node):
  134. """Replace a node after first checking integrity of node stack."""
  135. cur_node = self.cur_node
  136. nodestack = self.nodestack
  137. cur = nodestack.pop()
  138. prev = nodestack[-1]
  139. index = prev[-1] - 1
  140. oldnode, name = prev[-2][index]
  141. assert cur[0] is cur_node is oldnode, (cur[0], cur_node, prev[-2],
  142. index)
  143. parent = prev[0]
  144. if isinstance(parent, list):
  145. parent[index] = new_node
  146. else:
  147. setattr(parent, name, new_node)