1# -*- coding: utf-8 -*-
2"""API for traversing the AST nodes. Implemented by the compiler and
3meta introspection.
4"""
5from .nodes import Node
6
7
8class NodeVisitor(object):
9    """Walks the abstract syntax tree and call visitor functions for every
10    node found.  The visitor functions may return values which will be
11    forwarded by the `visit` method.
12
13    Per default the visitor functions for the nodes are ``'visit_'`` +
14    class name of the node.  So a `TryFinally` node visit function would
15    be `visit_TryFinally`.  This behavior can be changed by overriding
16    the `get_visitor` function.  If no visitor function exists for a node
17    (return value `None`) the `generic_visit` visitor is used instead.
18    """
19
20    def get_visitor(self, node):
21        """Return the visitor function for this node or `None` if no visitor
22        exists for this node.  In that case the generic visit function is
23        used instead.
24        """
25        method = "visit_" + node.__class__.__name__
26        return getattr(self, method, None)
27
28    def visit(self, node, *args, **kwargs):
29        """Visit a node."""
30        f = self.get_visitor(node)
31        if f is not None:
32            return f(node, *args, **kwargs)
33        return self.generic_visit(node, *args, **kwargs)
34
35    def generic_visit(self, node, *args, **kwargs):
36        """Called if no explicit visitor function exists for a node."""
37        for node in node.iter_child_nodes():
38            self.visit(node, *args, **kwargs)
39
40
41class NodeTransformer(NodeVisitor):
42    """Walks the abstract syntax tree and allows modifications of nodes.
43
44    The `NodeTransformer` will walk the AST and use the return value of the
45    visitor functions to replace or remove the old node.  If the return
46    value of the visitor function is `None` the node will be removed
47    from the previous location otherwise it's replaced with the return
48    value.  The return value may be the original node in which case no
49    replacement takes place.
50    """
51
52    def generic_visit(self, node, *args, **kwargs):
53        for field, old_value in node.iter_fields():
54            if isinstance(old_value, list):
55                new_values = []
56                for value in old_value:
57                    if isinstance(value, Node):
58                        value = self.visit(value, *args, **kwargs)
59                        if value is None:
60                            continue
61                        elif not isinstance(value, Node):
62                            new_values.extend(value)
63                            continue
64                    new_values.append(value)
65                old_value[:] = new_values
66            elif isinstance(old_value, Node):
67                new_node = self.visit(old_value, *args, **kwargs)
68                if new_node is None:
69                    delattr(node, field)
70                else:
71                    setattr(node, field, new_node)
72        return node
73
74    def visit_list(self, node, *args, **kwargs):
75        """As transformers may return lists in some places this method
76        can be used to enforce a list as return value.
77        """
78        rv = self.visit(node, *args, **kwargs)
79        if not isinstance(rv, list):
80            rv = [rv]
81        return rv
82