rtrip.py 6.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209
  1. #! /usr/bin/env python
  2. # -*- coding: utf-8 -*-
  3. """
  4. Part of the astor library for Python AST manipulation.
  5. License: 3-clause BSD
  6. Copyright (c) 2015 Patrick Maupin
  7. """
  8. import sys
  9. import os
  10. import ast
  11. import shutil
  12. import logging
  13. from astor.code_gen import to_source
  14. from astor.file_util import code_to_ast
  15. from astor.node_util import (allow_ast_comparison, dump_tree,
  16. strip_tree, fast_compare)
  17. dsttree = 'tmp_rtrip'
  18. # TODO: Remove this workaround once we remove version 2 support
  19. def out_prep(s, pre_encoded=(sys.version_info[0] == 2)):
  20. return s if pre_encoded else s.encode('utf-8')
  21. def convert(srctree, dsttree=dsttree, readonly=False, dumpall=False,
  22. ignore_exceptions=False, fullcomp=False):
  23. """Walk the srctree, and convert/copy all python files
  24. into the dsttree
  25. """
  26. if fullcomp:
  27. allow_ast_comparison()
  28. parse_file = code_to_ast.parse_file
  29. find_py_files = code_to_ast.find_py_files
  30. srctree = os.path.normpath(srctree)
  31. if not readonly:
  32. dsttree = os.path.normpath(dsttree)
  33. logging.info('')
  34. logging.info('Trashing ' + dsttree)
  35. shutil.rmtree(dsttree, True)
  36. unknown_src_nodes = set()
  37. unknown_dst_nodes = set()
  38. badfiles = set()
  39. broken = []
  40. oldpath = None
  41. allfiles = find_py_files(srctree, None if readonly else dsttree)
  42. for srcpath, fname in allfiles:
  43. # Create destination directory
  44. if not readonly and srcpath != oldpath:
  45. oldpath = srcpath
  46. if srcpath >= srctree:
  47. dstpath = srcpath.replace(srctree, dsttree, 1)
  48. if not dstpath.startswith(dsttree):
  49. raise ValueError("%s not a subdirectory of %s" %
  50. (dstpath, dsttree))
  51. else:
  52. assert srctree.startswith(srcpath)
  53. dstpath = dsttree
  54. os.makedirs(dstpath)
  55. srcfname = os.path.join(srcpath, fname)
  56. logging.info('Converting %s' % srcfname)
  57. try:
  58. srcast = parse_file(srcfname)
  59. except SyntaxError:
  60. badfiles.add(srcfname)
  61. continue
  62. try:
  63. dsttxt = to_source(srcast)
  64. except Exception:
  65. if not ignore_exceptions:
  66. raise
  67. dsttxt = ''
  68. if not readonly:
  69. dstfname = os.path.join(dstpath, fname)
  70. try:
  71. with open(dstfname, 'wb') as f:
  72. f.write(out_prep(dsttxt))
  73. except UnicodeEncodeError:
  74. badfiles.add(dstfname)
  75. # As a sanity check, make sure that ASTs themselves
  76. # round-trip OK
  77. try:
  78. dstast = ast.parse(dsttxt) if readonly else parse_file(dstfname)
  79. except SyntaxError:
  80. dstast = []
  81. if fullcomp:
  82. unknown_src_nodes.update(strip_tree(srcast))
  83. unknown_dst_nodes.update(strip_tree(dstast))
  84. bad = srcast != dstast
  85. else:
  86. bad = not fast_compare(srcast, dstast)
  87. if dumpall or bad:
  88. srcdump = dump_tree(srcast)
  89. dstdump = dump_tree(dstast)
  90. logging.warning(' calculating dump -- %s' %
  91. ('bad' if bad else 'OK'))
  92. if bad:
  93. broken.append(srcfname)
  94. if dumpall or bad:
  95. if not readonly:
  96. try:
  97. with open(dstfname[:-3] + '.srcdmp', 'wb') as f:
  98. f.write(out_prep(srcdump))
  99. except UnicodeEncodeError:
  100. badfiles.add(dstfname[:-3] + '.srcdmp')
  101. try:
  102. with open(dstfname[:-3] + '.dstdmp', 'wb') as f:
  103. f.write(out_prep(dstdump))
  104. except UnicodeEncodeError:
  105. badfiles.add(dstfname[:-3] + '.dstdmp')
  106. elif dumpall:
  107. sys.stdout.write('\n\nAST:\n\n ')
  108. sys.stdout.write(srcdump.replace('\n', '\n '))
  109. sys.stdout.write('\n\nDecompile:\n\n ')
  110. sys.stdout.write(dsttxt.replace('\n', '\n '))
  111. sys.stdout.write('\n\nNew AST:\n\n ')
  112. sys.stdout.write('(same as old)' if dstdump == srcdump
  113. else dstdump.replace('\n', '\n '))
  114. sys.stdout.write('\n')
  115. if badfiles:
  116. logging.warning('\nFiles not processed due to syntax errors:')
  117. for fname in sorted(badfiles):
  118. logging.warning(' %s' % fname)
  119. if broken:
  120. logging.warning('\nFiles failed to round-trip to AST:')
  121. for srcfname in broken:
  122. logging.warning(' %s' % srcfname)
  123. ok_to_strip = 'col_offset _precedence _use_parens lineno _p_op _pp'
  124. ok_to_strip = set(ok_to_strip.split())
  125. bad_nodes = (unknown_dst_nodes | unknown_src_nodes) - ok_to_strip
  126. if bad_nodes:
  127. logging.error('\nERROR -- UNKNOWN NODES STRIPPED: %s' % bad_nodes)
  128. logging.info('\n')
  129. return broken
  130. def usage(msg):
  131. raise SystemExit(textwrap.dedent("""
  132. Error: %s
  133. Usage:
  134. python -m astor.rtrip [readonly] [<source>]
  135. This utility tests round-tripping of Python source to AST
  136. and back to source.
  137. If readonly is specified, then the source will be tested,
  138. but no files will be written.
  139. if the source is specified to be "stdin" (without quotes)
  140. then any source entered at the command line will be compiled
  141. into an AST, converted back to text, and then compiled to
  142. an AST again, and the results will be displayed to stdout.
  143. If neither readonly nor stdin is specified, then rtrip
  144. will create a mirror directory named tmp_rtrip and will
  145. recursively round-trip all the Python source from the source
  146. into the tmp_rtrip dir, after compiling it and then reconstituting
  147. it through code_gen.to_source.
  148. If the source is not specified, the entire Python library will be used.
  149. """) % msg)
  150. if __name__ == '__main__':
  151. import textwrap
  152. args = sys.argv[1:]
  153. readonly = 'readonly' in args
  154. if readonly:
  155. args.remove('readonly')
  156. if not args:
  157. args = [os.path.dirname(textwrap.__file__)]
  158. if len(args) > 1:
  159. usage("Too many arguments")
  160. fname, = args
  161. dumpall = False
  162. if not os.path.exists(fname):
  163. dumpall = fname == 'stdin' or usage("Cannot find directory %s" % fname)
  164. logging.basicConfig(format='%(msg)s', level=logging.INFO)
  165. convert(fname, readonly=readonly or dumpall, dumpall=dumpall)