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 visit_Attribute(self, node):
   node = self.generic_visit(node)
   parent_val = anno.getanno(node.value, STATIC_VALUE, default=None)
   if parent_val is not None and tf_inspect.ismodule(parent_val):
     if hasattr(parent_val, node.attr):
       anno.setanno(node, STATIC_VALUE, getattr(parent_val, node.attr))
   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 accompany 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, NodeAnno.BODY_SCOPE, self.scope)
    self._exit_scope()

    self._exit_scope()
    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, NodeAnno.ARGS_SCOPE, self.scope)
   self._exit_scope()
   return node

 def _block_statement_live_out(self, node):
   successors = self.current_analyzer.graph.stmt_next[node]
   stmt_live_out = set()
   for s in successors:
     stmt_live_out.update(self.current_analyzer.in_[s])
   anno.setanno(node, anno.Static.LIVE_VARS_OUT, frozenset(stmt_live_out))
   return node

  def visit_Compare(self, node):
    node = self.generic_visit(node)

    if not all(self._has_matching_func(op) for op in node.ops):
      if len(node.ops) == 1:
        # Basic expressions are safe to leave as they are.
        return node
      else:
        raise NotImplementedError(
            'compound expression with at least one unsupported '
            'operator: {}'.format(node.ops))

    ops_and_comps = list(zip(node.ops, node.comparators))
    left = node.left
    op_tree = None

    # Repeated comparisons are converted to conjunctions:
    #   a < b < c   ->   a < b and b < c
    while ops_and_comps:
      op, right = ops_and_comps.pop(0)
      binary_comparison = self._as_function(
          self._matching_func(op), (left, right))
      if isinstance(left, gast.Name) and isinstance(right, gast.Name):
        anno.setanno(binary_comparison, SAFE_BOOLEAN_OPERAND, True)
      if op_tree:
        self._expect_simple_symbol(right)
        op_tree = self._as_function('tf.logical_and',
                                    (binary_comparison, op_tree))
      else:
        op_tree = binary_comparison
      left = right
    assert op_tree is not None
    return op_tree

  def _track_symbol(self,
                    node,
                    composite_writes_alter_parent=False,
                    writes_create_symbol=False):
    # A QN may be missing when we have an attribute (or subscript) on a function
    # call. Example: a().b
    if not anno.hasanno(node, anno.Basic.QN):
      return
    qn = anno.getanno(node, anno.Basic.QN)

    if isinstance(node.ctx, gast.Store):
      self.scope.mark_write(qn)
      if qn.is_composite and composite_writes_alter_parent:
        self.scope.mark_write(qn.parent)
      if writes_create_symbol:
        self.scope.mark_creation(qn, writes_create_symbol=True)
      if self._in_aug_assign:
        self.scope.mark_read(qn)
    elif isinstance(node.ctx, gast.Load):
      self.scope.mark_read(qn)
    elif isinstance(node.ctx, gast.Param):
      # Param contexts appear in function defs, so they have the meaning of
      # defining a variable.
      self.scope.mark_write(qn)
      self.scope.mark_param(qn, self.enclosing_entities[-1])
    else:
      raise ValueError('Unknown context %s for node %s.' % (type(node.ctx), qn))

    anno.setanno(node, NodeAnno.IS_LOCAL, self.scope.has(qn))

    if self._in_return_statement:
      self.scope.mark_returned(qn)

 def _aggregate_successors_live_in(self, node):
   successors = self.current_analyzer.graph.stmt_next[node]
   node_live_out = set()
   for s in successors:
     node_live_out.update(self.current_analyzer.in_[s])
   anno.setanno(node, anno.Static.LIVE_VARS_OUT, frozenset(node_live_out))
   node = self.generic_visit(node)
   return node

 def _as_function(self, func_name, args):
   template = """
     func_name(args)
   """
   replacement = templates.replace_as_expression(
       template, func_name=parser.parse_expression(func_name), args=args)
   anno.setanno(replacement, SAFE_BOOLEAN_OPERAND, True)
   return replacement

 def visit_Name(self, node):
   node = self.generic_visit(node)
   if isinstance(node.ctx, gast.Load):
     defs = anno.getanno(node, anno.Static.DEFINITIONS, ())
     is_defined = bool(defs)
     if not is_defined and node.id in self.ctx.info.namespace:
       anno.setanno(node, STATIC_VALUE, self.ctx.info.namespace[node.id])
   return node

 def _process_statement_directive(self, call_node, directive):
   if self.local_scope_level < 1:
     raise ValueError(
         '"%s" must be used inside a statement' % directive.__name__)
   target = self.get_local(ENCLOSING_LOOP)
   node_anno = anno.getanno(target, converter.AgAnno.DIRECTIVES, {})
   node_anno[directive] = _map_args(call_node, directive)
   anno.setanno(target, converter.AgAnno.DIRECTIVES, node_anno)
   return call_node

 def visit_Call(self, node):
   self._enter_scope(False)
   node.args = self.visit_block(node.args)
   node.keywords = self.visit_block(node.keywords)
   # TODO(mdan): Account starargs, kwargs
   anno.setanno(node, NodeAnno.ARGS_SCOPE, self.scope)
   self._exit_scope()
   node.func = self.visit(node.func)
   return node

 def visit(self, node):
   node = super(Annotator, self).visit(node)
   if (self.current_analyzer is not None and
       isinstance(node, gast.stmt) and
       node in self.current_analyzer.graph.index):
     cfg_node = self.current_analyzer.graph.index[node]
     anno.setanno(node, anno.Static.LIVE_VARS_IN,
                  frozenset(self.current_analyzer.in_[cfg_node]))
   return node

 def visit_While(self, node):
   self._enter_scope(False)
   node.test = self.visit(node.test)
   anno.setanno(node, NodeAnno.COND_SCOPE, self.scope)
   anno.setanno(node.test, anno.Static.SCOPE, self.scope)
   self._exit_scope()
   node = self._process_parallel_blocks(node,
                                        ((node.body, NodeAnno.BODY_SCOPE),
                                         (node.orelse, NodeAnno.ORELSE_SCOPE)))
   return node

  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 visit_For(self, node):
   self._enter_scope(False)
   node.target = self.visit(node.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, NodeAnno.BODY_SCOPE),
                                         (node.orelse, NodeAnno.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 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 = tf_inspect.getsourcelines(function)
    function_filepath = tf_inspect.getsourcefile(function)
  else:
    function_lineno = None
    function_filepath = None

  # TODO(mdan): 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 _block_statement_live_in(self, node, entry_node):
   if entry_node in self.current_analyzer.graph.index:
     cfg_node = self.current_analyzer.graph.index[entry_node]
     stmt_live_in = frozenset(self.current_analyzer.in_[cfg_node])
   else:
     assert anno.hasanno(entry_node, anno.Static.LIVE_VARS_IN), (
         'If not matching a CFG node, must be a block statement:'
         ' {}'.format(entry_node))
     stmt_live_in = anno.getanno(entry_node, anno.Static.LIVE_VARS_IN)
   anno.setanno(node, anno.Static.LIVE_VARS_IN, stmt_live_in)
   return node

 def _process_function_arg(self, arg_node):
   qn = anno.getanno(arg_node, anno.Basic.QN)
   arg_name = str(qn)
   self.scope.setval(qn, arg_node)
   if (len(self.enclosing_entities) == 1 and
       arg_name in self.entity_info.arg_types):
     # Forge a node to hold the type information, so that method calls on
     # it can resolve the type.
     type_string, type_obj = self.entity_info.arg_types[arg_name]
     anno.setanno(arg_node, 'type', type_obj)
     anno.setanno(arg_node, 'type_fqn', tuple(type_string.split('.')))