unparse.py

"Usage: unparse.py <path to source file>"
importsys
importast
importcStringIO
importos

# Large float and imaginary literals get turned into infinities in the AST.
# We unparse those infinities to INFSTR.
INFSTR="1e"+repr(sys.float_info.max_10_exp+1)

definterleave(inter, f, seq):
"""Call f on each item in seq, calling inter() in between.
 """
seq=iter(seq)
try:
f(next(seq))
exceptStopIteration:
pass
else:
forxinseq:
inter()
f(x)

classUnparser:
"""Methods in this class recursively traverse an AST and
 output source code for the abstract syntax; original formatting
 is disregarded. """

def__init__(self, tree, file=sys.stdout):
"""Unparser(tree, file=sys.stdout) -> None.
 Print the source for tree to file."""
self.f=file
self.future_imports= []
self._indent=0
self.dispatch(tree)
self.f.write("")
self.f.flush()

deffill(self, text=""):
"Indent a piece of text, according to the current indentation level"
self.f.write("\n"+" "*self._indent+text)

defwrite(self, text):
"Append a piece of text to the current line."
self.f.write(text)

defenter(self):
"Print ':', and increase the indentation."
self.write(":")
self._indent+=1

defleave(self):
"Decrease the indentation level."
self._indent-=1

defdispatch(self, tree):
"Dispatcher function, dispatching tree type T to method _T."
ifisinstance(tree, list):
fortintree:
self.dispatch(t)
return
meth=getattr(self, "_"+tree.__class__.__name__)
meth(tree)


############### Unparsing methods ######################
# There should be one method per concrete grammar type #
# Constructors should be grouped by sum type. Ideally, #
# this would follow the order in the grammar, but #
# currently doesn't. #
########################################################

def_Module(self, tree):
forstmtintree.body:
self.dispatch(stmt)

# stmt
def_Expr(self, tree):
self.fill()
self.dispatch(tree.value)

def_Import(self, t):
self.fill("import ")
interleave(lambda: self.write(", "), self.dispatch, t.names)

def_ImportFrom(self, t):
# A from __future__ import may affect unparsing, so record it.
ift.moduleandt.module=='__future__':
self.future_imports.extend(n.namefornint.names)

self.fill("from ")
self.write("."*t.level)
ift.module:
self.write(t.module)
self.write(" import ")
interleave(lambda: self.write(", "), self.dispatch, t.names)

def_Assign(self, t):
self.fill()
fortargetint.targets:
self.dispatch(target)
self.write(" = ")
self.dispatch(t.value)

def_AugAssign(self, t):
self.fill()
self.dispatch(t.target)
self.write(" "+self.binop[t.op.__class__.__name__]+"= ")
self.dispatch(t.value)

def_Return(self, t):
self.fill("return")
ift.value:
self.write(" ")
self.dispatch(t.value)

def_Pass(self, t):
self.fill("pass")

def_Break(self, t):
self.fill("break")

def_Continue(self, t):
self.fill("continue")

def_Delete(self, t):
self.fill("del ")
interleave(lambda: self.write(", "), self.dispatch, t.targets)

def_Assert(self, t):
self.fill("assert ")
self.dispatch(t.test)
ift.msg:
self.write(", ")
self.dispatch(t.msg)

def_Exec(self, t):
self.fill("exec ")
self.dispatch(t.body)
ift.globals:
self.write(" in ")
self.dispatch(t.globals)
ift.locals:
self.write(", ")
self.dispatch(t.locals)

def_Print(self, t):
self.fill("print ")
do_comma=False
ift.dest:
self.write(">>")
self.dispatch(t.dest)
do_comma=True
foreint.values:
ifdo_comma:self.write(", ")
else:do_comma=True
self.dispatch(e)
ifnott.nl:
self.write(",")

def_Global(self, t):
self.fill("global ")
interleave(lambda: self.write(", "), self.write, t.names)

def_Yield(self, t):
self.write("(")
self.write("yield")
ift.value:
self.write(" ")
self.dispatch(t.value)
self.write(")")

def_Raise(self, t):
self.fill('raise ')
ift.type:
self.dispatch(t.type)
ift.inst:
self.write(", ")
self.dispatch(t.inst)
ift.tback:
self.write(", ")
self.dispatch(t.tback)

def_TryExcept(self, t):
self.fill("try")
self.enter()
self.dispatch(t.body)
self.leave()

forexint.handlers:
self.dispatch(ex)
ift.orelse:
self.fill("else")
self.enter()
self.dispatch(t.orelse)
self.leave()

def_TryFinally(self, t):
iflen(t.body) ==1andisinstance(t.body[0], ast.TryExcept):
# try-except-finally
self.dispatch(t.body)
else:
self.fill("try")
self.enter()
self.dispatch(t.body)
self.leave()

self.fill("finally")
self.enter()
self.dispatch(t.finalbody)
self.leave()

def_ExceptHandler(self, t):
self.fill("except")
ift.type:
self.write(" ")
self.dispatch(t.type)
ift.name:
self.write(" as ")
self.dispatch(t.name)
self.enter()
self.dispatch(t.body)
self.leave()

