"""Convert compiler package AST nodes into Python source code.

The entry point is the L{pp} function.

This is mostly used to present certain nodes in the rendered documentation -- for
example any default values for a function's arguments are rendered using L{pp}.

The code was stolen from exarkun's svn.twistedmatrix.com sandbox:
U{http://twistedmatrix.com/trac/browser/sandbox/exarkun/ast/ast_pp.py}
"""


from cStringIO import StringIO
from compiler import parse, walk

class SourceWriter(object):
    _i = 0

    def __init__(self):
        self.s = StringIO()

    def w(self, s):
        self.s.write(s)

    def nl(self):
        self.s.write('\n')
        self.s.write(' ' * 4 * self._i)

    def indent(self):
        self._i += 1
        self.nl()

    def dedent(self):
        self._i -= 1
        self.nl()

    def visitModule(self, node):
        if node.doc is not None:
            self.wl(repr(node.doc))
        walk(node.node, self)

    def visitStmt(self, node):
        for n in node.getChildren():
            walk(n, self)

    def visitFunction(self, node):
        fmt = 'def %s(%s):'
        if node.defaults:
            nargs = len(node.argnames)
            ndefs = len(node.defaults)
            noDefaults = node.argnames[:nargs-ndefs]
            s = ', '.join(node.argnames[:noDefaults])
            if ndefs < nargs:
                argdefs = zip(node.argnames[noDefaults:], node.defaults)
                s = s + ', ' + ', '.join(['='.join(x) for x in argdefs])
        else:
            s = ', '.join(node.argnames)
        self.w(fmt % (node.name, s))
        self.indent()
        try:
            walk(node.code, self)
        finally:
            self.dedent()

    def visitAssign(self, node):
        walk(node.nodes[0], self)
        self.w(' = ')
        walk(node.expr, self)
        self.nl()

    def visitAssName(self, node):
        self.w(node.name)

    def visitCallFunc(self, node):
        walk(node.node, self)
        self.w('(')
        for a in node.args[:-1]:
            walk(a, self)
            self.w(', ')
        for a in node.args[-1:]:
            walk(a, self)
        self.w(')')

    def visitListComp(self, node):
        self.w('[')
        walk(node.expr, self)
        for q in node.quals:
            walk(q, self)
        self.w(']')

    def visitListCompFor(self, node):
        self.w(' for ')
        walk(node.assign, self)
        self.w(' in ')
        walk(node.list, self)
        for expr in node.ifs:
            self.w(' if ')
            walk(expr, self)

    def visitName(self, node):
        self.w(node.name)

    def visitDiscard(self, node):
        walk(node.expr, self)
        self.nl()

    def visitPrintnl(self, node):
        self.w('print ')
        if node.dest:
            self.w('>>')
            walk(node.dest, self)
            self.w(', ')
        for e in node.nodes:
            walk(e, self)
        self.nl()

    def visitGetattr(self, node):
        walk(node.expr, self)
        self.w('.')
        self.w(node.attrname)

    def visitImport(self, node):
        self.w('import ')
        for (mod, as_) in node.names:
            self.w(mod)
            if as_ is not None:
                self.w(' as ')
                self.w(as_)
            self.w(', ')
        self.nl()

    def visitFrom(self, node):
        self.w('from ')
        self.w(node.modname)
        self.w(' import ')
        for (mod, as_) in node.names:
            self.w(mod)
            if as_ is not None:
                self.w(' as ')
                self.w(as_)
            self.w(', ')
        self.nl()

    def visitConst(self, node):
        self.w(repr(node.value))

    def visitReturn(self, node):
        self.w('return ')
        walk(node.value, self)
        self.nl()

    def visitClass(self, node):
        self.w('class ')
        self.w(node.name)
        if node.bases:
            self.w('(')
            for b in node.bases:
                walk(b, self)
                self.w(', ')
            self.w('):')
        self.indent()
        try:
            if node.doc is not None:
                self.w(repr(node.doc))
            walk(node.code, self)
        finally:
            self.dedent()

    def visitAssAttr(self, node):
        walk(node.expr, self)
        self.w('.')
        self.w(node.attrname)

    def visitMul(self, node):
        walk(node.left, self)
        self.w(' * ')
        walk(node.right, self)

    def visitSub(self, node):
        walk(node.left, self)
        self.w(' - ')
        walk(node.right, self)

    def visitAdd(self, node):
        walk(node.left, self)
        self.w(' + ')
        walk(node.right, self)

    def visitMod(self, node):
        walk(node.left, self)
        self.w(' % ')
        walk(node.right, self)

    def visitAugAssign(self, node):
        walk(node.node, self)
        self.w(' ')
        self.w(node.op)
        self.w(' ')
        walk(node.expr, self)
        self.nl()

    def visitIf(self, node):
        keyword = 'if'
        for (cond, body) in node.tests:
            self.w(keyword)
            self.w(' ')
            walk(cond, self)
            self.w(':')
            self.indent()
            try:
                walk(body, self)
            finally:
                self.dedent()
            keyword = 'elif'
        if node.else_:
            self.w('else:')
            self.indent()
            try:
                walk(node.else_, self)
            finally:
                self.dedent()

    def visitCompare(self, node):
        walk(node.expr, self)
        for (op, arg) in node.ops:
            self.w(' ')
            self.w(op)
            self.w(' ')
            walk(arg, self)

    def visitFor(self, node):
        self.w('for ')
        walk(node.assign, self)
        self.w(' in ')
        walk(node.list, self)
        self.w(':')
        self.indent()
        try:
            walk(node.body, self)
        finally:
            self.dedent()
        if node.else_:
            self.w('else:')
            self.indent()
            try:
                walk(node.else_, self)
            finally:
                self.dedent()

    def visitSlice(self, node):
        walk(node.expr, self)
        self.w('[')
        if node.lower:
            walk(node.lower, self)
        self.w(':')
        if node.upper:
            walk(node.upper, self)
        self.w(']')

    def visitTuple(self, node):
        self.w('(')
        if len(node.nodes) == 0:
            pass
        elif len(node.nodes) == 1:
            walk(node.nodes[0], self)
            self.w(',')
        else:
            for expr in node.nodes[:-1]:
                walk(expr, self)
                self.w(', ')
            walk(node.nodes[-1], self)
        self.w(')')

    def visitTryFinally(self, node):
        self.w('try:')
        self.indent()
        try:
            walk(node.body, self)
        finally:
            self.dedent()
        self.w('finally:')
        self.indent()
        try:
            walk(node.final, self)
        finally:
            self.dedent()

    def visitSubscript(self, node):
        walk(node.expr, self)
        self.w('[')
        walk(node.subs[0], self)
        self.w(']')

    def visitUnarySub(self, node):
        self.w('-')
        walk(node.expr, self)

    def visitAssTuple(self, node):
        self.w('(')
        for expr in node.nodes:
            walk(expr, self)
            self.w(', ')
        self.w(')')

    def visitRaise(self, node):
        self.w('raise ')
        walk(node.expr1, self)
        if node.expr2:
            self.w(', ')
            walk(node.expr2, self)
            if node.expr3:
                self.w(', ')
                walk(node.expr3, self)
        self.nl()

    def visitDict(self, node):
        self.w('{')
        for (k, v) in node.items:
            walk(k, self)
            self.w(':')
            walk(v, self)
            self.w(',')
        self.w('}')

    def __str__(self):
        return self.s.getvalue()

def pp(ast):
    """Convert C{ast} to Python source.

    @param ast: The node to render into Python source.
    """
    sw = SourceWriter()
    walk(ast, sw)
    return sw.s.getvalue()
