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