def replace(template, **replacements): """Replaces placeholders in a Python template. AST Name and Tuple nodes always receive the context that inferred from the template. However, when replacing more complex nodes (that can potentially contain Name children), then the caller is responsible for setting the appropriate context. Args: template: A string representing Python code. Any symbol name can be used that appears in the template code can be used as placeholder. **replacements: A mapping from placeholder names to (lists of) AST nodes that these placeholders will be replaced by. String values are also supported as a shorthand for AST Name nodes with the respective ID. Returns: An AST node or list of AST nodes with the replacements made. If the template was a function, a list will be returned. If the template was a node, the same node will be returned. If the template was a string, an AST node will be returned (a `Module` node in the case of a multi-line string, an `Expr` node otherwise). Raises: ValueError: if the arguments are incorrect. """ if not isinstance(template, str): raise ValueError('Expected string template, got %s' % type(template)) tree = parsing.parse_str(textwrap.dedent(template)) for k in replacements: replacements[k] = _convert_to_ast(replacements[k]) results = ReplaceTransformer(replacements).visit(tree).body if isinstance(results, list): return [qual_names.resolve(r) for r in results] return qual_names.resolve(results)
def test_subscript_resolve(self): samples = """ x[i] x[i.b] a.b[c] a.b[x.y] a[z[c]] a[b[c[d]]] a[b].c a.b.c[d].e.f a.b[c[d]].e.f a.b[c[d.e.f].g].h """ nodes = qual_names.resolve(parsing.parse_str(textwrap.dedent(samples))) nodes = tuple(n.value for n in nodes.body) self.assertQNStringIs(nodes[0], 'x[i]') self.assertQNStringIs(nodes[1], 'x[i.b]') self.assertQNStringIs(nodes[2], 'a.b[c]') self.assertQNStringIs(nodes[3], 'a.b[x.y]') self.assertQNStringIs(nodes[4], 'a[z[c]]') self.assertQNStringIs(nodes[5], 'a[b[c[d]]]') self.assertQNStringIs(nodes[6], 'a[b].c') self.assertQNStringIs(nodes[7], 'a.b.c[d].e.f') self.assertQNStringIs(nodes[8], 'a.b[c[d]].e.f') self.assertQNStringIs(nodes[9], 'a.b[c[d.e.f].g].h')
def test_rename_symbols_attributes(self): node = parsing.parse_str('b.c = b.c.d') node = qual_names.resolve(node) node = ast_util.rename_symbols( node, {qual_names.from_str('b.c'): qual_names.QN('renamed_b_c')}) source = parsing.ast_to_source(node) self.assertEqual(source.strip(), 'renamed_b_c = renamed_b_c.d')
def test_rename_symbols_annotations(self): node = parsing.parse_str('a[i]') node = qual_names.resolve(node) anno.setanno(node, 'foo', 'bar') orig_anno = anno.getanno(node, 'foo') node = ast_util.rename_symbols(node, {qual_names.QN('a'): qual_names.QN('b')}) self.assertIs(anno.getanno(node, 'foo'), orig_anno)
def test_rename_symbols_basic(self): node = parsing.parse_str('a + b') node = qual_names.resolve(node) node = ast_util.rename_symbols( node, {qual_names.QN('a'): qual_names.QN('renamed_a')}) self.assertIsInstance(node.body[0].value.left.id, str) source = parsing.ast_to_source(node) self.assertEqual(source.strip(), 'renamed_a + b')
def _parse_and_analyze(self, test_fn): node, source = parsing.parse_entity(test_fn) entity_info = transformer.EntityInfo( source_code=source, source_file=None, namespace={}, arg_values=None, arg_types=None, owner_type=None) node = qual_names.resolve(node) ctx = transformer.Context(entity_info) node = activity.resolve(node, ctx) return node, entity_info
def replace_as_expression(template, **replacements): """Variant of replace that generates expressions, instead of code blocks.""" replacement = replace(template, **replacements) if len(replacement) != 1: raise ValueError( 'single expression expected; for more general templates use replace') node = replacement[0] node = qual_names.resolve(node) if isinstance(node, gast.Expr): return node.value elif isinstance(node, gast.Name): return node raise ValueError( 'the template is expected to generate an expression or a name node;' ' instead found %s' % node)
def test_function_calls(self): samples = """ a.b a.b() a().b z[i] z[i]() z()[i] """ nodes = qual_names.resolve(parsing.parse_str(textwrap.dedent(samples))) nodes = tuple(n.value for n in nodes.body) self.assertQNStringIs(nodes[0], 'a.b') self.assertQNStringIs(nodes[1].func, 'a.b') self.assertQNStringIs(nodes[2].value.func, 'a') self.assertQNStringIs(nodes[3], 'z[i]') self.assertQNStringIs(nodes[4].func, 'z[i]') self.assertQNStringIs(nodes[5].value.func, 'z')
def test_resolve(self): samples = """ a a.b (c, d.e) [f, (g.h.i)] j(k, l) """ nodes = qual_names.resolve(parsing.parse_str(textwrap.dedent(samples))) nodes = tuple(n.value for n in nodes.body) self.assertQNStringIs(nodes[0], 'a') self.assertQNStringIs(nodes[1], 'a.b') self.assertQNStringIs(nodes[2].elts[0], 'c') self.assertQNStringIs(nodes[2].elts[1], 'd.e') self.assertQNStringIs(nodes[3].elts[0], 'f') self.assertQNStringIs(nodes[3].elts[1], 'g.h.i') self.assertQNStringIs(nodes[4].func, 'j') self.assertQNStringIs(nodes[4].args[0], 'k') self.assertQNStringIs(nodes[4].args[1], 'l')
def transform(node, ctx, overload): node = qual_names.resolve(node) node = activity.resolve(node, ctx, parent_scope=None, overload=overload) node = ControlFlowTransformer(ctx, overload).visit(node) return node