pytree_visitor.py 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135
  1. # Copyright 2015 Google Inc. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. """Generic visitor pattern for pytrees.
  15. The lib2to3 parser produces a "pytree" - syntax tree consisting of Node
  16. and Leaf types. This module implements a visitor pattern for such trees.
  17. It also exports a basic "dumping" visitor that dumps a textual representation of
  18. a pytree into a stream.
  19. PyTreeVisitor: a generic visitor pattern for pytrees.
  20. PyTreeDumper: a configurable "dumper" for displaying pytrees.
  21. DumpPyTree(): a convenience function to dump a pytree.
  22. """
  23. import sys
  24. from yapf_third_party._ylib2to3 import pytree
  25. from yapf.pytree import pytree_utils
  26. class PyTreeVisitor(object):
  27. """Visitor pattern for pytree trees.
  28. Methods named Visit_XXX will be invoked when a node with type XXX is
  29. encountered in the tree. The type is either a token type (for Leaf nodes) or
  30. grammar symbols (for Node nodes). The return value of Visit_XXX methods is
  31. ignored by the visitor.
  32. Visitors can modify node contents but must not change the tree structure
  33. (e.g. add/remove children and move nodes around).
  34. This is a very common visitor pattern in Python code; it's also used in the
  35. Python standard library ast module for providing AST visitors.
  36. Note: this makes names that aren't style conformant, so such visitor methods
  37. need to be marked with # pylint: disable=invalid-name We don't have a choice
  38. here, because lib2to3 nodes have under_separated names.
  39. For more complex behavior, the visit, DefaultNodeVisit and DefaultLeafVisit
  40. methods can be overridden. Don't forget to invoke DefaultNodeVisit for nodes
  41. that may have children - otherwise the children will not be visited.
  42. """
  43. def Visit(self, node):
  44. """Visit a node."""
  45. method = 'Visit_{0}'.format(pytree_utils.NodeName(node))
  46. if hasattr(self, method):
  47. # Found a specific visitor for this node
  48. getattr(self, method)(node)
  49. else:
  50. if isinstance(node, pytree.Leaf):
  51. self.DefaultLeafVisit(node)
  52. else:
  53. self.DefaultNodeVisit(node)
  54. def DefaultNodeVisit(self, node):
  55. """Default visitor for Node: visits the node's children depth-first.
  56. This method is invoked when no specific visitor for the node is defined.
  57. Arguments:
  58. node: the node to visit
  59. """
  60. for child in node.children:
  61. self.Visit(child)
  62. def DefaultLeafVisit(self, leaf):
  63. """Default visitor for Leaf: no-op.
  64. This method is invoked when no specific visitor for the leaf is defined.
  65. Arguments:
  66. leaf: the leaf to visit
  67. """
  68. pass
  69. def DumpPyTree(tree, target_stream=sys.stdout):
  70. """Convenience function for dumping a given pytree.
  71. This function presents a very minimal interface. For more configurability (for
  72. example, controlling how specific node types are displayed), use PyTreeDumper
  73. directly.
  74. Arguments:
  75. tree: the tree to dump.
  76. target_stream: the stream to dump the tree to. A file-like object. By
  77. default will dump into stdout.
  78. """
  79. dumper = PyTreeDumper(target_stream)
  80. dumper.Visit(tree)
  81. class PyTreeDumper(PyTreeVisitor):
  82. """Visitor that dumps the tree to a stream.
  83. Implements the PyTreeVisitor interface.
  84. """
  85. def __init__(self, target_stream=sys.stdout):
  86. """Create a tree dumper.
  87. Arguments:
  88. target_stream: the stream to dump the tree to. A file-like object. By
  89. default will dump into stdout.
  90. """
  91. self._target_stream = target_stream
  92. self._current_indent = 0
  93. def _DumpString(self, s):
  94. self._target_stream.write('{0}{1}\n'.format(' ' * self._current_indent, s))
  95. def DefaultNodeVisit(self, node):
  96. # Dump information about the current node, and then use the generic
  97. # DefaultNodeVisit visitor to dump each of its children.
  98. self._DumpString(pytree_utils.DumpNodeToString(node))
  99. self._current_indent += 2
  100. super(PyTreeDumper, self).DefaultNodeVisit(node)
  101. self._current_indent -= 2
  102. def DefaultLeafVisit(self, leaf):
  103. self._DumpString(pytree_utils.DumpNodeToString(leaf))