Armin Ronacher | e791c2a | 2008-04-07 18:39:54 +0200 | [diff] [blame] | 1 | # -*- coding: utf-8 -*- |
| 2 | """ |
| 3 | jinja2.visitor |
| 4 | ~~~~~~~~~~~~~~ |
| 5 | |
| 6 | This module implements a visitor for the nodes. |
| 7 | |
| 8 | :copyright: Copyright 2008 by Armin Ronacher. |
| 9 | :license: GNU GPL. |
| 10 | """ |
| 11 | from jinja2.nodes import Node |
| 12 | |
| 13 | |
| 14 | class NodeVisitor(object): |
Armin Ronacher | 26c0f51 | 2008-04-17 11:16:01 +0200 | [diff] [blame] | 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. |
Armin Ronacher | e791c2a | 2008-04-07 18:39:54 +0200 | [diff] [blame] | 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): |
Armin Ronacher | 963f97d | 2008-04-25 11:44:59 +0200 | [diff] [blame] | 27 | """Return the visitor function for this node or `None` if no visitor |
Armin Ronacher | e791c2a | 2008-04-07 18:39:54 +0200 | [diff] [blame] | 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): |
Armin Ronacher | 26c0f51 | 2008-04-17 11:16:01 +0200 | [diff] [blame] | 48 | """Walks the abstract syntax tree and allows modifications of nodes. |
Armin Ronacher | e791c2a | 2008-04-07 18:39:54 +0200 | [diff] [blame] | 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 |
Armin Ronacher | 4f7d2d5 | 2008-04-22 10:40:26 +0200 | [diff] [blame] | 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 |