def test_parallel_walk_inconsistent_trees(self):
node_1 = parser.parse_str(
textwrap.dedent("""
def f(a):
return a + 1
"""))
node_2 = parser.parse_str(
textwrap.dedent("""
def f(a):
return a + (a * 2)
"""))
node_3 = parser.parse_str(
textwrap.dedent("""
def f(a):
return a + 2
"""))
with self.assertRaises(ValueError):
for _ in ast_util.parallel_walk(node_1, node_2):
pass
# There is not particular reason to reject trees that differ only in the
# value of a constant.
# TODO(mdan): This should probably be allowed.
with self.assertRaises(ValueError):
for _ in ast_util.parallel_walk(node_1, node_3):
pass
def test_parallel_walk_inconsistent_trees(self):
node_1 = parser.parse_str(
textwrap.dedent("""
def f(a):
return a + 1
"""))
node_2 = parser.parse_str(
textwrap.dedent("""
def f(a):
return a + (a * 2)
"""))
node_3 = parser.parse_str(
textwrap.dedent("""
def f(a):
return a + 2
"""))
with self.assertRaises(ValueError):
for _ in ast_util.parallel_walk(node_1, node_2):
pass
# There is not particular reason to reject trees that differ only in the
# value of a constant.
# TODO(mdan): This should probably be allowed.
with self.assertRaises(ValueError):
for _ in ast_util.parallel_walk(node_1, node_3):
pass
def test_parse_str(self):
mod = parser.parse_str(
textwrap.dedent("""
def f(x):
return x + 1
"""))
self.assertEqual('f', mod.body[0].name)
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 = resolve(parser.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_create_source_map_multiple_nodes(self):
source = """
from __future__ import print_function
def test_fn(x):
return x + 1
"""
source = textwrap.dedent(source)
nodes = parser.parse_str(source, single_node=False)
fake_import_origin = origin_info.OriginInfo(
loc=origin_info.Location('fake_filename', 3, 7),
function_name='fake_function_name',
source_code_line='fake source line',
comment=None)
anno.setanno(nodes[0], anno.Basic.ORIGIN, fake_import_origin)
fake_function_origin = origin_info.OriginInfo(
loc=origin_info.Location('fake_filename', 3, 7),
function_name='fake_function_name',
source_code_line='fake source line',
comment=None)
anno.setanno(nodes[1], anno.Basic.ORIGIN, fake_function_origin)
source_map = origin_info.create_source_map(nodes, source,
'test_filename')
loc = origin_info.LineLocation('test_filename', 2)
self.assertIn(loc, source_map)
self.assertIs(source_map[loc], fake_import_origin)
loc = origin_info.LineLocation('test_filename', 3)
self.assertIn(loc, source_map)
self.assertIs(source_map[loc], fake_function_origin)
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 = parser.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_resolve(self):
source = """
def test_fn(x):
'''Docstring.'''
return x # comment
"""
source = textwrap.dedent(source)
node = parser.parse_str(source)
origin_info.resolve(node, source)
origin = anno.getanno(node, anno.Basic.ORIGIN)
self.assertEqual(origin.loc.lineno, 2)
self.assertEqual(origin.loc.col_offset, 0)
self.assertEqual(origin.source_code_line, 'def test_fn(x):')
self.assertIsNone(origin.comment)
origin = anno.getanno(node.body[0], anno.Basic.ORIGIN)
self.assertEqual(origin.loc.lineno, 3)
self.assertEqual(origin.loc.col_offset, 2)
self.assertEqual(origin.source_code_line, " '''Docstring.'''")
self.assertIsNone(origin.comment)
origin = anno.getanno(node.body[1], anno.Basic.ORIGIN)
self.assertEqual(origin.loc.lineno, 4)
self.assertEqual(origin.loc.col_offset, 2)
self.assertEqual(origin.source_code_line, ' return x # comment')
self.assertEqual(origin.comment, 'comment')
def test_create_source_map_multiple_nodes(self):
source = """
from __future__ import print_function
def test_fn(x):
return x + 1
"""
source = textwrap.dedent(source)
nodes = parser.parse_str(source, single_node=False)
fake_import_origin = origin_info.OriginInfo(
loc=origin_info.Location('fake_filename', 3, 7),
function_name='fake_function_name',
source_code_line='fake source line',
comment=None)
anno.setanno(nodes[0], anno.Basic.ORIGIN, fake_import_origin)
fake_function_origin = origin_info.OriginInfo(
loc=origin_info.Location('fake_filename', 3, 7),
function_name='fake_function_name',
source_code_line='fake source line',
comment=None)
anno.setanno(nodes[1], anno.Basic.ORIGIN, fake_function_origin)
source_map = origin_info.create_source_map(nodes, source, 'test_filename')
loc = origin_info.LineLocation('test_filename', 2)
self.assertIn(loc, source_map)
self.assertIs(source_map[loc], fake_import_origin)
loc = origin_info.LineLocation('test_filename', 3)
self.assertIn(loc, source_map)
self.assertIs(source_map[loc], fake_function_origin)
def test_resolve(self):
source = """
def test_fn(x):
'''Docstring.'''
return x # comment
"""
source = textwrap.dedent(source)
node = parser.parse_str(source)
origin_info.resolve(node, source)
origin = anno.getanno(node, anno.Basic.ORIGIN)
self.assertEqual(origin.loc.lineno, 2)
self.assertEqual(origin.loc.col_offset, 0)
self.assertEqual(origin.source_code_line, 'def test_fn(x):')
self.assertIsNone(origin.comment)
origin = anno.getanno(node.body[0], anno.Basic.ORIGIN)
self.assertEqual(origin.loc.lineno, 3)
self.assertEqual(origin.loc.col_offset, 2)
self.assertEqual(origin.source_code_line, " '''Docstring.'''")
self.assertIsNone(origin.comment)
origin = anno.getanno(node.body[1], anno.Basic.ORIGIN)
self.assertEqual(origin.loc.lineno, 4)
self.assertEqual(origin.loc.col_offset, 2)
self.assertEqual(origin.source_code_line, ' return x # comment')
self.assertEqual(origin.comment, 'comment')
def test_resolve(self):
source = """
def test_fn(x):
'''Docstring.'''
return x # comment
"""
source = textwrap.dedent(source)
node = parser.parse_str(source)
origin_info.resolve(node, source, 'test_file', 10, 10)
def_origin = anno.getanno(node, anno.Basic.ORIGIN)
self.assertEqual(def_origin.loc.filename, 'test_file')
self.assertEqual(def_origin.loc.lineno, 10)
self.assertEqual(def_origin.loc.col_offset, 10)
self.assertEqual(def_origin.source_code_line, 'def test_fn(x):')
self.assertIsNone(def_origin.comment)
docstring_origin = anno.getanno(node.body[0], anno.Basic.ORIGIN)
self.assertEqual(def_origin.loc.filename, 'test_file')
self.assertEqual(docstring_origin.loc.lineno, 11)
self.assertEqual(docstring_origin.loc.col_offset, 12)
self.assertEqual(docstring_origin.source_code_line,
" '''Docstring.'''")
self.assertIsNone(docstring_origin.comment)
ret_origin = anno.getanno(node.body[1], anno.Basic.ORIGIN)
self.assertEqual(def_origin.loc.filename, 'test_file')
self.assertEqual(ret_origin.loc.lineno, 12)
self.assertEqual(ret_origin.loc.col_offset, 12)
self.assertEqual(ret_origin.source_code_line, ' return x # comment')
self.assertEqual(ret_origin.comment, 'comment')
def test_parse_str(self):
mod = parser.parse_str(
textwrap.dedent("""
def f(x):
return x + 1
"""))
self.assertEqual('f', mod.body[0].name)
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 = parser.parse_str(textwrap.dedent(samples), single_node=False)
nodes = tuple(resolve(node).value for node in nodes)
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 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 = parser.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_find_matching_definitions_lambda(self):
node = parser.parse_str(
textwrap.dedent("""
f = lambda x: 1
"""))
f = lambda x: x
nodes = ast_util.find_matching_definitions(node, f)
self.assertLambdaNodes(nodes, ('1',))
def test_parallel_walk(self):
node = parser.parse_str(
textwrap.dedent("""
def f(a):
return a + 1
"""))
for child_a, child_b in ast_util.parallel_walk(node, node):
self.assertEqual(child_a, child_b)
def test_find_matching_definitions_lambda_multiple_matches(self):
node = parser.parse_str(
textwrap.dedent("""
f = lambda x: 1, lambda x: 2
"""))
f = lambda x: x
nodes = ast_util.find_matching_definitions(node, f)
self.assertLambdaNodes(nodes, ('(1)', '(2)'))
def test_find_matching_definitions_lambda_multiple_matches(self):
node = parser.parse_str(
textwrap.dedent("""
f = lambda x: 1, lambda x: 2
"""))
f = lambda x: x
nodes = ast_util.find_matching_definitions(node, f)
self.assertLambdaNodes(nodes, ('(1)', '(2)'))
def test_parallel_walk_string_leaves(self):
src = """
def f(a):
global g
"""
node = parser.parse_str(textwrap.dedent(src))
for child_a, child_b in ast_util.parallel_walk(node, node):
self.assertEqual(child_a, child_b)
def test_parallel_walk(self):
node = parser.parse_str(
textwrap.dedent("""
def f(a):
return a + 1
"""))
for child_a, child_b in ast_util.parallel_walk(node, node):
self.assertEqual(child_a, child_b)
def test_to_code_basic(self):
def test_fn(x, s):
while tf.reduce_sum(x) > s:
x /= 2
return x
# Just check that the output is parseable Python code.
self.assertIsNotNone(parser.parse_str(api.to_code(test_fn)))
def test_to_code_basic(self):
def test_fn(x, s):
while tf.reduce_sum(x) > s:
x /= 2
return x
# Just check that the output is parseable Python code.
self.assertIsNotNone(parser.parse_str(api.to_code(test_fn)))
def test_apply_to_single_assignments_static_unpack(self):
node = parser.parse_str('a, b, c = d, e, f')
ast_util.apply_to_single_assignments(node.targets, node.value,
self._mock_apply_fn)
self.assertDictEqual(self._invocation_counts, {
('a', 'd'): 1,
('b', 'e'): 1,
('c', 'f'): 1,
})
def test_rename_symbols_attributes(self):
node = parser.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 = compiler.ast_to_source(node)
self.assertEqual(source.strip(), 'renamed_b_c = renamed_b_c.d')
def test_keywords_to_dict(self):
keywords = parser.parse_expression('f(a=b, c=1, d=\'e\')').keywords
d = ast_util.keywords_to_dict(keywords)
# Make sure we generate a usable dict node by attaching it to a variable and
# compiling everything.
node = parser.parse_str('def f(b): pass').body[0]
node.body.append(ast.Return(d))
result, _ = compiler.ast_to_object(node)
self.assertDictEqual(result.f(3), {'a': 3, 'c': 1, 'd': 'e'})
def test_keywords_to_dict(self):
keywords = parser.parse_expression('f(a=b, c=1, d=\'e\')').keywords
d = ast_util.keywords_to_dict(keywords)
# Make sure we generate a usable dict node by attaching it to a variable and
# compiling everything.
node = parser.parse_str('def f(b): pass').body[0]
node.body.append(ast.Return(d))
result, _ = compiler.ast_to_object(node)
self.assertDictEqual(result.f(3), {'a': 3, 'c': 1, 'd': 'e'})
def test_rename_symbols_attributes(self):
node = parser.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 = compiler.ast_to_source(node)
self.assertEqual(source.strip(), 'renamed_b_c = renamed_b_c.d')
def test_to_code_basic(self):
def test_fn(x, s):
while tf.reduce_sum(x) > s:
x /= 2
return x
compiled_code = api.to_code(test_fn)
# Just check that it is parseable Python code.
self.assertIsNotNone(parser.parse_str(compiled_code))
def test_rename_symbols_basic(self):
node = parser.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 = compiler.ast_to_source(node)
self.assertEqual(source.strip(), 'renamed_a + b')
def test_copy_clean(self):
node = parser.parse_str(
textwrap.dedent("""
def f(a):
return a + 1
"""))
setattr(node, '__foo', 'bar')
new_node = ast_util.copy_clean(node)
self.assertIsNot(new_node, node)
self.assertFalse(hasattr(new_node, '__foo'))
def test_rename_symbols_annotations(self):
node = parser.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 = parser.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 = compiler.ast_to_source(node)
self.assertEqual(source.strip(), 'renamed_a + b')
def test_rename_symbols_annotations(self):
node = parser.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_apply_to_single_assignments_static_unpack(self):
node = parser.parse_str('a, b, c = d, e, f')
node = node.body[0]
ast_util.apply_to_single_assignments(node.targets, node.value,
self._mock_apply_fn)
self.assertDictEqual(self._invocation_counts, {
('a', 'd'): 1,
('b', 'e'): 1,
('c', 'f'): 1,
})
def test_apply_to_single_assignments_dynamic_unpack(self):
node = parser.parse_str('a, b, c = d')
node = node.body[0]
ast_util.apply_to_single_assignments(node.targets, node.value,
self._mock_apply_fn)
self.assertDictEqual(self._invocation_counts, {
('a', 'd[0]'): 1,
('b', 'd[1]'): 1,
('c', 'd[2]'): 1,
})
def test_rename_symbols_basic(self):
node = parser.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.value.left.id, str)
source = parser.unparse(node, include_encoding_marker=False)
self.assertEqual(source.strip(), '(renamed_a + b)')
def test_source_map(self):
def test_fn(x):
if x > 0:
x += 1
return x
node, source = parser.parse_entity(test_fn)
fn_node = node.body[0]
origin_info.resolve(fn_node, source)
# Insert a traced line.
new_node = parser.parse_str('x = abs(x)').body[0]
anno.copyanno(fn_node.body[0], new_node, anno.Basic.ORIGIN)
fn_node.body.insert(0, new_node)
# Insert an untraced line.
fn_node.body.insert(0, parser.parse_str('x = 0').body[0])
modified_source = compiler.ast_to_source(fn_node)
source_map = origin_info.source_map(fn_node, modified_source,
'test_filename', [0])
loc = origin_info.LineLocation('test_filename', 1)
origin = source_map[loc]
self.assertEqual(origin.source_code_line, 'def test_fn(x):')
self.assertEqual(origin.loc.lineno, 1)
# The untraced line, inserted second.
loc = origin_info.LineLocation('test_filename', 2)
self.assertFalse(loc in source_map)
# The traced line, inserted first.
loc = origin_info.LineLocation('test_filename', 3)
origin = source_map[loc]
self.assertEqual(origin.source_code_line, ' if x > 0:')
self.assertEqual(origin.loc.lineno, 2)
loc = origin_info.LineLocation('test_filename', 4)
origin = source_map[loc]
self.assertEqual(origin.source_code_line, ' if x > 0:')
self.assertEqual(origin.loc.lineno, 2)
def test_copy_clean_preserves_annotations(self):
node = parser.parse_str(
textwrap.dedent("""
def f(a):
return a + 1
"""))
anno.setanno(node.body[0], 'foo', 'bar')
anno.setanno(node.body[0], 'baz', 1)
new_node = ast_util.copy_clean(node, preserve_annos={'foo'})
self.assertEqual(anno.getanno(new_node.body[0], 'foo'), 'bar')
self.assertFalse(anno.hasanno(new_node.body[0], 'baz'))
def test_to_code_basic(self):
def test_fn(x, s):
while tf.reduce_sum(x) > s:
x /= 2
return x
compiled_code = api.to_code(test_fn)
# Just check that it is parseable Python code.
self.assertIsNotNone(parser.parse_str(compiled_code))
def test_copy_clean(self):
node = parser.parse_str(
textwrap.dedent("""
def f(a):
return a + 1
"""))
setattr(node.body[0], '__foo', 'bar')
new_node = ast_util.copy_clean(node)
self.assertIsNot(new_node, node)
self.assertIsNot(new_node.body[0], node.body[0])
self.assertFalse(hasattr(new_node.body[0], '__foo'))
def test_copy_clean_preserves_annotations(self):
node = parser.parse_str(
textwrap.dedent("""
def f(a):
return a + 1
"""))
anno.setanno(node.body[0], 'foo', 'bar')
anno.setanno(node.body[0], 'baz', 1)
new_node = ast_util.copy_clean(node, preserve_annos={'foo'})
self.assertEqual(anno.getanno(new_node.body[0], 'foo'), 'bar')
self.assertFalse(anno.hasanno(new_node.body[0], 'baz'))
def test_find_matching_definitions_lambda_uses_arg_names(self):
node = parser.parse_str(
textwrap.dedent("""
f = lambda x: 1, lambda y: 2
"""))
f = lambda x: x
nodes = ast_util.find_matching_definitions(node, f)
self.assertLambdaNodes(nodes, ('(1)',))
f = lambda y: y
nodes = ast_util.find_matching_definitions(node, f)
self.assertLambdaNodes(nodes, ('(2)',))
def test_find_matching_definitions_lambda_uses_arg_names(self):
node = parser.parse_str(
textwrap.dedent("""
f = lambda x: 1, lambda y: 2
"""))
f = lambda x: x
nodes = ast_util.find_matching_definitions(node, f)
self.assertLambdaNodes(nodes, ('(1)',))
f = lambda y: y
nodes = ast_util.find_matching_definitions(node, f)
self.assertLambdaNodes(nodes, ('(2)',))
def test_find_matching_definitions_function(self):
node = parser.parse_str(
textwrap.dedent("""
def f(x):
return 1
"""))
def f(x):
return x
nodes = ast_util.find_matching_definitions(node, f)
self.assertFunctionDefNodes(nodes, ('return 1', ))
def test_find_matching_definitions_function(self):
node = parser.parse_str(
textwrap.dedent("""
def f(x):
return 1
"""))
def f(x):
return x
nodes = ast_util.find_matching_definitions(node, f)
self.assertFunctionDefNodes(nodes, ('return 1',))
def test_source_map_no_origin(self):
source = """
def test_fn(x):
return x + 1
"""
source = textwrap.dedent(source)
node = parser.parse_str(source)
source_map = origin_info.create_source_map(node, source, 'test_filename')
self.assertEmpty(source_map)
def test_find_matching_definitions_decorated_compatible(self):
node = parser.parse_str(
textwrap.dedent("""
@sneaky_decorator
def f(x, *args, **kwargs):
return 1
"""))
def f(a, b, c, d=1):
return a + b + c + d
nodes = ast_util.find_matching_definitions(node, f)
self.assertFunctionDefNodes(nodes, ('return 1', ))
def test_find_matching_definitions_decorated_compatible(self):
node = parser.parse_str(
textwrap.dedent("""
@sneaky_decorator
def f(x, *args, **kwargs):
return 1
"""))
def f(a, b, c, d=1):
return a + b + c + d
nodes = ast_util.find_matching_definitions(node, f)
self.assertFunctionDefNodes(nodes, ('return 1',))
def create_source_map(nodes, code, filename, indices_in_code):
"""Creates a source map between an annotated AST and the code it compiles to.
Args:
nodes: Iterable[ast.AST, ...]
code: Text
filename: Optional[Text]
indices_in_code: Union[int, Iterable[int, ...]], the positions at which
nodes appear in code. The parser always returns a module when parsing
code. This argument indicates the position in that module's body at
which the corresponding of node should appear.
Returns:
Dict[CodeLocation, OriginInfo], mapping locations in code to locations
indicated by origin annotations in node.
"""
reparsed_nodes = parser.parse_str(code)
reparsed_nodes = [reparsed_nodes.body[i] for i in indices_in_code]
resolve(reparsed_nodes, code)
result = {}
for before, after in ast_util.parallel_walk(nodes, reparsed_nodes):
# Note: generated code might not be mapped back to its origin.
# TODO(mdan): Generated code should always be mapped to something.
origin_info = anno.getanno(before, anno.Basic.ORIGIN, default=None)
final_info = anno.getanno(after, anno.Basic.ORIGIN, default=None)
if origin_info is None or final_info is None:
continue
line_loc = LineLocation(filename, final_info.loc.lineno)
existing_origin = result.get(line_loc)
if existing_origin is not None:
# Overlaps may exist because of child nodes, but almost never to
# different line locations. Exception make decorated functions, where
# both lines are mapped to the same line in the AST.
# Line overlaps: keep bottom node.
if existing_origin.loc.line_loc == origin_info.loc.line_loc:
if existing_origin.loc.lineno >= origin_info.loc.lineno:
continue
# In case of overlaps, keep the leftmost node.
if existing_origin.loc.col_offset <= origin_info.loc.col_offset:
continue
result[line_loc] = origin_info
return result
def test_find_matching_definitions_nested_functions_same_name(self):
node = parser.parse_str(
textwrap.dedent("""
def f(x, *args, **kwargs):
def f(x, y):
return 1
return 2
"""))
def f(x, y):
return x + y
nodes = ast_util.find_matching_definitions(node, f)
self.assertFunctionDefNodes(nodes, ('return 1', ))
def test_source_map_no_origin(self):
source = """
def test_fn(x):
return x + 1
"""
source = textwrap.dedent(source)
node = parser.parse_str(source)
source_map = origin_info.create_source_map(node, source,
'test_filename')
self.assertEmpty(source_map)
def test_find_matching_definitions_nested_functions_same_name(self):
node = parser.parse_str(
textwrap.dedent("""
def f(x, *args, **kwargs):
def f(x, y):
return 1
return 2
"""))
def f(x, y):
return x + y
nodes = ast_util.find_matching_definitions(node, f)
self.assertFunctionDefNodes(nodes, ('return 1',))
def create_source_map(nodes, code, filename, indices_in_code):
"""Creates a source map between an annotated AST and the code it compiles to.
Args:
nodes: Iterable[ast.AST, ...]
code: Text
filename: Optional[Text]
indices_in_code: Union[int, Iterable[int, ...]], the positions at which
nodes appear in code. The parser always returns a module when parsing
code. This argument indicates the position in that module's body at
which the corresponding of node should appear.
Returns:
Dict[CodeLocation, OriginInfo], mapping locations in code to locations
indicated by origin annotations in node.
"""
reparsed_nodes = parser.parse_str(code)
reparsed_nodes = [reparsed_nodes.body[i] for i in indices_in_code]
resolve(reparsed_nodes, code)
result = {}
for before, after in ast_util.parallel_walk(nodes, reparsed_nodes):
# Note: generated code might not be mapped back to its origin.
# TODO(mdan): Generated code should always be mapped to something.
origin_info = anno.getanno(before, anno.Basic.ORIGIN, default=None)
final_info = anno.getanno(after, anno.Basic.ORIGIN, default=None)
if origin_info is None or final_info is None:
continue
line_loc = LineLocation(filename, final_info.loc.lineno)
existing_origin = result.get(line_loc)
if existing_origin is not None:
# Overlaps may exist because of child nodes, but almost never to
# different line locations. Exception make decorated functions, where
# both lines are mapped to the same line in the AST.
# Line overlaps: keep bottom node.
if existing_origin.loc.line_loc == origin_info.loc.line_loc:
if existing_origin.loc.lineno >= origin_info.loc.lineno:
continue
# In case of overlaps, keep the leftmost node.
if existing_origin.loc.col_offset <= origin_info.loc.col_offset:
continue
result[line_loc] = origin_info
return result
def test_function_calls(self):
samples = """
a.b
a.b()
a().b
z[i]
z[i]()
z()[i]
"""
nodes = resolve(parser.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_function_calls(self):
samples = """
a.b
a.b()
a().b
z[i]
z[i]()
z()[i]
"""
nodes = parser.parse_str(textwrap.dedent(samples), single_node=False)
nodes = tuple(resolve(node).value for node in nodes)
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 matches(node, pattern):
"""Basic pattern matcher for AST.
The pattern may contain wildcards represented by the symbol '_'. A node
matches a pattern if for every node in the tree, either there is a node of
the same type in pattern, or a Name node with id='_'.
Args:
node: ast.AST
pattern: ast.AST
Returns:
bool
"""
if isinstance(pattern, str):
pattern = parser.parse_str(pattern)
matcher = PatternMatcher(pattern)
matcher.visit(node)
return matcher.matches
def test_resolve(self):
samples = """
a
a.b
(c, d.e)
[f, (g.h.i)]
j(k, l)
"""
nodes = resolve(parser.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 test_create_source_map(self):
source = """
def test_fn(x):
return x + 1
"""
source = textwrap.dedent(source)
node = parser.parse_str(source)
fake_origin = origin_info.OriginInfo(
loc=origin_info.Location('fake_filename', 3, 7),
function_name='fake_function_name',
source_code_line='fake source line',
comment=None)
anno.setanno(node, anno.Basic.ORIGIN, fake_origin)
source_map = origin_info.create_source_map(node, source, 'test_filename')
loc = origin_info.LineLocation('test_filename', 2)
self.assertIn(loc, source_map)
self.assertIs(source_map[loc], fake_origin)
def create_source_map(nodes, code, filename):
"""Creates a source map between an annotated AST and the code it compiles to.
Args:
nodes: Iterable[ast.AST, ...]
code: Text
filename: Optional[Text]
Returns:
Dict[LineLocation, OriginInfo], mapping locations in code to locations
indicated by origin annotations in node.
"""
reparsed_nodes = parser.parse_str(code, preamble_len=0, single_node=False)
for node in reparsed_nodes:
resolve(node, code)
result = {}
try:
for before, after in ast_util.parallel_walk(nodes, reparsed_nodes):
# Note: generated code might not be mapped back to its origin.
# TODO(mdan): Generated code should always be mapped to something.
origin_info = anno.getanno(before, anno.Basic.ORIGIN, default=None)
final_info = anno.getanno(after, anno.Basic.ORIGIN, default=None)
if origin_info is None or final_info is None:
continue
line_loc = LineLocation(filename, final_info.loc.lineno)
existing_origin = result.get(line_loc)
if existing_origin is not None:
# Overlaps may exist because of child nodes, but almost never to
# different line locations. Exception make decorated functions, where
# both lines are mapped to the same line in the AST.
# Line overlaps: keep bottom node.
if existing_origin.loc.line_loc == origin_info.loc.line_loc:
if existing_origin.loc.lineno >= origin_info.loc.lineno:
continue
# In case of overlaps, keep the leftmost node.
if existing_origin.loc.col_offset <= origin_info.loc.col_offset:
continue
result[line_loc] = origin_info
except ValueError:
if logging.has_verbosity(3):
for n, rn in zip(nodes, reparsed_nodes):
nodes_str = pretty_printer.fmt(n, color=False, noanno=True)
reparsed_nodes_str = pretty_printer.fmt(rn, color=False, noanno=True)
diff = difflib.context_diff(
nodes_str.split('\n'),
reparsed_nodes_str.split('\n'),
fromfile='Original nodes',
tofile='Reparsed nodes',
n=7)
diff = '\n'.join(diff)
logging.log(3, 'AST seems to lack integrity. Diff:\n%s', diff)
raise
return result