| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273 |
- # -*- coding: utf-8 -*-
- """
- Part of the astor library for Python AST manipulation.
- License: 3-clause BSD
- Copyright (c) 2015 Patrick Maupin
- Pretty-print source -- post-process for the decompiler
- The goals of the initial cut of this engine are:
- 1) Do a passable, if not PEP8, job of line-wrapping.
- 2) Serve as an example of an interface to the decompiler
- for anybody who wants to do a better job. :)
- """
- def pretty_source(source):
- """ Prettify the source.
- """
- return ''.join(split_lines(source))
- def split_lines(source, maxline=79):
- """Split inputs according to lines.
- If a line is short enough, just yield it.
- Otherwise, fix it.
- """
- result = []
- extend = result.extend
- append = result.append
- line = []
- multiline = False
- count = 0
- for item in source:
- newline = type(item)('\n')
- index = item.find(newline)
- if index:
- line.append(item)
- multiline = index > 0
- count += len(item)
- else:
- if line:
- if count <= maxline or multiline:
- extend(line)
- else:
- wrap_line(line, maxline, result)
- count = 0
- multiline = False
- line = []
- append(item)
- return result
- def count(group, slen=str.__len__):
- return sum([slen(x) for x in group])
- def wrap_line(line, maxline=79, result=[], count=count):
- """ We have a line that is too long,
- so we're going to try to wrap it.
- """
- # Extract the indentation
- append = result.append
- extend = result.extend
- indentation = line[0]
- lenfirst = len(indentation)
- indent = lenfirst - len(indentation.lstrip())
- assert indent in (0, lenfirst)
- indentation = line.pop(0) if indent else ''
- # Get splittable/non-splittable groups
- dgroups = list(delimiter_groups(line))
- unsplittable = dgroups[::2]
- splittable = dgroups[1::2]
- # If the largest non-splittable group won't fit
- # on a line, try to add parentheses to the line.
- if max(count(x) for x in unsplittable) > maxline - indent:
- line = add_parens(line, maxline, indent)
- dgroups = list(delimiter_groups(line))
- unsplittable = dgroups[::2]
- splittable = dgroups[1::2]
- # Deal with the first (always unsplittable) group, and
- # then set up to deal with the remainder in pairs.
- first = unsplittable[0]
- append(indentation)
- extend(first)
- if not splittable:
- return result
- pos = indent + count(first)
- indentation += ' '
- indent += 4
- if indent >= maxline / 2:
- maxline = maxline / 2 + indent
- for sg, nsg in zip(splittable, unsplittable[1:]):
- if sg:
- # If we already have stuff on the line and even
- # the very first item won't fit, start a new line
- if pos > indent and pos + len(sg[0]) > maxline:
- append('\n')
- append(indentation)
- pos = indent
- # Dump lines out of the splittable group
- # until the entire thing fits
- csg = count(sg)
- while pos + csg > maxline:
- ready, sg = split_group(sg, pos, maxline)
- if ready[-1].endswith(' '):
- ready[-1] = ready[-1][:-1]
- extend(ready)
- append('\n')
- append(indentation)
- pos = indent
- csg = count(sg)
- # Dump the remainder of the splittable group
- if sg:
- extend(sg)
- pos += csg
- # Dump the unsplittable group, optionally
- # preceded by a linefeed.
- cnsg = count(nsg)
- if pos > indent and pos + cnsg > maxline:
- append('\n')
- append(indentation)
- pos = indent
- extend(nsg)
- pos += cnsg
- def split_group(source, pos, maxline):
- """ Split a group into two subgroups. The
- first will be appended to the current
- line, the second will start the new line.
- Note that the first group must always
- contain at least one item.
- The original group may be destroyed.
- """
- first = []
- source.reverse()
- while source:
- tok = source.pop()
- first.append(tok)
- pos += len(tok)
- if source:
- tok = source[-1]
- allowed = (maxline + 1) if tok.endswith(' ') else (maxline - 4)
- if pos + len(tok) > allowed:
- break
- source.reverse()
- return first, source
- begin_delim = set('([{')
- end_delim = set(')]}')
- end_delim.add('):')
- def delimiter_groups(line, begin_delim=begin_delim,
- end_delim=end_delim):
- """Split a line into alternating groups.
- The first group cannot have a line feed inserted,
- the next one can, etc.
- """
- text = []
- line = iter(line)
- while True:
- # First build and yield an unsplittable group
- for item in line:
- text.append(item)
- if item in begin_delim:
- break
- if not text:
- break
- yield text
- # Now build and yield a splittable group
- level = 0
- text = []
- for item in line:
- if item in begin_delim:
- level += 1
- elif item in end_delim:
- level -= 1
- if level < 0:
- yield text
- text = [item]
- break
- text.append(item)
- else:
- assert not text, text
- break
- statements = set(['del ', 'return', 'yield ', 'if ', 'while '])
- def add_parens(line, maxline, indent, statements=statements, count=count):
- """Attempt to add parentheses around the line
- in order to make it splittable.
- """
- if line[0] in statements:
- index = 1
- if not line[0].endswith(' '):
- index = 2
- assert line[1] == ' '
- line.insert(index, '(')
- if line[-1] == ':':
- line.insert(-1, ')')
- else:
- line.append(')')
- # That was the easy stuff. Now for assignments.
- groups = list(get_assign_groups(line))
- if len(groups) == 1:
- # So sad, too bad
- return line
- counts = list(count(x) for x in groups)
- didwrap = False
- # If the LHS is large, wrap it first
- if sum(counts[:-1]) >= maxline - indent - 4:
- for group in groups[:-1]:
- didwrap = False # Only want to know about last group
- if len(group) > 1:
- group.insert(0, '(')
- group.insert(-1, ')')
- didwrap = True
- # Might not need to wrap the RHS if wrapped the LHS
- if not didwrap or counts[-1] > maxline - indent - 10:
- groups[-1].insert(0, '(')
- groups[-1].append(')')
- return [item for group in groups for item in group]
- # Assignment operators
- ops = list('|^&+-*/%@~') + '<< >> // **'.split() + ['']
- ops = set(' %s= ' % x for x in ops)
- def get_assign_groups(line, ops=ops):
- """ Split a line into groups by assignment (including
- augmented assignment)
- """
- group = []
- for item in line:
- group.append(item)
- if item in ops:
- yield group
- group = []
- yield group
|