def visit_Attribute(self, node):
node = self.generic_visit(node)
if anno.hasanno(node.value, anno.Basic.QN):
anno.setanno(
node, anno.Basic.QN,
QN(anno.getanno(node.value, anno.Basic.QN), attr=node.attr))
return node
def visit_Print(self, node):
self._enter_scope(False)
node.values = self.visit_block(node.values)
anno.setanno(node, anno.Static.SCOPE, self.scope)
anno.setanno(node, anno.Static.ARGS_SCOPE, self.scope)
self._exit_scope()
return node
def visit_FunctionDef(self, node):
# The FunctionDef node itself has a Scope object that tracks the creation
# of its name, along with the usage of any decorator accompanying it.
self._enter_scope(False)
node.decorator_list = self.visit_block(node.decorator_list)
self.scope.mark_modified(qual_names.QN(node.name))
anno.setanno(node, anno.Static.SCOPE, self.scope)
self._exit_scope()
# A separate Scope tracks the actual function definition.
self._enter_scope(True)
assert not (self._in_function_def_args or self.state[_Lambda].level)
self._in_function_def_args = True
node.args = self.visit(node.args)
self._in_function_def_args = False
# Track the body separately. This is for compatibility reasons, it may not
# be strictly needed.
self._enter_scope(False)
node.body = self.visit_block(node.body)
anno.setanno(node, anno.Static.BODY_SCOPE, self.scope)
self._exit_scope()
self._exit_scope()
return node
def visit_While(self, node):
self._enter_scope(False)
node.test = self.visit(node.test)
anno.setanno(node, anno.Static.COND_SCOPE, self.scope)
anno.setanno(node.test, anno.Static.SCOPE, self.scope)
self._exit_scope()
node = self._process_parallel_blocks(
node, ((node.body, anno.Static.BODY_SCOPE),
(node.orelse, anno.Static.ORELSE_SCOPE)))
return node
def test_copy(self):
node_1 = ast.Name()
anno.setanno(node_1, 'foo', 3)
node_2 = ast.Name()
anno.copyanno(node_1, node_2, 'foo')
anno.copyanno(node_1, node_2, 'bar')
self.assertTrue(anno.hasanno(node_2, 'foo'))
self.assertFalse(anno.hasanno(node_2, 'bar'))
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 resolve(nodes, source, function=None):
"""Adds an origin information to all nodes inside the body of function.
Args:
nodes: Union[ast.AST, Iterable[ast.AST, ...]]
source: Text, the source code string for the function whose body nodes will
be annotated.
function: Callable, the function that will have all nodes inside of it
annotation with an OriginInfo annotation with key anno.Basic.ORIGIN. If
it is None then only the line numbers and column offset will be set in the
annotation, with the rest of the information being None.
Returns:
A tuple of the AST node for function and a String containing its source
code.
"""
if not isinstance(nodes, (list, tuple)):
nodes = (nodes,)
if function:
_, function_lineno = inspect.getsourcelines(function)
function_filepath = inspect.getsourcefile(function)
else:
function_lineno = None
function_filepath = None
# TODO(mdanatg): Pull this to a separate utility.
code_reader = six.StringIO(source)
comment_map = {}
for token in tokenize.generate_tokens(code_reader.readline):
tok_type, tok_string, loc, _, _ = token
srow, _ = loc
if tok_type == tokenize.COMMENT:
comment_map[srow] = tok_string.strip()[1:].strip()
source_lines = source.split('\n')
for node in nodes:
for n in gast.walk(node):
if not hasattr(n, 'lineno'):
continue
lineno_in_body = n.lineno
source_code_line = source_lines[lineno_in_body - 1]
if function:
source_lineno = function_lineno + lineno_in_body
function_name = function.__name__
else:
source_lineno = lineno_in_body
function_name = None
location = Location(function_filepath, source_lineno, n.col_offset)
origin = OriginInfo(location, function_name, source_code_line,
comment_map.get(source_lineno))
anno.setanno(n, anno.Basic.ORIGIN, origin)
def test_copy_clean_preserves_annotations(self):
node = parsing.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 visit_For(self, node):
self._enter_scope(False)
if isinstance(node.target, gast.Name):
node.target = self.visit(node.target)
else:
for target in node.target:
node.target = self.visit(target)
node.iter = self.visit(node.iter)
anno.setanno(node.iter, anno.Static.SCOPE, self.scope)
self._exit_scope()
node = self._process_parallel_blocks(
node, ((node.body, anno.Static.BODY_SCOPE),
(node.orelse, anno.Static.ORELSE_SCOPE)))
return node
def visit_Call(self, node):
self._enter_scope(False)
if self._is_overloaded_assign(node):
qn = anno.getanno(node.args[0], anno.Basic.QN)
self.scope.mark_modified(qn)
# We've already marked the first argument (lhs) as modified. Revisiting it
# will mark it as read, which is incorrect.
node.args = [node.args[0]] + self.visit_block(node.args[1:])
else:
node.args = self.visit_block(node.args)
node.keywords = self.visit_block(node.keywords)
# TODO(mdanatg): Account starargs, kwargs
anno.setanno(node, anno.Static.ARGS_SCOPE, self.scope)
self._exit_scope()
node.func = self.visit(node.func)
return node
def test_create_source_map(self):
def test_fn(x):
return x + 1
node, _ = parsing.parse_entity(test_fn)
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)
fn_node = node.body[0]
anno.setanno(fn_node.body[0], anno.Basic.ORIGIN, fake_origin)
converted_code = parsing.ast_to_source(fn_node)
source_map = origin_info.create_source_map(fn_node, converted_code,
'test_filename', [0])
loc = origin_info.LineLocation('test_filename', 2)
self.assertIn(loc, source_map)
self.assertIs(source_map[loc], fake_origin)
def test_basic(self):
node = ast.Name()
self.assertEqual(anno.keys(node), set())
self.assertFalse(anno.hasanno(node, 'foo'))
with self.assertRaises(AttributeError):
anno.getanno(node, 'foo')
anno.setanno(node, 'foo', 3)
self.assertEqual(anno.keys(node), {'foo'})
self.assertTrue(anno.hasanno(node, 'foo'))
self.assertEqual(anno.getanno(node, 'foo'), 3)
self.assertEqual(anno.getanno(node, 'bar', default=7), 7)
anno.delanno(node, 'foo')
self.assertEqual(anno.keys(node), set())
self.assertFalse(anno.hasanno(node, 'foo'))
with self.assertRaises(AttributeError):
anno.getanno(node, 'foo')
self.assertIsNone(anno.getanno(node, 'foo', default=None))
def visit_Subscript(self, node):
# TODO(mdanatg): This may no longer apply if we overload getitem.
node = self.generic_visit(node)
s = node.slice
if not isinstance(s, gast.Index):
# TODO(mdanatg): Support range and multi-dimensional indices.
# Continuing silently because some demos use these.
return node
if isinstance(s.value, gast.Num):
subscript = QN(NumberLiteral(s.value.n))
elif isinstance(s.value, gast.Str):
subscript = QN(StringLiteral(s.value.s))
else:
# The index may be an expression, case in which a name doesn't make sense.
if anno.hasanno(node.slice.value, anno.Basic.QN):
subscript = anno.getanno(node.slice.value, anno.Basic.QN)
else:
return node
if anno.hasanno(node.value, anno.Basic.QN):
anno.setanno(
node, anno.Basic.QN,
QN(anno.getanno(node.value, anno.Basic.QN),
subscript=subscript))
return node
def test_duplicate(self):
node = ast.If(test=ast.Num(1),
body=[ast.Expr(ast.Name('bar', ast.Load()))],
orelse=[])
anno.setanno(node, 'spam', 1)
anno.setanno(node, 'ham', 1)
anno.setanno(node.body[0], 'ham', 1)
anno.dup(node, {'spam': 'eggs'})
self.assertTrue(anno.hasanno(node, 'spam'))
self.assertTrue(anno.hasanno(node, 'ham'))
self.assertTrue(anno.hasanno(node, 'eggs'))
self.assertFalse(anno.hasanno(node.body[0], 'eggs'))
def visit_BinOp(self, node):
anno.setanno(node, 'enclosing_entities',
self.enclosing_entities)
return self.generic_visit(node)
def visit_Name(self, node):
node = self.generic_visit(node)
anno.setanno(node, anno.Basic.QN, QN(node.id))
return node
def visit_With(self, node):
self._enter_scope(False)
node = self.generic_visit(node)
anno.setanno(node, anno.Static.BODY_SCOPE, self.scope)
self._exit_scope()
return node
def visit(self, node):
anno.setanno(node, 'loop_state', self.state[LoopState].value)
anno.setanno(node, 'cond_state', self.state[CondState].value)
return super(TestTransformer, self).visit(node)
def _process_block_node(self, node, block, scope_name):
self._enter_scope(False)
block = self.visit_block(block)
anno.setanno(node, scope_name, self.scope)
self._exit_scope()
return node
def _process_statement(self, node):
self._enter_scope(False)
node = self.generic_visit(node)
anno.setanno(node, anno.Static.SCOPE, self.scope)
self._exit_scope()
return node
