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_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_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_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_parse_str(self):
   mod = parser.parse_str(
       textwrap.dedent("""
           def f(x):
             return x + 1
   """))
   self.assertEqual('f', mod.body[0].name)

 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(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_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

    # Just check that the output is parseable Python code.
    self.assertIsNotNone(parser.parse_str(api.to_code(test_fn)))

 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_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_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_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_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_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_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_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)