Replace AST nodes.
| 106 | |
| 107 | |
| 108 | class ReplaceTransformer(gast.NodeTransformer): |
| 109 | """Replace AST nodes.""" |
| 110 | |
| 111 | def __init__(self, replacements): |
| 112 | """Create a new ReplaceTransformer. |
| 113 | |
| 114 | Args: |
| 115 | replacements: A mapping from placeholder names to (lists of) AST nodes |
| 116 | that these placeholders will be replaced by. |
| 117 | """ |
| 118 | self.replacements = replacements |
| 119 | self.in_replacements = False |
| 120 | self.preserved_annos = { |
| 121 | anno.Basic.ORIGIN, |
| 122 | anno.Basic.SKIP_PROCESSING, |
| 123 | anno.Static.ORIG_DEFINITIONS, |
| 124 | 'extra_test', |
| 125 | 'function_context_name', |
| 126 | } |
| 127 | |
| 128 | def _prepare_replacement(self, replaced, key): |
| 129 | """Prepares a replacement AST that's safe to swap in for a node. |
| 130 | |
| 131 | Args: |
| 132 | replaced: ast.AST, the node being replaced |
| 133 | key: Hashable, the key of the replacement AST |
| 134 | Returns: |
| 135 | ast.AST, the replacement AST |
| 136 | """ |
| 137 | repl = self.replacements[key] |
| 138 | |
| 139 | new_nodes = ast_util.copy_clean(repl, preserve_annos=self.preserved_annos) |
| 140 | if isinstance(new_nodes, gast.AST): |
| 141 | new_nodes = [new_nodes] |
| 142 | |
| 143 | return new_nodes |
| 144 | |
| 145 | def visit_Expr(self, node): |
| 146 | # When replacing a placeholder with an entire statement, the replacement |
| 147 | # must stand on its own and not be wrapped in an Expr. |
| 148 | new_value = self.visit(node.value) |
| 149 | if new_value is node.value: |
| 150 | return node |
| 151 | return new_value |
| 152 | |
| 153 | def visit_keyword(self, node): |
| 154 | if node.arg not in self.replacements: |
| 155 | return self.generic_visit(node) |
| 156 | |
| 157 | repl = self._prepare_replacement(node, node.arg) |
| 158 | if isinstance(repl, gast.keyword): |
| 159 | return repl |
| 160 | elif (repl and isinstance(repl, (list, tuple)) and |
| 161 | all(isinstance(r, gast.keyword) for r in repl)): |
| 162 | return repl |
| 163 | # TODO(mdan): We may allow replacing with a string as well. |
| 164 | # For example, if one wanted to replace foo with bar in foo=baz, then |
| 165 | # we could allow changing just node arg, so that we end up with bar=baz. |