def_ClassDef(self, t):
self.write("\n")
fordecoint.decorator_list:
self.fill("@")
self.dispatch(deco)
self.fill("class "+t.name)
ift.bases:
self.write("(")
foraint.bases:
self.dispatch(a)
self.write(", ")
self.write(")")
self.enter()
self.dispatch(t.body)
self.leave()

def_FunctionDef(self, t):
self.write("\n")
fordecoint.decorator_list:
self.fill("@")
self.dispatch(deco)
self.fill("def "+t.name+"(")
self.dispatch(t.args)
self.write(")")
self.enter()
self.dispatch(t.body)
self.leave()

def_For(self, t):
self.fill("for ")
self.dispatch(t.target)
self.write(" in ")
self.dispatch(t.iter)
self.enter()
self.dispatch(t.body)
self.leave()
ift.orelse:
self.fill("else")
self.enter()
self.dispatch(t.orelse)
self.leave()

def_If(self, t):
self.fill("if ")
self.dispatch(t.test)
self.enter()
self.dispatch(t.body)
self.leave()
# collapse nested ifs into equivalent elifs.
while (t.orelseandlen(t.orelse) ==1and
isinstance(t.orelse[0], ast.If)):
t=t.orelse[0]
self.fill("elif ")
self.dispatch(t.test)
self.enter()
self.dispatch(t.body)
self.leave()
# final else
ift.orelse:
self.fill("else")
self.enter()
self.dispatch(t.orelse)
self.leave()

def_While(self, t):
self.fill("while ")
self.dispatch(t.test)
self.enter()
self.dispatch(t.body)
self.leave()
ift.orelse:
self.fill("else")
self.enter()
self.dispatch(t.orelse)
self.leave()

def_With(self, t):
self.fill("with ")
self.dispatch(t.context_expr)
ift.optional_vars:
self.write(" as ")
self.dispatch(t.optional_vars)
self.enter()
self.dispatch(t.body)
self.leave()

# expr
def_Str(self, tree):
# if from __future__ import unicode_literals is in effect,
# then we want to output string literals using a 'b' prefix
# and unicode literals with no prefix.
if"unicode_literals"notinself.future_imports:
self.write(repr(tree.s))
elifisinstance(tree.s, str):
self.write("b"+repr(tree.s))
elifisinstance(tree.s, unicode):
self.write(repr(tree.s).lstrip("u"))
else:
assertFalse, "shouldn't get here"

def_Name(self, t):
self.write(t.id)

def_Repr(self, t):
self.write("`")
self.dispatch(t.value)
self.write("`")

def_Num(self, t):
repr_n=repr(t.n)
# Parenthesize negative numbers, to avoid turning (-1)**2 into -1**2.
ifrepr_n.startswith("-"):
self.write("(")
# Substitute overflowing decimal literal for AST infinities.
self.write(repr_n.replace("inf", INFSTR))
ifrepr_n.startswith("-"):
self.write(")")

def_List(self, t):
self.write("[")
interleave(lambda: self.write(", "), self.dispatch, t.elts)
self.write("]")

def_ListComp(self, t):
self.write("[")
self.dispatch(t.elt)
forgenint.generators:
self.dispatch(gen)
self.write("]")

def_GeneratorExp(self, t):
self.write("(")
self.dispatch(t.elt)
forgenint.generators:
self.dispatch(gen)
self.write(")")

def_SetComp(self, t):
self.write("{")
self.dispatch(t.elt)
forgenint.generators:
self.dispatch(gen)
self.write("}")

def_DictComp(self, t):
self.write("{")
self.dispatch(t.key)
self.write(": ")
self.dispatch(t.value)
forgenint.generators:
self.dispatch(gen)
self.write("}")

def_comprehension(self, t):
self.write(" for ")
self.dispatch(t.target)
self.write(" in ")
self.dispatch(t.iter)
forif_clauseint.ifs:
self.write(" if ")
self.dispatch(if_clause)

def_IfExp(self, t):
self.write("(")
self.dispatch(t.body)
self.write(" if ")
self.dispatch(t.test)
self.write(" else ")
self.dispatch(t.orelse)
self.write(")")

def_Set(self, t):
assert(t.elts) # should be at least one element
self.write("{")
interleave(lambda: self.write(", "), self.dispatch, t.elts)
self.write("}")

def_Dict(self, t):
self.write("{")
defwrite_pair(pair):
 (k, v) =pair
self.dispatch(k)
self.write(": ")
self.dispatch(v)
interleave(lambda: self.write(", "), write_pair, zip(t.keys, t.values))
self.write("}")

def_Tuple(self, t):
self.write("(")
iflen(t.elts) ==1:
 (elt,) =t.elts
self.dispatch(elt)
self.write(",")
else:
interleave(lambda: self.write(", "), self.dispatch, t.elts)
self.write(")")

