astdump.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112
  1. # Copyright (C) 2022 The Qt Company Ltd.
  2. # SPDX-License-Identifier: LicenseRef-Qt-Commercial OR LGPL-3.0-only OR GPL-2.0-only OR GPL-3.0-only
  3. from __future__ import annotations
  4. """Tool to dump a Python AST"""
  5. import ast
  6. import tokenize
  7. from argparse import ArgumentParser, RawTextHelpFormatter
  8. from enum import Enum
  9. from nodedump import debug_format_node
  10. DESCRIPTION = "Tool to dump a Python AST"
  11. _source_lines = []
  12. _opt_verbose = False
  13. def first_non_space(s):
  14. for i, c in enumerate(s):
  15. if c != ' ':
  16. return i
  17. return 0
  18. class NodeType(Enum):
  19. IGNORE = 1
  20. PRINT_ONE_LINE = 2 # Print as a one liner, do not visit children
  21. PRINT = 3 # Print with opening closing tag, visit children
  22. PRINT_WITH_SOURCE = 4 # Like PRINT, but print source line above
  23. def get_node_type(node):
  24. if isinstance(node, (ast.Load, ast.Store, ast.Delete)):
  25. return NodeType.IGNORE
  26. if isinstance(node, (ast.Add, ast.alias, ast.arg, ast.Eq, ast.Gt, ast.Lt,
  27. ast.Mult, ast.Name, ast.NotEq, ast.NameConstant, ast.Not,
  28. ast.Num, ast.Str)):
  29. return NodeType.PRINT_ONE_LINE
  30. if not hasattr(node, 'lineno'):
  31. return NodeType.PRINT
  32. if isinstance(node, (ast.Attribute)):
  33. return NodeType.PRINT_ONE_LINE if isinstance(node.value, ast.Name) else NodeType.PRINT
  34. return NodeType.PRINT_WITH_SOURCE
  35. class DumpVisitor(ast.NodeVisitor):
  36. def __init__(self):
  37. ast.NodeVisitor.__init__(self)
  38. self._indent = 0
  39. self._printed_source_lines = {-1}
  40. def generic_visit(self, node):
  41. node_type = get_node_type(node)
  42. if _opt_verbose and node_type in (NodeType.IGNORE, NodeType.PRINT_ONE_LINE):
  43. node_type = NodeType.PRINT
  44. if node_type == NodeType.IGNORE:
  45. return
  46. self._indent = self._indent + 1
  47. indent = ' ' * self._indent
  48. if node_type == NodeType.PRINT_WITH_SOURCE:
  49. line_number = node.lineno - 1
  50. if line_number not in self._printed_source_lines:
  51. self._printed_source_lines.add(line_number)
  52. line = _source_lines[line_number]
  53. non_space = first_non_space(line)
  54. print('{:04d} {}{}'.format(line_number, '_' * non_space,
  55. line[non_space:]))
  56. if node_type == NodeType.PRINT_ONE_LINE:
  57. print(indent, debug_format_node(node))
  58. else:
  59. print(indent, '>', debug_format_node(node))
  60. ast.NodeVisitor.generic_visit(self, node)
  61. print(indent, '<', type(node).__name__)
  62. self._indent = self._indent - 1
  63. def parse_ast(filename):
  64. node = None
  65. with tokenize.open(filename) as f:
  66. global _source_lines
  67. source = f.read()
  68. _source_lines = source.split('\n')
  69. node = ast.parse(source, mode="exec")
  70. return node
  71. def create_arg_parser(desc):
  72. parser = ArgumentParser(description=desc,
  73. formatter_class=RawTextHelpFormatter)
  74. parser.add_argument('--verbose', '-v', action='store_true',
  75. help='Verbose')
  76. parser.add_argument('source', type=str, help='Python source')
  77. return parser
  78. if __name__ == '__main__':
  79. arg_parser = create_arg_parser(DESCRIPTION)
  80. options = arg_parser.parse_args()
  81. _opt_verbose = options.verbose
  82. title = f'AST tree for {options.source}'
  83. print('=' * len(title))
  84. print(title)
  85. print('=' * len(title))
  86. tree = parse_ast(options.source)
  87. DumpVisitor().visit(tree)