"""Visitor classes for iterating over Python AST (Abstract Syntax Trees)

This is a visitor class that implements the visitor pattern to 'visit'
all the nodes in an Abstract Syntax Tree (AST) in either a preorder or
postorder syntax.  It can be used in two different ways.  The first is
to extend ASTVisitor so that it has it's own methods that follow the
pattern visitXXX where XXX can be any name of a valid class in an
AST.  In short, names such as visitFunction and visitIf are desired.

You can also use ASTVisitor as a "Walker" as the original author of
this module had intended.  You can define and instanciate any
arbitrary class that has methods following the pattern visitXXX
like above, and pass this as an arguement to the preorder and
postorder methods of ASTVisitor.

The function walk can be used as a shortcut to a preorder
traversal with a seperate Walker class.

The class ShowUnimplementedVisitor is a useful and practical example
of how to extend ASTVisitor as it's own Visitor class.  It does not
actually implement any visitXXX methods.  Rather the first time a
node of a type is encountered, it prints a quick message to standard
out, noting that the hereto mentioned node type visitXXX method is
not (yet ?) implemented. This can be useful to developers, as they
can see which Nodes are relevant to a chunk of Python code rather
than implementing every visitXXX method when coding.  Finally,
when the code is production ready, switching the parent node to
ASTVisitor will remove all 'unimplemented' messages transparently.

So while developing, your code could be:

class DementedVisitor(ShowUnimplementedVisitor):
    def visitFunction:
        pass


and when finished, it would be:

class DementedVisitor(ASTVisitor):
    def visitFunction:
        pass


thus ensuring all your demented visitors are quiet house guests
when being shown before society in general.

TODO: Implement a visitor that "pretty prints" out the nodes
similat to the form seen in the Python Library Reference,
as an independant Visitor class.
"""
from compiler import ast

# XXX should probably rename ASTVisitor to ASTWalker
# XXX can it be made even more generic?

class ASTVisitor(object):
    """Performs a depth-first walk of the AST

    The ASTVisitor will walk the AST, performing either a preorder or
    postorder traversal depending on which method is called.

    methods:
    preorder(tree, visitor)
    postorder(tree, visitor)
        tree: an instance of ast.Node
        visitor: an (optional) instance with visitXXX methods

    Additionally ASTVisitor maintains two callback lists, the
    "defaults" and the "others".  They are called when either
    processing a node and it's children, or when a visitXXX
    method for a node does not exist, respectively. The
    following methods are relevant:
    register_default(function)
    register_other(function)
        function: a callback to be put at the beginning of the list

    unregister_default(function)
    unregister_other(function)
        function: removes first instance of the function from the list

    (Keep in mind, in the odd situation where a function is meant to be
    called twice, it must be registered and unregistered twice as well.)

    The ASTVisitor is responsible for walking over the tree in the
    correct order.  For each node, it checks the visitor argument for
    a method named 'visitNodeType' where NodeType is the name of the
    node's class, e.g. Class.  If the method exists, it is called
    with the node as its sole argument.

    The visitor method for a particular node type can control how
    child nodes are visited during a preorder walk.  ASTVisitor
    has two constants, ASTVisitor.CONTINUE and ASTVisitor.NO_CONTINUE
    that can be returned to ASTVisitor by visit methods, in order to
    control whether it should in fact continue or not. (It can't control
    the order during a postorder walk, because it is called _after_
    the walk has occurred.)  
    """

    CONTINUE = True
    NO_CONTINUE = False

    PREORDER = True
    POSTORDER = False

    def __init__(self):
        self._cache = {}
        self._defaults = [self.default]
        self._others = [self.default]

    def register_default(self, function):
        self._defaults = [function] + self._defaults

    def unregister_default(self, function):
        self._defaults.remove(function)

    def register_other(self, function):
        self._others = [function] + self._others

    def unregister_other(self, function):
        self._others.remove(function)

    def _iterate_over_defaults(self, node, *args):
        for function in self._defaults:
            if function(node, *args) == self.NO_CONTINUE:
                return

    def _iterate_over_others(self, node, *args):
        for function in self._others:
            if function(node, *args) == self.NO_CONTINUE:
                return
        
    def default(self, node, *args):
        for child in node.getChildNodes():
            self.dispatch(child, *args)

    def dispatch(self, node, *args):
        klass = node.__class__
        meth = self._get_method_for_className(klass)
        if meth is None: #in this case, pre/post order don't matter much
            self._iterate_over_others(node, *args)
        else:
            if self_order == self.POSTORDER:
                self._iterate_over_defaults(node, *args)
                meth(node, *args)
            elif meth(node, *args) != self.NO_CONTINUE:
                self._iterate_over_defaults(node, *args)
            else:
                pass #the method said stop, and it's preorder

    def preorder(self, tree, visitor = None, *args):
        """Do preorder walk of tree using visitor pattern"""
        if visitor is None:
            visitor = self
        self.visitor = visitor
        self._order = self.PREORDER
        self.dispatch(tree, *args) # XXX *args make sense?

    def postorder(self, tree, visitor = None, *args):
        """Do postorder walk of tree using visitor pattern"""
        if visitor is None:
            visitor = self
        self.visitor = visitor
        self._order = self.POSTORDER
        self.dispatch(tree, *args) # XXX *args make sense?
    
    def _get_method_for_className(self, klass):
        meth = self._cache.get(klass, None)
        if meth is None:
            meth = getattr(self.visitor, self._get_visitor_name_for_class(klass), None)
            self._cache[klass] = meth
        return meth

    def _get_visitor_name_for_class(self, klass):
        return 'visit' + klass.__name__

class ShowUnimplementedVisitor(ASTVisitor):
    """Prints examples of the nodes that aren't visited

    This visitor-driver is only useful for development, when it's
    helpful to develop a visitor incrementally, and get feedback on what
    you still have to do.
    """
    def __init__(self):
        ASTVisitor.__init__(self)
        self.register_other(self.no_visitor)
        self.visited = {}

    def visited_p(self, klass):
        return self.visited.get(klass, False)

    def set_visited(self, klass):
        self.visited[klass] = True

    def no_visitor(self, node, *args):
        klass = node.__class__
        if self.visited_p(klass) == False:
            print "%s is not yet implemented" % self._get_visitor_name_for_class(klass)
            self.set_visited(klass)

##This code could be a little braindead all things considered, but needs
##how we want to do all this.
_walker = ASTVisitor
def walk(tree, visitor, verbose=None):
    if visitor is not None:
        _walker = visitor
    walker = _walker()
    if verbose is not None:
        walker.VERBOSE = verbose
    walker.preorder(tree)
    return walker

def dumpNode(node): ##Needed?
    print node.__class__
    for attr in dir(node):
        if attr[0] != '_':
            print "\t", "%-10.10s" % attr, getattr(node, attr)