"Usage: unparse.py <path to source file>"

import sys

import _ast

import cStringIO

import os

def interleave(inter, f, seq):

    """Call f on each item in seq, calling inter() in between.

    """

    seq = iter(seq)

    try:

        f(seq.next())

    except StopIteration:

        pass

    else:

        for x in seq:

            inter()

            f(x)

class Unparser:

    """Methods in this class recursively traverse an AST and

    output source code for the abstract syntax; original formatting

    is disregarged. """

    def __init__(self, tree, file = sys.stdout):

        """Unparser(tree, file=sys.stdout) -> None.

         Print the source for tree to file."""

        self.f = file

        self._indent = 0

        self.dispatch(tree)

        print >>self.f,""

        self.f.flush()

    def fill(self, text = ""):

        "Indent a piece of text, according to the current indentation level"

        self.f.write("\n"+"    "*self._indent + text)

    def write(self, text):

        "Append a piece of text to the current line."

        self.f.write(text)

    def enter(self):

        "Print ':', and increase the indentation."

        self.write(":")

        self._indent += 1

    def leave(self):

        "Decrease the indentation level."

        self._indent -= 1

    def dispatch(self, tree):

        "Dispatcher function, dispatching tree type T to method _T."

        if isinstance(tree, list):

            for t in tree:

                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):

        for stmt in tree.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):

        self.fill("from ")

        self.write(t.module)

        self.write(" import ")

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

        # XXX(jpe) what is level for?

    def _Assign(self, t):

        self.fill()

        for target in t.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")

        if t.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 ")

        self.dispatch(t.targets)

    def _Assert(self, t):

        self.fill("assert ")

        self.dispatch(t.test)

        if t.msg:

            self.write(", ")

            self.dispatch(t.msg)

    def _Exec(self, t):

        self.fill("exec ")

        self.dispatch(t.body)

        if t.globals:

            self.write(" in ")

            self.dispatch(t.globals)

        if t.locals:

            self.write(", ")

            self.dispatch(t.locals)

    def _Print(self, t):

        self.fill("print ")

        do_comma = False

        if t.dest:

            self.write(">>")

            self.dispatch(t.dest)

            do_comma = True

        for e in t.values:

            if do_comma:self.write(", ")

            else:do_comma=True

            self.dispatch(e)

        if not t.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")

        if t.value:

            self.write(" ")

            self.dispatch(t.value)

        self.write(")")

    def _Raise(self, t):

        self.fill('raise ')

        if t.type:

            self.dispatch(t.type)

        if t.inst:

            self.write(", ")

            self.dispatch(t.inst)

        if t.tback:

            self.write(", ")

            self.dispatch(t.tback)

    def _TryExcept(self, t):

        self.fill("try")

        self.enter()

        self.dispatch(t.body)

        self.leave()

        for ex in t.handlers:

            self.dispatch(ex)

        if t.orelse:

            self.fill("else")

            self.enter()

            self.dispatch(t.orelse)

            self.leave()

    def _TryFinally(self, t):

        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")

        if t.type:

            self.write(" ")

            self.dispatch(t.type)

        if t.name:

            self.write(", ")

            self.dispatch(t.name)

        self.enter()

        self.dispatch(t.body)

        self.leave()

    def _ClassDef(self, t):

        self.write("\n")

        self.fill("class "+t.name)

        if t.bases:

            self.write("(")

            for a in t.bases:

                self.dispatch(a)

                self.write(", ")

            self.write(")")

        self.enter()

        self.dispatch(t.body)

        self.leave()

    def _FunctionDef(self, t):

        self.write("\n")

        for deco in t.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()

        if t.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()

        # XXX elif?

        self.dispatch(t.body)

        self.leave()

        if t.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()

        if t.orelse:

            self.fill("else")

            self.enter()

            self.dispatch(t.orelse)

            self.leave

    def _With(self, t):

        self.fill("with ")

        self.dispatch(t.context_expr)

        if t.optional_vars:

            self.write(" as ")

            self.dispatch(t.optional_vars)

        self.enter()

        self.dispatch(t.body)

        self.leave()

    # expr

    def _Str(self, tree):

        self.write(repr(tree.s))

    def _Name(self, t):

        self.write(t.id)

    def _Repr(self, t):

        self.write("`")

        self.dispatch(t.value)

        self.write("`")

    def _Num(self, t):

        self.write(repr(t.n))

    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)

        for gen in t.generators:

            self.dispatch(gen)

        self.write("]")

    def _GeneratorExp(self, t):

        self.write("(")

        self.dispatch(t.elt)

        for gen in t.generators:

            self.dispatch(gen)

        self.write(")")

    def _comprehension(self, t):

        self.write(" for ")

        self.dispatch(t.target)

        self.write(" in ")

        self.dispatch(t.iter)

        for if_clause in t.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 _Dict(self, t):

        self.write("{")

        def writem((k, v)):

            self.dispatch(k)

            self.write(": ")

            self.dispatch(v)

        interleave(lambda: self.write(", "), writem, zip(t.keys, t.values))

        self.write("}")

    def _Tuple(self, t):

        self.write("(")

        if len(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.unop[t.op.__class__.__name__])

        self.write("(")

        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)

        for o, e in zip(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)

        self.write(".")

        self.write(t.attr)

    def _Call(self, t):

        self.dispatch(t.func)

        self.write("(")

        comma = False

        for e in t.args:

            if comma: self.write(", ")

            else: comma = True

            self.dispatch(e)

        for e in t.keywords:

            if comma: self.write(", ")

            else: comma = True

            self.dispatch(e)

        if t.starargs:

            if comma: self.write(", ")

            else: comma = True

            self.write("*")

            self.dispatch(t.starargs)

        if t.kwargs:

            if comma: 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):

        if t.lower:

            self.dispatch(t.lower)

        self.write(":")

        if t.upper:

            self.dispatch(t.upper)

        if t.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

        nonDef = len(t.args)-len(t.defaults)

        for a in t.args[0:nonDef]:

            if first:first = False

            else: self.write(", ")

            self.dispatch(a)

        for a,d in zip(t.args[nonDef:], t.defaults):

            if first:first = False

            else: self.write(", ")

            self.dispatch(a),

            self.write("=")

            self.dispatch(d)

        if t.vararg:

            if first:first = False

            else: self.write(", ")

            self.write("*"+t.vararg)

        if t.kwarg:

            if first: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("lambda ")

        self.dispatch(t.args)

        self.write(": ")

        self.dispatch(t.body)

    def _alias(self, t):

        self.write(t.name)

        if t.asname:

            self.write(" as "+t.asname)

def roundtrip(filename, output=sys.stdout):

    source = open(filename).read()

    tree = compile(source, filename, "exec", _ast.PyCF_ONLY_AST)

    Unparser(tree, output)

def testdir(a):

    try:

        names = [n for n in os.listdir(a) if n.endswith('.py')]

    except OSError:

        print >> sys.stderr, "Directory not readable: %s" % a

    else:

        for n in names:

            fullname = os.path.join(a, n)

            if os.path.isfile(fullname):

                output = cStringIO.StringIO()

                print 'Testing %s' % fullname

                try:

                    roundtrip(fullname, output)

                except Exception, e:

                    print '  Failed to compile, exception is %s' % repr(e)

            elif os.path.isdir(fullname):

                testdir(fullname)

def main(args):

    if args[0] == '--testdir':

        for a in args[1:]:

            testdir(a)

    else:

        for a in args:

            roundtrip(a)

if __name__=='__main__':

    main(sys.argv[1:])