# -*- coding: utf-8 -*- """ Part of the astor library for Python AST manipulation. License: 3-clause BSD Copyright 2012 (c) Patrick Maupin Copyright 2013 (c) Berker Peksag This file contains a TreeWalk class that views a node tree as a unified whole and allows several modes of traversal. """ from .node_util import iter_node class MetaFlatten(type): """This metaclass is used to flatten classes to remove class hierarchy. This makes it easier to manipulate classes (find attributes in a single dict, etc.) """ def __new__(clstype, name, bases, clsdict): newbases = (object,) newdict = {} for base in reversed(bases): if base not in newbases: newdict.update(vars(base)) newdict.update(clsdict) # Delegate the real work to type return type.__new__(clstype, name, newbases, newdict) MetaFlatten = MetaFlatten('MetaFlatten', (object,), {}) class TreeWalk(MetaFlatten): """The TreeWalk class can be used as a superclass in order to walk an AST or similar tree. Unlike other treewalkers, this class can walk a tree either recursively or non-recursively. Subclasses can define methods with the following signatures:: def pre_xxx(self): pass def post_xxx(self): pass def init_xxx(self): pass Where 'xxx' is one of: - A class name - An attribute member name concatenated with '_name' For example, 'pre_targets_name' will process nodes that are referenced by the name 'targets' in their parent's node. - An attribute member name concatenated with '_item' For example, 'pre_targets_item' will process nodes that are in a list that is the targets attribute of some node. pre_xxx will process a node before processing any of its subnodes. if the return value from pre_xxx evalates to true, then walk will not process any of the subnodes. Those can be manually processed, if desired, by calling self.walk(node) on the subnodes before returning True. post_xxx will process a node after processing all its subnodes. init_xxx methods can decorate the class instance with subclass-specific information. A single init_whatever method could be written, but to make it easy to keep initialization with use, any number of init_xxx methods can be written. They will be called in alphabetical order. """ def __init__(self, node=None): self.nodestack = [] self.setup() if node is not None: self.walk(node) def setup(self): """All the node-specific handlers are setup at object initialization time. """ self.pre_handlers = pre_handlers = {} self.post_handlers = post_handlers = {} for name in sorted(vars(type(self))): if name.startswith('init_'): getattr(self, name)() elif name.startswith('pre_'): pre_handlers[name[4:]] = getattr(self, name) elif name.startswith('post_'): post_handlers[name[5:]] = getattr(self, name) def walk(self, node, name='', list=list, len=len, type=type): """Walk the tree starting at a given node. Maintain a stack of nodes. """ pre_handlers = self.pre_handlers.get post_handlers = self.post_handlers.get nodestack = self.nodestack emptystack = len(nodestack) append, pop = nodestack.append, nodestack.pop append([node, name, list(iter_node(node, name + '_item')), -1]) while len(nodestack) > emptystack: node, name, subnodes, index = nodestack[-1] if index >= len(subnodes): handler = (post_handlers(type(node).__name__) or post_handlers(name + '_name')) if handler is None: pop() continue self.cur_node = node self.cur_name = name handler() current = nodestack and nodestack[-1] popstack = current and current[0] is node if popstack and current[-1] >= len(current[-2]): pop() continue nodestack[-1][-1] = index + 1 if index < 0: handler = (pre_handlers(type(node).__name__) or pre_handlers(name + '_name')) if handler is not None: self.cur_node = node self.cur_name = name if handler(): pop() else: node, name = subnodes[index] append([node, name, list(iter_node(node, name + '_item')), -1]) @property def parent(self): """Return the parent node of the current node.""" nodestack = self.nodestack if len(nodestack) < 2: return None return nodestack[-2][0] @property def parent_name(self): """Return the parent node and name.""" nodestack = self.nodestack if len(nodestack) < 2: return None return nodestack[-2][:2] def replace(self, new_node): """Replace a node after first checking integrity of node stack.""" cur_node = self.cur_node nodestack = self.nodestack cur = nodestack.pop() prev = nodestack[-1] index = prev[-1] - 1 oldnode, name = prev[-2][index] assert cur[0] is cur_node is oldnode, (cur[0], cur_node, prev[-2], index) parent = prev[0] if isinstance(parent, list): parent[index] = new_node else: setattr(parent, name, new_node)