unop= {"Invert":"~", "Not": "not", "UAdd":"+", "USub":"-"}
def_UnaryOp(self, t):
self.write("(")
self.write(self.unop[t.op.__class__.__name__])
self.write(" ")
# If we're applying unary minus to a number, parenthesize the number.
# This is necessary: -2147483648 is different from -(2147483648) on
# a 32-bit machine (the first is an int, the second a long), and
# -7j is different from -(7j). (The first has real part 0.0, the second
# has real part -0.0.)
ifisinstance(t.op, ast.USub) andisinstance(t.operand, ast.Num):
self.write("(")
self.dispatch(t.operand)
self.write(")")
else:
self.dispatch(t.operand)
self.write(")")

binop= { "Add":"+", "Sub":"-", "Mult":"*", "Div":"/", "Mod":"%",
"LShift":"<<", "RShift":">>", "BitOr":"|", "BitXor":"^", "BitAnd":"&",
"FloorDiv":"//", "Pow": "**"}
def_BinOp(self, t):
self.write("(")
self.dispatch(t.left)
self.write(" "+self.binop[t.op.__class__.__name__] +" ")
self.dispatch(t.right)
self.write(")")

cmpops= {"Eq":"==", "NotEq":"!=", "Lt":"<", "LtE":"<=", "Gt":">", "GtE":">=",
"Is":"is", "IsNot":"is not", "In":"in", "NotIn":"not in"}
def_Compare(self, t):
self.write("(")
self.dispatch(t.left)
foro, einzip(t.ops, t.comparators):
self.write(" "+self.cmpops[o.__class__.__name__] +" ")
self.dispatch(e)
self.write(")")

boolops= {ast.And: 'and', ast.Or: 'or'}
def_BoolOp(self, t):
self.write("(")
s=" %s "%self.boolops[t.op.__class__]
interleave(lambda: self.write(s), self.dispatch, t.values)
self.write(")")

def_Attribute(self,t):
self.dispatch(t.value)
# Special case: 3.__abs__() is a syntax error, so if t.value
# is an integer literal then we need to either parenthesize
# it or add an extra space to get 3 .__abs__().
ifisinstance(t.value, ast.Num) andisinstance(t.value.n, int):
self.write(" ")
self.write(".")
self.write(t.attr)

def_Call(self, t):
self.dispatch(t.func)
self.write("(")
comma=False
foreint.args:
ifcomma: self.write(", ")
else: comma=True
self.dispatch(e)
foreint.keywords:
ifcomma: self.write(", ")
else: comma=True
self.dispatch(e)
ift.starargs:
ifcomma: self.write(", ")
else: comma=True
self.write("*")
self.dispatch(t.starargs)
ift.kwargs:
ifcomma: self.write(", ")
else: comma=True
self.write("**")
self.dispatch(t.kwargs)
self.write(")")

def_Subscript(self, t):
self.dispatch(t.value)
self.write("[")
self.dispatch(t.slice)
self.write("]")

# slice
def_Ellipsis(self, t):
self.write("...")

def_Index(self, t):
self.dispatch(t.value)

def_Slice(self, t):
ift.lower:
self.dispatch(t.lower)
self.write(":")
ift.upper:
self.dispatch(t.upper)
ift.step:
self.write(":")
self.dispatch(t.step)

def_ExtSlice(self, t):
interleave(lambda: self.write(', '), self.dispatch, t.dims)

# others
def_arguments(self, t):
first=True
# normal arguments
defaults= [None] * (len(t.args) -len(t.defaults)) +t.defaults
fora,dinzip(t.args, defaults):
iffirst:first=False
else: self.write(", ")
self.dispatch(a),
ifd:
self.write("=")
self.dispatch(d)

# varargs
ift.vararg:
iffirst:first=False
else: self.write(", ")
self.write("*")
self.write(t.vararg)

# kwargs
ift.kwarg:
iffirst:first=False
else: self.write(", ")
self.write("**"+t.kwarg)

def_keyword(self, t):
self.write(t.arg)
self.write("=")
self.dispatch(t.value)

def_Lambda(self, t):
self.write("(")
self.write("lambda ")
self.dispatch(t.args)
self.write(": ")
self.dispatch(t.body)
self.write(")")

def_alias(self, t):
self.write(t.name)
ift.asname:
self.write(" as "+t.asname)

defroundtrip(filename, output=sys.stdout):
withopen(filename, "r") aspyfile:
source=pyfile.read()
tree=compile(source, filename, "exec", ast.PyCF_ONLY_AST)
Unparser(tree, output)



deftestdir(a):
try:
names= [nforninos.listdir(a) ifn.endswith('.py')]
exceptOSError:
sys.stderr.write("Directory not readable: %s"%a)
else:
forninnames:
fullname=os.path.join(a, n)
ifos.path.isfile(fullname):
output=cStringIO.StringIO()
print'Testing %s'%fullname
try:
roundtrip(fullname, output)
exceptExceptionase:
print' Failed to compile, exception is %s'%repr(e)
elifos.path.isdir(fullname):
testdir(fullname)

defmain(args):
ifargs[0] =='--testdir':
forainargs[1:]:
testdir(a)
else:
forainargs:
roundtrip(a)

if__name__=='__main__':
main(sys.argv[1:])