source_repr.py 7.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273
  1. # -*- coding: utf-8 -*-
  2. """
  3. Part of the astor library for Python AST manipulation.
  4. License: 3-clause BSD
  5. Copyright (c) 2015 Patrick Maupin
  6. Pretty-print source -- post-process for the decompiler
  7. The goals of the initial cut of this engine are:
  8. 1) Do a passable, if not PEP8, job of line-wrapping.
  9. 2) Serve as an example of an interface to the decompiler
  10. for anybody who wants to do a better job. :)
  11. """
  12. def pretty_source(source):
  13. """ Prettify the source.
  14. """
  15. return ''.join(split_lines(source))
  16. def split_lines(source, maxline=79):
  17. """Split inputs according to lines.
  18. If a line is short enough, just yield it.
  19. Otherwise, fix it.
  20. """
  21. result = []
  22. extend = result.extend
  23. append = result.append
  24. line = []
  25. multiline = False
  26. count = 0
  27. for item in source:
  28. newline = type(item)('\n')
  29. index = item.find(newline)
  30. if index:
  31. line.append(item)
  32. multiline = index > 0
  33. count += len(item)
  34. else:
  35. if line:
  36. if count <= maxline or multiline:
  37. extend(line)
  38. else:
  39. wrap_line(line, maxline, result)
  40. count = 0
  41. multiline = False
  42. line = []
  43. append(item)
  44. return result
  45. def count(group, slen=str.__len__):
  46. return sum([slen(x) for x in group])
  47. def wrap_line(line, maxline=79, result=[], count=count):
  48. """ We have a line that is too long,
  49. so we're going to try to wrap it.
  50. """
  51. # Extract the indentation
  52. append = result.append
  53. extend = result.extend
  54. indentation = line[0]
  55. lenfirst = len(indentation)
  56. indent = lenfirst - len(indentation.lstrip())
  57. assert indent in (0, lenfirst)
  58. indentation = line.pop(0) if indent else ''
  59. # Get splittable/non-splittable groups
  60. dgroups = list(delimiter_groups(line))
  61. unsplittable = dgroups[::2]
  62. splittable = dgroups[1::2]
  63. # If the largest non-splittable group won't fit
  64. # on a line, try to add parentheses to the line.
  65. if max(count(x) for x in unsplittable) > maxline - indent:
  66. line = add_parens(line, maxline, indent)
  67. dgroups = list(delimiter_groups(line))
  68. unsplittable = dgroups[::2]
  69. splittable = dgroups[1::2]
  70. # Deal with the first (always unsplittable) group, and
  71. # then set up to deal with the remainder in pairs.
  72. first = unsplittable[0]
  73. append(indentation)
  74. extend(first)
  75. if not splittable:
  76. return result
  77. pos = indent + count(first)
  78. indentation += ' '
  79. indent += 4
  80. if indent >= maxline / 2:
  81. maxline = maxline / 2 + indent
  82. for sg, nsg in zip(splittable, unsplittable[1:]):
  83. if sg:
  84. # If we already have stuff on the line and even
  85. # the very first item won't fit, start a new line
  86. if pos > indent and pos + len(sg[0]) > maxline:
  87. append('\n')
  88. append(indentation)
  89. pos = indent
  90. # Dump lines out of the splittable group
  91. # until the entire thing fits
  92. csg = count(sg)
  93. while pos + csg > maxline:
  94. ready, sg = split_group(sg, pos, maxline)
  95. if ready[-1].endswith(' '):
  96. ready[-1] = ready[-1][:-1]
  97. extend(ready)
  98. append('\n')
  99. append(indentation)
  100. pos = indent
  101. csg = count(sg)
  102. # Dump the remainder of the splittable group
  103. if sg:
  104. extend(sg)
  105. pos += csg
  106. # Dump the unsplittable group, optionally
  107. # preceded by a linefeed.
  108. cnsg = count(nsg)
  109. if pos > indent and pos + cnsg > maxline:
  110. append('\n')
  111. append(indentation)
  112. pos = indent
  113. extend(nsg)
  114. pos += cnsg
  115. def split_group(source, pos, maxline):
  116. """ Split a group into two subgroups. The
  117. first will be appended to the current
  118. line, the second will start the new line.
  119. Note that the first group must always
  120. contain at least one item.
  121. The original group may be destroyed.
  122. """
  123. first = []
  124. source.reverse()
  125. while source:
  126. tok = source.pop()
  127. first.append(tok)
  128. pos += len(tok)
  129. if source:
  130. tok = source[-1]
  131. allowed = (maxline + 1) if tok.endswith(' ') else (maxline - 4)
  132. if pos + len(tok) > allowed:
  133. break
  134. source.reverse()
  135. return first, source
  136. begin_delim = set('([{')
  137. end_delim = set(')]}')
  138. end_delim.add('):')
  139. def delimiter_groups(line, begin_delim=begin_delim,
  140. end_delim=end_delim):
  141. """Split a line into alternating groups.
  142. The first group cannot have a line feed inserted,
  143. the next one can, etc.
  144. """
  145. text = []
  146. line = iter(line)
  147. while True:
  148. # First build and yield an unsplittable group
  149. for item in line:
  150. text.append(item)
  151. if item in begin_delim:
  152. break
  153. if not text:
  154. break
  155. yield text
  156. # Now build and yield a splittable group
  157. level = 0
  158. text = []
  159. for item in line:
  160. if item in begin_delim:
  161. level += 1
  162. elif item in end_delim:
  163. level -= 1
  164. if level < 0:
  165. yield text
  166. text = [item]
  167. break
  168. text.append(item)
  169. else:
  170. assert not text, text
  171. break
  172. statements = set(['del ', 'return', 'yield ', 'if ', 'while '])
  173. def add_parens(line, maxline, indent, statements=statements, count=count):
  174. """Attempt to add parentheses around the line
  175. in order to make it splittable.
  176. """
  177. if line[0] in statements:
  178. index = 1
  179. if not line[0].endswith(' '):
  180. index = 2
  181. assert line[1] == ' '
  182. line.insert(index, '(')
  183. if line[-1] == ':':
  184. line.insert(-1, ')')
  185. else:
  186. line.append(')')
  187. # That was the easy stuff. Now for assignments.
  188. groups = list(get_assign_groups(line))
  189. if len(groups) == 1:
  190. # So sad, too bad
  191. return line
  192. counts = list(count(x) for x in groups)
  193. didwrap = False
  194. # If the LHS is large, wrap it first
  195. if sum(counts[:-1]) >= maxline - indent - 4:
  196. for group in groups[:-1]:
  197. didwrap = False # Only want to know about last group
  198. if len(group) > 1:
  199. group.insert(0, '(')
  200. group.insert(-1, ')')
  201. didwrap = True
  202. # Might not need to wrap the RHS if wrapped the LHS
  203. if not didwrap or counts[-1] > maxline - indent - 10:
  204. groups[-1].insert(0, '(')
  205. groups[-1].append(')')
  206. return [item for group in groups for item in group]
  207. # Assignment operators
  208. ops = list('|^&+-*/%@~') + '<< >> // **'.split() + ['']
  209. ops = set(' %s= ' % x for x in ops)
  210. def get_assign_groups(line, ops=ops):
  211. """ Split a line into groups by assignment (including
  212. augmented assignment)
  213. """
  214. group = []
  215. for item in line:
  216. group.append(item)
  217. if item in ops:
  218. yield group
  219. group = []
  220. yield group